v: add array.count as a method that accepts a predicate, similar to f…

…ilter, but returning just the number of matches (vlang#23054)

vlib/builtin/array.v

@@ -807,6 +807,13 @@ pub fn (a array) filter(predicate fn (voidptr) bool) array
 // Example: array.any(it.name == 'Bob') // will yield `true` if any element has `.name == 'Bob'`
 pub fn (a array) any(predicate fn (voidptr) bool) bool
 
+// count counts how many elements in array pass the test.
+// Ignore the function signature. `count` does not take an actual callback. Rather, it
+// takes an `it` expression.
+//
+// Example: array.count(it % 2 == 1) // will return how many elements are odd
+pub fn (a array) count(predicate fn (voidptr) bool) int
+
 // all tests whether all elements in the array pass the test.
 // Ignore the function signature. `all` does not take an actual callback. Rather, it
 // takes an `it` expression.

vlib/v/checker/checker.v

@@ -30,10 +30,10 @@ const generic_fn_postprocess_iterations_cutoff_limit = 1_000_000
 // Note that methods that do not return anything, or that return known types, are not listed here, since they are just ordinary non generic methods.
 pub const array_builtin_methods = ['filter', 'clone', 'repeat', 'reverse', 'map', 'slice', 'sort',
 	'sort_with_compare', 'sorted', 'sorted_with_compare', 'contains', 'index', 'wait', 'any', 'all',
-	'first', 'last', 'pop', 'delete', 'insert', 'prepend']
+	'first', 'last', 'pop', 'delete', 'insert', 'prepend', 'count']
 pub const array_builtin_methods_chk = token.new_keywords_matcher_from_array_trie(array_builtin_methods)
 pub const fixed_array_builtin_methods = ['contains', 'index', 'any', 'all', 'wait', 'map', 'sort',
-	'sorted', 'sort_with_compare', 'sorted_with_compare', 'reverse', 'reverse_in_place']
+	'sorted', 'sort_with_compare', 'sorted_with_compare', 'reverse', 'reverse_in_place', 'count']
 pub const fixed_array_builtin_methods_chk = token.new_keywords_matcher_from_array_trie(fixed_array_builtin_methods)
 // TODO: remove `byte` from this list when it is no longer supported
 pub const reserved_type_names = ['byte', 'bool', 'char', 'i8', 'i16', 'int', 'i64', 'u8', 'u16',

vlib/v/checker/fn.v

@@ -3118,7 +3118,7 @@ fn (mut c Checker) fn_call_error_have_want(p HaveWantParams) {
 	c.error('expected ${p.nr_params} ${args_plural}, but got ${p.nr_args}', p.pos)
 }
 
-fn (mut c Checker) check_map_and_filter(is_map bool, elem_typ ast.Type, node ast.CallExpr) {
+fn (mut c Checker) check_predicate_param(is_map bool, elem_typ ast.Type, node ast.CallExpr) {
 	if node.args.len != 1 {
 		c.error('expected 1 argument, but got ${node.args.len}', node.pos)
 		// Finish early so that it doesn't fail later
@@ -3327,7 +3327,7 @@ fn (mut c Checker) array_builtin_method_call(mut node ast.CallExpr, left_type as
 		c.table.sym(unaliased_left_type).info as ast.Array
 	}
 	elem_typ = array_info.elem_type
-	if method_name in ['filter', 'map', 'any', 'all'] {
+	if method_name in ['filter', 'map', 'any', 'all', 'count'] {
 		if node.args.len > 0 && mut node.args[0].expr is ast.LambdaExpr {
 			if node.args[0].expr.params.len != 1 {
 				c.error('lambda expressions used in the builtin array methods require exactly 1 parameter',
@@ -3513,7 +3513,7 @@ fn (mut c Checker) array_builtin_method_call(mut node ast.CallExpr, left_type as
 	if method_name == 'map' {
 		// eprintln('>>>>>>> map node.args[0].expr: ${node.args[0].expr}, left_type: ${left_type} | elem_typ: ${elem_typ} | arg_type: ${arg_type}')
 		// check fn
-		c.check_map_and_filter(true, elem_typ, node)
+		c.check_predicate_param(true, elem_typ, node)
 		arg_sym := c.table.sym(arg_type)
 		ret_type := match arg_sym.info {
 			ast.FnType {
@@ -3542,10 +3542,13 @@ fn (mut c Checker) array_builtin_method_call(mut node ast.CallExpr, left_type as
 		} else if node.left.is_auto_deref_var() {
 			node.return_type = node.return_type.deref()
 		}
-		c.check_map_and_filter(false, elem_typ, node)
+		c.check_predicate_param(false, elem_typ, node)
 	} else if method_name in ['any', 'all'] {
-		c.check_map_and_filter(false, elem_typ, node)
+		c.check_predicate_param(false, elem_typ, node)
 		node.return_type = ast.bool_type
+	} else if method_name == 'count' {
+		c.check_predicate_param(false, elem_typ, node)
+		node.return_type = ast.int_type
 	} else if method_name == 'clone' {
 		if node.args.len != 0 {
 			c.error('`.clone()` does not have any arguments', node.args[0].pos)
@@ -3715,8 +3718,28 @@ fn (mut c Checker) fixed_array_builtin_method_call(mut node ast.CallExpr, left_t
 			scope_register_it(mut node.scope, node.pos, elem_typ)
 		}
 		c.expr(mut node.args[0].expr)
-		c.check_map_and_filter(false, elem_typ, node)
+		c.check_predicate_param(false, elem_typ, node)
 		node.return_type = ast.bool_type
+	} else if method_name == 'count' {
+		if node.args.len != 1 {
+			c.error('`.${method_name}` expected 1 argument, but got ${node.args.len}',
+				node.pos)
+			return ast.bool_type
+		}
+		if node.args.len > 0 && mut node.args[0].expr is ast.LambdaExpr {
+			if node.args[0].expr.params.len != 1 {
+				c.error('lambda expressions used in the builtin array methods require exactly 1 parameter',
+					node.args[0].expr.pos)
+				return ast.bool_type
+			}
+			c.support_lambda_expr_one_param(elem_typ, ast.bool_type, mut node.args[0].expr)
+		} else {
+			// position of `it` doesn't matter
+			scope_register_it(mut node.scope, node.pos, elem_typ)
+		}
+		c.expr(mut node.args[0].expr)
+		c.check_predicate_param(false, elem_typ, node)
+		node.return_type = ast.int_type
 	} else if method_name == 'wait' {
 		elem_sym := c.table.sym(elem_typ)
 		if elem_sym.kind == .thread {
@@ -3757,7 +3780,7 @@ fn (mut c Checker) fixed_array_builtin_method_call(mut node ast.CallExpr, left_t
 			scope_register_it(mut node.scope, node.pos, elem_typ)
 		}
 
-		c.check_map_and_filter(true, elem_typ, node)
+		c.check_predicate_param(true, elem_typ, node)
 		arg_type := c.check_expr_option_or_result_call(node.args[0].expr, c.expr(mut node.args[0].expr))
 		arg_sym := c.table.sym(arg_type)
 		ret_type := match arg_sym.info {

vlib/v/checker/tests/array_count_err.out

@@ -0,0 +1,20 @@
+vlib/v/checker/tests/array_count_err.vv:4:4: error: expected 1 argument, but got 2
+    2 |     a := []int{}
+    3 |     a.count(1)
+    4 |     a.count(1, 2)
+      |       ~~~~~~~~~~~
+    5 |     a.count('')
+    6 |     a.count()
+vlib/v/checker/tests/array_count_err.vv:5:10: error: type mismatch, should use e.g. `count(it > 2)`
+    3 |     a.count(1)
+    4 |     a.count(1, 2)
+    5 |     a.count('')
+      |             ~~
+    6 |     a.count()
+    7 | }
+vlib/v/checker/tests/array_count_err.vv:6:4: error: expected 1 argument, but got 0
+    4 |     a.count(1, 2)
+    5 |     a.count('')
+    6 |     a.count()
+      |       ~~~~~~~
+    7 | }

vlib/v/checker/tests/array_count_err.vv

@@ -0,0 +1,7 @@
+fn main() {
+	a := []int{}
+	a.count(1)
+	a.count(1, 2)
+	a.count('')
+	a.count()
+}

vlib/v/fmt/fmt.v

@@ -2058,7 +2058,7 @@ fn (mut f Fmt) write_static_method(name string, short_name string) {
 pub fn (mut f Fmt) call_expr(node ast.CallExpr) {
 	mut is_method_newline := false
 	if node.is_method {
-		if node.name in ['map', 'filter', 'all', 'any'] {
+		if node.name in ['map', 'filter', 'all', 'any', 'count'] {
 			f.in_lambda_depth++
 			defer { f.in_lambda_depth-- }
 		}

vlib/v/gen/c/array.v

@@ -558,7 +558,7 @@ fn (mut g Gen) gen_array_map(node ast.CallExpr) {
 			}
 		}
 		ast.CallExpr {
-			if expr.name in ['map', 'filter', 'all', 'any'] {
+			if expr.name in ['map', 'filter', 'all', 'any', 'count'] {
 				is_embed_map_filter = true
 				g.set_current_pos_as_last_stmt_pos()
 			}
@@ -940,7 +940,7 @@ fn (mut g Gen) gen_array_filter(node ast.CallExpr) {
 			}
 		}
 		ast.CallExpr {
-			if expr.name in ['map', 'filter', 'all', 'any'] {
+			if expr.name in ['map', 'filter', 'all', 'any', 'count'] {
 				is_embed_map_filter = true
 				g.set_current_pos_as_last_stmt_pos()
 			}
@@ -1439,7 +1439,7 @@ fn (mut g Gen) gen_array_any(node ast.CallExpr) {
 			}
 		}
 		ast.CallExpr {
-			if expr.name in ['map', 'filter', 'all', 'any'] {
+			if expr.name in ['map', 'filter', 'all', 'any', 'count'] {
 				is_embed_map_filter = true
 				g.set_current_pos_as_last_stmt_pos()
 			}
@@ -1469,6 +1469,96 @@ fn (mut g Gen) gen_array_any(node ast.CallExpr) {
 	}
 }
 
+fn (mut g Gen) gen_array_count(node ast.CallExpr) {
+	past := g.past_tmp_var_new()
+	defer {
+		g.past_tmp_var_done(past)
+	}
+
+	sym := g.table.final_sym(node.left_type)
+	left_is_array := sym.kind == .array
+	elem_type := if left_is_array {
+		(sym.info as ast.Array).elem_type
+	} else {
+		(sym.info as ast.ArrayFixed).elem_type
+	}
+	elem_type_str := g.styp(elem_type)
+	has_infix_left_var_name := g.write_prepared_tmp_value(past.tmp_var, node, 'int', '0')
+
+	mut expr := node.args[0].expr
+	var_name := g.get_array_expr_param_name(mut expr)
+
+	mut closure_var := ''
+	if mut expr is ast.AnonFn {
+		if expr.inherited_vars.len > 0 {
+			closure_var = g.new_tmp_var()
+			g.declare_closure_fn(mut expr, closure_var)
+		}
+	}
+	i := g.new_tmp_var()
+	g.writeln('for (int ${i} = 0; ${i} < ${past.tmp_var}_len; ++${i}) {')
+	g.indent++
+
+	g.write_prepared_var(var_name, elem_type, elem_type_str, past.tmp_var, i, left_is_array)
+	g.set_current_pos_as_last_stmt_pos()
+	mut is_embed_map_filter := false
+	match mut expr {
+		ast.AnonFn {
+			g.write('if (')
+			if expr.inherited_vars.len > 0 {
+				g.write_closure_fn(mut expr, var_name, closure_var)
+			} else {
+				g.gen_anon_fn_decl(mut expr)
+				g.write('${expr.decl.name}(${var_name})')
+			}
+		}
+		ast.Ident {
+			g.write('if (')
+			if expr.kind == .function {
+				g.write('${c_name(expr.name)}(${var_name})')
+			} else if expr.kind == .variable {
+				var_info := expr.var_info()
+				sym_t := g.table.sym(var_info.typ)
+				if sym_t.kind == .function {
+					g.write('${c_name(expr.name)}(${var_name})')
+				} else {
+					g.expr(expr)
+				}
+			} else {
+				g.expr(expr)
+			}
+		}
+		ast.CallExpr {
+			if expr.name in ['map', 'filter', 'all', 'any', 'count'] {
+				is_embed_map_filter = true
+				g.set_current_pos_as_last_stmt_pos()
+			}
+			g.write('if (')
+			g.expr(expr)
+		}
+		ast.LambdaExpr {
+			g.write('if (')
+			g.expr(expr.expr)
+		}
+		else {
+			g.write('if (')
+			g.expr(expr)
+		}
+	}
+	g.writeln2(') {', '\t++${past.tmp_var};')
+	g.writeln('}')
+	g.indent--
+	g.writeln('}')
+	if !is_embed_map_filter {
+		g.set_current_pos_as_last_stmt_pos()
+	}
+	if has_infix_left_var_name {
+		g.indent--
+		g.writeln('}')
+		g.set_current_pos_as_last_stmt_pos()
+	}
+}
+
 fn (mut g Gen) gen_array_all(node ast.CallExpr) {
 	past := g.past_tmp_var_new()
 	defer {
@@ -1532,7 +1622,7 @@ fn (mut g Gen) gen_array_all(node ast.CallExpr) {
 			}
 		}
 		ast.CallExpr {
-			if expr.name in ['map', 'filter', 'all', 'any'] {
+			if expr.name in ['map', 'filter', 'all', 'any', 'count'] {
 				is_embed_map_filter = true
 				g.set_current_pos_as_last_stmt_pos()
 			}

vlib/v/gen/c/fn.v

@@ -1155,6 +1155,9 @@ fn (mut g Gen) gen_array_method_call(node ast.CallExpr, left_type ast.Type, left
 		'any' {
 			g.gen_array_any(node)
 		}
+		'count' {
+			g.gen_array_count(node)
+		}
 		'all' {
 			g.gen_array_all(node)
 		}
@@ -1251,6 +1254,9 @@ fn (mut g Gen) gen_fixed_array_method_call(node ast.CallExpr, left_type ast.Type
 		'any' {
 			g.gen_array_any(node)
 		}
+		'count' {
+			g.gen_array_count(node)
+		}
 		'all' {
 			g.gen_array_all(node)
 		}

vlib/v/gen/c/infix.v

@@ -1057,7 +1057,7 @@ fn (mut g Gen) need_tmp_var_in_array_call(node ast.Expr) bool {
 	match node {
 		ast.CallExpr {
 			if node.left_type != 0 && g.table.sym(node.left_type).kind == .array
-				&& node.name in ['all', 'any', 'filter', 'map'] {
+				&& node.name in ['all', 'any', 'filter', 'map', 'count'] {
 				return true
 			}
 		}

vlib/v/gen/golang/golang.v

@@ -1419,7 +1419,7 @@ pub fn (mut f Gen) call_expr(node ast.CallExpr) {
 	// for arg in node.args {}
 	mut is_method_newline := false
 	if node.is_method {
-		if node.name in ['map', 'filter', 'all', 'any'] {
+		if node.name in ['map', 'filter', 'all', 'any', 'count'] {
 			f.in_lambda_depth++
 			defer {
 				f.in_lambda_depth--

vlib/v/parser/parser.v

@@ -3332,7 +3332,7 @@ fn (mut p Parser) dot_expr(left ast.Expr) ast.Expr {
 	} else {
 		p.name_error = true
 	}
-	is_filter := field_name in ['filter', 'map', 'any', 'all']
+	is_filter := field_name in ['filter', 'map', 'any', 'all', 'count']
 	if is_filter || field_name == 'sort' || field_name == 'sorted' {
 		p.open_scope()
 		defer {

vlib/v/tests/builtin_arrays/array_count_test.v

@@ -0,0 +1,36 @@
+fn test_main() {
+	a := []int{len: 10, init: index}
+	assert a.count(it % 2) == 5
+
+	b := [10]int{init: index}
+	assert a.count(it % 2) == 5
+}
+
+fn test_zero() {
+	a := []int{len: 10, init: index}
+	assert a.count(it == 1000) == 0
+
+	b := [10]int{init: index}
+	assert a.count(it == 1000) == 0
+}
+
+fn test_struct() {
+	struct Abc {
+		x int
+		y int
+		z string
+	}
+
+	a := [Abc{}, Abc{1, 2, 'abc'}, Abc{100, 2, 'def'}, Abc{0, 0, 'a'}]
+
+	assert dump(a.count(it.z.starts_with('a'))) == 2
+	assert dump(a.count(it.y == 2)) == 2
+	assert dump(a.count(it.z.len == 1)) == 1
+	assert dump(a.count(it.z.len < 3)) == 2
+
+	sa := ['aa', 'bb', 'ccc']
+	dump(sa)
+	assert dump(sa.count(it.len < 3)) == 2
+	assert dump(sa.count(it == 'aa')) == 1
+	assert dump(sa.count(it.len == 3)) == 1
+}