Merge pull request #2002 from rtfeldman/builtins-list-intersperse

Add builtin `List.intersperse`

compiler/builtins/src/std.rs

@@ -1159,6 +1159,13 @@ pub fn types() -> MutMap<Symbol, (SolvedType, Region)> {
         )
     }
 
+    // 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

compiler/can/src/builtins.rs

@@ -97,6 +97,7 @@ pub fn builtin_defs_map(symbol: Symbol, var_store: &mut VarStore) -> Option<Def>
         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();

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"
+        50 LIST_INTERSPERSE: "intersperse"
-        51 LIST_SPLIT_CLOS: "#splitClos"
+        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

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

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<i64>);
 }
 
+#[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<i64>
+    );
+    assert_evals_to!(
+        indoc!(
+            r#"
+                    List.intersperse [] 1
+                "#
+        ),
+        RocList::from_slice(&[]),
+        RocList<i64>
+    );
+}
+
 #[test]
 #[cfg(any(feature = "gen-llvm"))]
 fn list_drop_at_shared() {