diff --git a/compiler/builtins/src/std.rs b/compiler/builtins/src/std.rs index 41ed828beb..784927b135 100644 --- a/compiler/builtins/src/std.rs +++ b/compiler/builtins/src/std.rs @@ -1159,6 +1159,13 @@ pub fn types() -> MutMap { ) } + // intersperse : List elem, elem -> List elem + add_top_level_function_type!( + Symbol::LIST_INTERSPERSE, + vec![list_type(flex(TVAR1)), flex(TVAR1)], + Box::new(list_type(flex(TVAR1))), + ); + // Dict module // len : Dict * * -> Nat diff --git a/compiler/can/src/builtins.rs b/compiler/can/src/builtins.rs index bda105589d..163a82fd92 100644 --- a/compiler/can/src/builtins.rs +++ b/compiler/can/src/builtins.rs @@ -97,6 +97,7 @@ pub fn builtin_defs_map(symbol: Symbol, var_store: &mut VarStore) -> Option LIST_TAKE_LAST => list_take_last, LIST_SUBLIST => list_sublist, LIST_SPLIT => list_split, + LIST_INTERSPERSE => list_intersperse, LIST_DROP => list_drop, LIST_DROP_AT => list_drop_at, LIST_DROP_FIRST => list_drop_first, @@ -2148,6 +2149,84 @@ fn list_sublist(symbol: Symbol, var_store: &mut VarStore) -> Def { ) } +/// List.intersperse : List elem, elem -> List elem +fn list_intersperse(symbol: Symbol, var_store: &mut VarStore) -> Def { + let list_var = var_store.fresh(); + let sep_var = var_store.fresh(); + + let list_sym = Symbol::ARG_1; + let sep_sym = Symbol::ARG_2; + + let clos_var = var_store.fresh(); + let clos_acc_var = var_store.fresh(); + + let clos_sym = Symbol::LIST_INTERSPERSE_CLOS; + let clos_acc_sym = Symbol::ARG_3; + let clos_elem_sym = Symbol::ARG_4; + + let int_var = var_store.fresh(); + let zero = int(int_var, Variable::NATURAL, 0); + + // \acc, elem -> acc |> List.append sep |> List.append elem + let clos = Closure(ClosureData { + function_type: clos_var, + closure_type: var_store.fresh(), + closure_ext_var: var_store.fresh(), + return_type: clos_acc_var, + name: clos_sym, + recursive: Recursive::NotRecursive, + captured_symbols: vec![(sep_sym, sep_var)], + arguments: vec![ + (clos_acc_var, no_region(Pattern::Identifier(clos_acc_sym))), + (sep_var, no_region(Pattern::Identifier(clos_elem_sym))), + ], + loc_body: { + let append_sep = RunLowLevel { + op: LowLevel::ListAppend, + args: vec![(clos_acc_var, Var(clos_acc_sym)), (sep_var, Var(sep_sym))], + ret_var: clos_acc_var, + }; + + Box::new(no_region(RunLowLevel { + op: LowLevel::ListAppend, + args: vec![(clos_acc_var, append_sep), (sep_var, Var(clos_elem_sym))], + ret_var: clos_acc_var, + })) + }, + }); + + // List.walk [] l (\acc, elem -> acc |> List.append sep |> List.append elem) + let acc = RunLowLevel { + op: LowLevel::ListWalk, + args: vec![ + (list_var, Var(list_sym)), + ( + clos_acc_var, + List { + elem_var: sep_var, + loc_elems: vec![], + }, + ), + (clos_var, clos), + ], + ret_var: clos_acc_var, + }; + + let body = RunLowLevel { + op: LowLevel::ListDropAt, + args: vec![(clos_acc_var, acc), (int_var, zero)], + ret_var: clos_acc_var, + }; + + defn( + symbol, + vec![(list_var, list_sym), (sep_var, sep_sym)], + var_store, + body, + clos_acc_var, + ) +} + /// List.split : List elem, Nat -> { before: List elem, others: List elem } fn list_split(symbol: Symbol, var_store: &mut VarStore) -> Def { let list_var = var_store.fresh(); diff --git a/compiler/module/src/symbol.rs b/compiler/module/src/symbol.rs index 77e790b68d..d4446ce2cf 100644 --- a/compiler/module/src/symbol.rs +++ b/compiler/module/src/symbol.rs @@ -1075,8 +1075,10 @@ define_builtins! { 47 LIST_FIND: "find" 48 LIST_FIND_RESULT: "#find_result" // symbol used in the definition of List.find 49 LIST_SUBLIST: "sublist" - 50 LIST_SPLIT: "split" - 51 LIST_SPLIT_CLOS: "#splitClos" + 50 LIST_INTERSPERSE: "intersperse" + 51 LIST_INTERSPERSE_CLOS: "#intersperseClos" + 52 LIST_SPLIT: "split" + 53 LIST_SPLIT_CLOS: "#splitClos" } 5 RESULT: "Result" => { 0 RESULT_RESULT: "Result" imported // the Result.Result type alias diff --git a/compiler/solve/tests/solve_expr.rs b/compiler/solve/tests/solve_expr.rs index 1e12191375..90cf426dc1 100644 --- a/compiler/solve/tests/solve_expr.rs +++ b/compiler/solve/tests/solve_expr.rs @@ -3813,6 +3813,17 @@ mod solve_expr { ); } + #[test] + fn list_intersperse() { + infer_eq_without_problem( + indoc!( + r#" + List.intersperse + "# + ), + "List a, a -> List a", + ); + } #[test] fn function_that_captures_nothing_is_not_captured() { // we should make sure that a function that doesn't capture anything it not itself captured diff --git a/compiler/test_gen/src/gen_list.rs b/compiler/test_gen/src/gen_list.rs index 86aba0e539..9a8e5d0790 100644 --- a/compiler/test_gen/src/gen_list.rs +++ b/compiler/test_gen/src/gen_list.rs @@ -318,6 +318,29 @@ fn list_drop_at() { assert_evals_to!("List.dropAt [0] 0", RocList::from_slice(&[]), RocList); } +#[test] +#[cfg(any(feature = "gen-llvm"))] +fn list_intersperse() { + assert_evals_to!( + indoc!( + r#" + List.intersperse [0, 0, 0] 1 + "# + ), + RocList::from_slice(&[0, 1, 0, 1, 0]), + RocList + ); + assert_evals_to!( + indoc!( + r#" + List.intersperse [] 1 + "# + ), + RocList::from_slice(&[]), + RocList + ); +} + #[test] #[cfg(any(feature = "gen-llvm"))] fn list_drop_at_shared() {