diff --git a/compiler/builtins/bitcode/src/list.zig b/compiler/builtins/bitcode/src/list.zig index d012e92e26..1cc2f36f89 100644 --- a/compiler/builtins/bitcode/src/list.zig +++ b/compiler/builtins/bitcode/src/list.zig @@ -870,6 +870,9 @@ pub fn listSublist( len: usize, dec: Dec, ) callconv(.C) RocList { + if (len == 0) { + return RocList.empty(); + } if (list.bytes) |source_ptr| { const size = list.len(); diff --git a/compiler/builtins/src/std.rs b/compiler/builtins/src/std.rs index 0f8d1fe154..41ed828beb 100644 --- a/compiler/builtins/src/std.rs +++ b/compiler/builtins/src/std.rs @@ -1015,6 +1015,25 @@ pub fn types() -> MutMap { Box::new(list_type(flex(TVAR1))), ); + // split : List elem, Nat -> { before: List elem, others: List elem } + add_top_level_function_type!( + Symbol::LIST_SPLIT, + vec![list_type(flex(TVAR1)), nat_type(),], + Box::new(SolvedType::Record { + fields: vec![ + ( + "before".into(), + RecordField::Required(list_type(flex(TVAR1))) + ), + ( + "others".into(), + RecordField::Required(list_type(flex(TVAR1))) + ), + ], + ext: Box::new(SolvedType::EmptyRecord), + },), + ); + // drop : List elem, Nat -> List elem add_top_level_function_type!( Symbol::LIST_DROP, diff --git a/compiler/can/src/builtins.rs b/compiler/can/src/builtins.rs index 6e2342ca3e..bda105589d 100644 --- a/compiler/can/src/builtins.rs +++ b/compiler/can/src/builtins.rs @@ -1,9 +1,9 @@ use crate::def::Def; use crate::expr::{ClosureData, Expr::*}; -use crate::expr::{Expr, Recursive, WhenBranch}; +use crate::expr::{Expr, Field, Recursive, WhenBranch}; use crate::pattern::Pattern; use roc_collections::all::SendMap; -use roc_module::ident::TagName; +use roc_module::ident::{Lowercase, TagName}; use roc_module::low_level::LowLevel; use roc_module::operator::CalledVia; use roc_module::symbol::Symbol; @@ -96,6 +96,7 @@ pub fn builtin_defs_map(symbol: Symbol, var_store: &mut VarStore) -> Option LIST_TAKE_FIRST => list_take_first, LIST_TAKE_LAST => list_take_last, LIST_SUBLIST => list_sublist, + LIST_SPLIT => list_split, LIST_DROP => list_drop, LIST_DROP_AT => list_drop_at, LIST_DROP_FIRST => list_drop_first, @@ -2147,6 +2148,116 @@ fn list_sublist(symbol: Symbol, var_store: &mut VarStore) -> Def { ) } +/// 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(); + let index_var = var_store.fresh(); + + let list_sym = Symbol::ARG_1; + let index_sym = Symbol::ARG_2; + + let clos_sym = Symbol::LIST_SPLIT_CLOS; + let clos_start_sym = Symbol::ARG_3; + let clos_len_sym = Symbol::ARG_4; + + let clos_fun_var = var_store.fresh(); + let clos_start_var = var_store.fresh(); + let clos_len_var = var_store.fresh(); + let clos_ret_var = var_store.fresh(); + + let ret_var = var_store.fresh(); + let zero = int(index_var, Variable::NATURAL, 0); + + let clos = Closure(ClosureData { + function_type: clos_fun_var, + closure_type: var_store.fresh(), + closure_ext_var: var_store.fresh(), + return_type: clos_ret_var, + name: clos_sym, + recursive: Recursive::NotRecursive, + captured_symbols: vec![(list_sym, clos_ret_var)], + arguments: vec![ + ( + clos_start_var, + no_region(Pattern::Identifier(clos_start_sym)), + ), + (clos_len_var, no_region(Pattern::Identifier(clos_len_sym))), + ], + loc_body: { + Box::new(no_region(RunLowLevel { + op: LowLevel::ListSublist, + args: vec![ + (clos_ret_var, Var(list_sym)), + (clos_start_var, Var(clos_start_sym)), + (clos_len_var, Var(clos_len_sym)), + ], + ret_var: clos_ret_var, + })) + }, + }); + + let fun = Box::new(( + clos_fun_var, + no_region(clos), + var_store.fresh(), + clos_ret_var, + )); + + let get_before = Call( + fun.clone(), + vec![ + (index_var, no_region(zero)), + (index_var, no_region(Var(index_sym))), + ], + CalledVia::Space, + ); + + let get_list_len = RunLowLevel { + op: LowLevel::ListLen, + args: vec![(list_var, Var(list_sym))], + ret_var: index_var, + }; + + let get_others_len = RunLowLevel { + op: LowLevel::NumSubWrap, + args: vec![(index_var, get_list_len), (index_var, Var(index_sym))], + ret_var: index_var, + }; + + let get_others = Call( + fun, + vec![ + (index_var, no_region(Var(index_sym))), + (index_var, no_region(get_others_len)), + ], + CalledVia::Space, + ); + + let before = Field { + var: clos_ret_var, + region: Region::zero(), + loc_expr: Box::new(no_region(get_before)), + }; + let others = Field { + var: clos_ret_var, + region: Region::zero(), + loc_expr: Box::new(no_region(get_others)), + }; + + let body = record( + vec![("before".into(), before), ("others".into(), others)], + var_store, + ); + + defn( + symbol, + vec![(list_var, list_sym), (index_var, index_sym)], + var_store, + body, + ret_var, + ) +} + /// List.drop : List elem, Nat -> List elem fn list_drop(symbol: Symbol, var_store: &mut VarStore) -> Def { let list_var = var_store.fresh(); @@ -4309,17 +4420,17 @@ fn tag(name: &'static str, args: Vec, var_store: &mut VarStore) -> Expr { } } -// #[inline(always)] -// fn record(fields: Vec<(Lowercase, Field)>, var_store: &mut VarStore) -> Expr { -// let mut send_map = SendMap::default(); -// for (k, v) in fields { -// send_map.insert(k, v); -// } -// Expr::Record { -// record_var: var_store.fresh(), -// fields: send_map, -// } -// } +#[inline(always)] +fn record(fields: Vec<(Lowercase, Field)>, var_store: &mut VarStore) -> Expr { + let mut send_map = SendMap::default(); + for (k, v) in fields { + send_map.insert(k, v); + } + Expr::Record { + record_var: var_store.fresh(), + fields: send_map, + } +} #[inline(always)] fn defn( diff --git a/compiler/module/src/symbol.rs b/compiler/module/src/symbol.rs index 44bf51012a..65b8d366cf 100644 --- a/compiler/module/src/symbol.rs +++ b/compiler/module/src/symbol.rs @@ -1072,6 +1072,8 @@ 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" } 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 dcf74d79bb..1e12191375 100644 --- a/compiler/solve/tests/solve_expr.rs +++ b/compiler/solve/tests/solve_expr.rs @@ -3793,6 +3793,14 @@ mod solve_expr { ); } + #[test] + fn list_split() { + infer_eq_without_problem( + indoc!("List.split"), + "List a, Nat -> { before : List a, others : List a }", + ); + } + #[test] fn list_drop_last() { infer_eq_without_problem( diff --git a/compiler/test_gen/src/gen_list.rs b/compiler/test_gen/src/gen_list.rs index 5e0f737e72..86aba0e539 100644 --- a/compiler/test_gen/src/gen_list.rs +++ b/compiler/test_gen/src/gen_list.rs @@ -248,6 +248,47 @@ fn list_sublist() { ); } +#[test] +#[cfg(any(feature = "gen-llvm"))] +fn list_split() { + assert_evals_to!( + r#" + list = List.split [1, 2, 3] 0 + list.before + "#, + RocList::from_slice(&[]), + RocList + ); + assert_evals_to!( + r#" + list = List.split [1, 2, 3] 0 + list.others + "#, + RocList::from_slice(&[1, 2, 3]), + RocList + ); + + assert_evals_to!( + "List.split [1, 2, 3] 1", + (RocList::from_slice(&[1]), RocList::from_slice(&[2, 3]),), + (RocList, RocList,) + ); + assert_evals_to!( + "List.split [1, 2, 3] 3", + (RocList::from_slice(&[1, 2, 3]), RocList::from_slice(&[]),), + (RocList, RocList,) + ); + assert_evals_to!( + "List.split [1, 2, 3] 4", + (RocList::from_slice(&[1, 2, 3]), RocList::from_slice(&[]),), + (RocList, RocList,) + ); + assert_evals_to!( + "List.split [] 1", + (RocList::from_slice(&[]), RocList::from_slice(&[]),), + (RocList, RocList,) + ); +} #[test] #[cfg(any(feature = "gen-llvm"))] fn list_drop() {