From 1eb3867f4e164f55d252bc3305e8b6b6decac5e7 Mon Sep 17 00:00:00 2001 From: Felipe Pena Date: Tue, 3 Dec 2024 15:22:14 -0300 Subject: [PATCH] v: add array.count as a method that accepts a predicate, similar to filter, but returning just the number of matches (#23054) --- vlib/builtin/array.v | 7 ++ vlib/v/checker/checker.v | 4 +- vlib/v/checker/fn.v | 37 +++++-- vlib/v/checker/tests/array_count_err.out | 20 ++++ vlib/v/checker/tests/array_count_err.vv | 7 ++ vlib/v/fmt/fmt.v | 2 +- vlib/v/gen/c/array.v | 98 ++++++++++++++++++- vlib/v/gen/c/fn.v | 6 ++ vlib/v/gen/c/infix.v | 2 +- vlib/v/gen/golang/golang.v | 2 +- vlib/v/parser/parser.v | 2 +- .../v/tests/builtin_arrays/array_count_test.v | 36 +++++++ 12 files changed, 206 insertions(+), 17 deletions(-) create mode 100644 vlib/v/checker/tests/array_count_err.out create mode 100644 vlib/v/checker/tests/array_count_err.vv create mode 100644 vlib/v/tests/builtin_arrays/array_count_test.v diff --git a/vlib/builtin/array.v b/vlib/builtin/array.v index 6fd5c596b1924a..4766e192229d89 100644 --- a/vlib/builtin/array.v +++ b/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. diff --git a/vlib/v/checker/checker.v b/vlib/v/checker/checker.v index 00f4d766f05e8d..dc752671dd635f 100644 --- a/vlib/v/checker/checker.v +++ b/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', diff --git a/vlib/v/checker/fn.v b/vlib/v/checker/fn.v index b87aa24ac052a0..ef1a7bed8fc1a2 100644 --- a/vlib/v/checker/fn.v +++ b/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 { diff --git a/vlib/v/checker/tests/array_count_err.out b/vlib/v/checker/tests/array_count_err.out new file mode 100644 index 00000000000000..e9521a705bd0ca --- /dev/null +++ b/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 | } diff --git a/vlib/v/checker/tests/array_count_err.vv b/vlib/v/checker/tests/array_count_err.vv new file mode 100644 index 00000000000000..531d701899d962 --- /dev/null +++ b/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() +} diff --git a/vlib/v/fmt/fmt.v b/vlib/v/fmt/fmt.v index 66cac0d42e594b..4d5d59ccead4a6 100644 --- a/vlib/v/fmt/fmt.v +++ b/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-- } } diff --git a/vlib/v/gen/c/array.v b/vlib/v/gen/c/array.v index 4130bcc5654b26..a3827c9b3c0a63 100644 --- a/vlib/v/gen/c/array.v +++ b/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() } diff --git a/vlib/v/gen/c/fn.v b/vlib/v/gen/c/fn.v index a490f7f58f339e..484acb3f772a9b 100644 --- a/vlib/v/gen/c/fn.v +++ b/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) } diff --git a/vlib/v/gen/c/infix.v b/vlib/v/gen/c/infix.v index bc4d0b28817f4f..349455b3189283 100644 --- a/vlib/v/gen/c/infix.v +++ b/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 } } diff --git a/vlib/v/gen/golang/golang.v b/vlib/v/gen/golang/golang.v index 8076a12bf1844c..a452f0659f579d 100644 --- a/vlib/v/gen/golang/golang.v +++ b/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-- diff --git a/vlib/v/parser/parser.v b/vlib/v/parser/parser.v index 2b3122d333f720..79efa142171b5f 100644 --- a/vlib/v/parser/parser.v +++ b/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 { diff --git a/vlib/v/tests/builtin_arrays/array_count_test.v b/vlib/v/tests/builtin_arrays/array_count_test.v new file mode 100644 index 00000000000000..5a9488569ee63b --- /dev/null +++ b/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 +}