diff --git a/crates/compiler/builtins/roc/List.roc b/crates/compiler/builtins/roc/List.roc index ac65a8c5f0..7d5b90d40e 100644 --- a/crates/compiler/builtins/roc/List.roc +++ b/crates/compiler/builtins/roc/List.roc @@ -62,6 +62,7 @@ interface List sortDesc, reserve, walkBackwardsUntil, + countIf, ] imports [ Bool.{ Bool }, @@ -528,6 +529,18 @@ dropIf : List a, (a -> Bool) -> List a dropIf = \list, predicate -> List.keepIf list (\e -> Bool.not (predicate e)) +## Run the given function on each element of a list, and return the +## number of elements for which the function returned `Bool.true`. +countIf : List a, (a -> Bool) -> Nat +countIf = \list, predicate -> + walkState = \state, elem -> + if predicate elem then + state + 1 + else + state + + List.walk list 0 walkState + ## This works like [List.map], except only the transformed values that are ## wrapped in `Ok` are kept. Any that are wrapped in `Err` are dropped. ## diff --git a/crates/compiler/module/src/symbol.rs b/crates/compiler/module/src/symbol.rs index 8f1af3238a..cd01b2b4c6 100644 --- a/crates/compiler/module/src/symbol.rs +++ b/crates/compiler/module/src/symbol.rs @@ -1377,6 +1377,7 @@ define_builtins! { 74 LIST_MAP_TRY: "mapTry" 75 LIST_WALK_TRY: "walkTry" 76 LIST_WALK_BACKWARDS_UNTIL: "walkBackwardsUntil" + 77 LIST_COUNT_IF: "countIf" } 7 RESULT: "Result" => { 0 RESULT_RESULT: "Result" exposed_type=true // the Result.Result type alias diff --git a/crates/compiler/test_gen/src/gen_list.rs b/crates/compiler/test_gen/src/gen_list.rs index 91c43b1282..9d5985caa3 100644 --- a/crates/compiler/test_gen/src/gen_list.rs +++ b/crates/compiler/test_gen/src/gen_list.rs @@ -1065,6 +1065,92 @@ fn list_keep_if_str_is_hello() { ); } +#[test] +#[cfg(any(feature = "gen-llvm", feature = "gen-wasm"))] +fn list_count_if_empty_list() { + assert_evals_to!( + indoc!( + r#" + List.countIf [] \_ -> Bool.true + "# + ), + 0, + usize + ); +} + +#[test] +#[cfg(any(feature = "gen-llvm", feature = "gen-wasm"))] +fn list_count_if_always_true_for_non_empty_list() { + assert_evals_to!( + indoc!( + r#" + alwaysTrue : I64 -> Bool + alwaysTrue = \_ -> + Bool.true + + oneThroughEight : List I64 + oneThroughEight = + [1,2,3,4,5,6,7,8] + + List.countIf oneThroughEight alwaysTrue + "# + ), + 8, + usize + ); +} + +#[test] +#[cfg(any(feature = "gen-llvm", feature = "gen-wasm"))] +fn list_count_if_always_false_for_non_empty_list() { + assert_evals_to!( + indoc!( + r#" + alwaysFalse : I64 -> Bool + alwaysFalse = \_ -> + Bool.false + + List.countIf [1,2,3,4,5,6,7,8] alwaysFalse + "# + ), + 0, + usize + ); +} + +#[test] +#[cfg(any(feature = "gen-llvm", feature = "gen-wasm"))] +fn list_count_if_condition() { + assert_evals_to!( + indoc!( + r#" + intIsLessThanThree : I64 -> Bool + intIsLessThanThree = \i -> + i < 3 + + List.countIf [1,2,3,4,5,6,7,8] intIsLessThanThree + "# + ), + 2, + usize + ); +} + +#[test] +#[cfg(any(feature = "gen-llvm", feature = "gen-wasm"))] +fn list_count_if_str() { + assert_evals_to!( + indoc!( + r#" + List.countIf ["x", "y", "x"] (\x -> x == "x") + "# + ), + 2, + usize + ); +} + #[test] #[cfg(any(feature = "gen-llvm", feature = "gen-wasm"))] fn list_map_on_empty_list_with_int_layout() {