diff --git a/compiler/builtins/src/std.rs b/compiler/builtins/src/std.rs index a78af7e971..f934e83143 100644 --- a/compiler/builtins/src/std.rs +++ b/compiler/builtins/src/std.rs @@ -877,6 +877,19 @@ pub fn types() -> MutMap { Box::new(list_type(int_type(flex(TVAR1)))), ); + // joinMap : List before, (before -> List after) -> List after + { + let_tvars! { cvar, before, after } + add_top_level_function_type!( + Symbol::LIST_JOIN_MAP, + vec![ + list_type(flex(before)), + closure(vec![flex(before)], cvar, Box::new(list_type(flex(after)))), + ], + Box::new(list_type(flex(after))), + ); + } + // map : List before, (before -> after) -> List after add_top_level_function_type!( Symbol::LIST_MAP, diff --git a/compiler/can/src/builtins.rs b/compiler/can/src/builtins.rs index d89fc247e7..f4a2b2de33 100644 --- a/compiler/can/src/builtins.rs +++ b/compiler/can/src/builtins.rs @@ -86,6 +86,7 @@ pub fn builtin_defs_map(symbol: Symbol, var_store: &mut VarStore) -> Option LIST_PRODUCT => list_product, LIST_PREPEND => list_prepend, LIST_JOIN => list_join, + LIST_JOIN_MAP => list_join_map, LIST_MAP => list_map, LIST_MAP2 => list_map2, LIST_MAP3 => list_map3, @@ -2200,6 +2201,83 @@ fn list_walk_until(symbol: Symbol, var_store: &mut VarStore) -> Def { lowlevel_3(symbol, LowLevel::ListWalkUntil, var_store) } +/// List.joinMap : List before, (before -> List after) -> List after +fn list_join_map(symbol: Symbol, var_store: &mut VarStore) -> Def { + let before = var_store.fresh(); + let list_before = var_store.fresh(); + let after = var_store.fresh(); + let list_after = var_store.fresh(); + let before2list_after = var_store.fresh(); + let t_concat_clos = var_store.fresh(); + let mapper_lambda_set = var_store.fresh(); + + // \state, elem -> List.concat state (mapper elem) + let concat_clos = Closure(ClosureData { + function_type: t_concat_clos, + closure_type: var_store.fresh(), + closure_ext_var: var_store.fresh(), + return_type: list_after, + name: Symbol::LIST_JOIN_MAP_CONCAT, + recursive: Recursive::NotRecursive, + captured_symbols: vec![(Symbol::ARG_2, before2list_after)], + arguments: vec![ + (list_after, no_region(Pattern::Identifier(Symbol::ARG_3))), + (before, no_region(Pattern::Identifier(Symbol::ARG_4))), + ], + loc_body: { + let mapper = Box::new(( + before2list_after, + no_region(Var(Symbol::ARG_2)), + mapper_lambda_set, + list_after, // return type + )); + // (mapper elem) + let mapper_elem = Call( + mapper, + vec![(before, no_region(Var(Symbol::ARG_4)))], + CalledVia::Space, + ); + Box::new(no_region(RunLowLevel { + op: LowLevel::ListConcat, + args: vec![(list_after, Var(Symbol::ARG_3)), (list_after, mapper_elem)], + ret_var: list_after, + })) + }, + }); + + // List.joinMap = \input_list, mapper -> + // List.walk [] input_list (\state, elem -> List.concat state (mapper elem)) + let body = RunLowLevel { + op: LowLevel::ListWalk, + args: vec![ + // input_list : List before + (list_before, Var(Symbol::ARG_1)), + // [] : List after + ( + list_after, + List { + elem_var: after, + loc_elems: vec![], + }, + ), + // \state, elem -> List.concat state (mapper elem) + (t_concat_clos, concat_clos), + ], + ret_var: list_after, + }; + + defn( + symbol, + vec![ + (list_before, Symbol::ARG_1), + (before2list_after, Symbol::ARG_2), + ], + var_store, + body, + list_after, + ) +} + // min : List (Num a) -> Result (Num a) [ ListWasEmpty ]* fn list_min(symbol: Symbol, var_store: &mut VarStore) -> Def { let arg_var = var_store.fresh(); diff --git a/compiler/module/src/symbol.rs b/compiler/module/src/symbol.rs index 36aa6e6675..22a3521b1b 100644 --- a/compiler/module/src/symbol.rs +++ b/compiler/module/src/symbol.rs @@ -1062,6 +1062,8 @@ define_builtins! { 39 LIST_MAX_GT: "#maxGt" 40 LIST_MAP4: "map4" 41 LIST_DROP_FIRST: "dropFirst" + 42 LIST_JOIN_MAP: "joinMap" + 43 LIST_JOIN_MAP_CONCAT: "#joinMapConcat" } 5 RESULT: "Result" => { 0 RESULT_RESULT: "Result" imported // the Result.Result type alias diff --git a/compiler/test_gen/src/gen_list.rs b/compiler/test_gen/src/gen_list.rs index 81a86804bf..cfaaf122cb 100644 --- a/compiler/test_gen/src/gen_list.rs +++ b/compiler/test_gen/src/gen_list.rs @@ -2221,3 +2221,35 @@ fn empty_list_of_function_type() { RocStr ); } + +#[test] +fn list_join_map() { + assert_evals_to!( + indoc!( + r#" + List.joinMap ["guava,apple,pear", "bailey,cyrus"] (\s -> Str.split s ",") + "# + ), + RocList::from_slice(&[ + RocStr::from_slice("guava".as_bytes()), + RocStr::from_slice("apple".as_bytes()), + RocStr::from_slice("pear".as_bytes()), + RocStr::from_slice("bailey".as_bytes()), + RocStr::from_slice("cyrus".as_bytes()), + ]), + RocList + ); +} + +#[test] +fn list_join_map_empty() { + assert_evals_to!( + indoc!( + r#" + List.joinMap [] (\s -> Str.split s ",") + "# + ), + RocList::from_slice(&[]), + RocList + ); +}