diff --git a/compiler/builtins/bitcode/src/list.zig b/compiler/builtins/bitcode/src/list.zig index 1791a8f61b..bd31d71951 100644 --- a/compiler/builtins/bitcode/src/list.zig +++ b/compiler/builtins/bitcode/src/list.zig @@ -665,6 +665,42 @@ pub fn listAppend(list: RocList, alignment: usize, element: Opaque, element_widt return output; } +pub fn listDrop( + list: RocList, + alignment: usize, + element_width: usize, + drop_count: usize, + dec: Dec, +) callconv(.C) RocList { + if (list.bytes) |source_ptr| { + const size = list.len(); + const keep_count = size - drop_count; + + var i: usize = 0; + const iterations = std.math.min(drop_count, size); + + while (i < iterations) : (i += 1) { + const element = source_ptr + i * element_width; + dec(element); + } + + if (drop_count >= size) { + return RocList.empty(); + } + + const output = RocList.allocate(std.heap.c_allocator, alignment, keep_count, element_width); + const target_ptr = output.bytes orelse unreachable; + + @memcpy(target_ptr, source_ptr + drop_count * element_width, keep_count * element_width); + + utils.decref(std.heap.c_allocator, alignment, list.bytes, size * element_width); + + return output; + } else { + return RocList.empty(); + } +} + pub fn listRange(width: utils.IntWidth, low: Opaque, high: Opaque) callconv(.C) RocList { const allocator = std.heap.c_allocator; const IntWidth = utils.IntWidth; diff --git a/compiler/builtins/bitcode/src/main.zig b/compiler/builtins/bitcode/src/main.zig index 18db41131c..6eb689bd66 100644 --- a/compiler/builtins/bitcode/src/main.zig +++ b/compiler/builtins/bitcode/src/main.zig @@ -25,6 +25,7 @@ comptime { exportListFn(list.listReverse, "reverse"); exportListFn(list.listSortWith, "sort_with"); exportListFn(list.listConcat, "concat"); + exportListFn(list.listDrop, "drop"); } // Dict Module diff --git a/compiler/builtins/src/bitcode.rs b/compiler/builtins/src/bitcode.rs index 2032e13700..7c3d82dc4d 100644 --- a/compiler/builtins/src/bitcode.rs +++ b/compiler/builtins/src/bitcode.rs @@ -76,6 +76,7 @@ pub const LIST_WALK_BACKWARDS: &str = "roc_builtins.list.walk_backwards"; pub const LIST_CONTAINS: &str = "roc_builtins.list.contains"; pub const LIST_REPEAT: &str = "roc_builtins.list.repeat"; pub const LIST_APPEND: &str = "roc_builtins.list.append"; +pub const LIST_DROP: &str = "roc_builtins.list.drop"; pub const LIST_SINGLE: &str = "roc_builtins.list.single"; pub const LIST_JOIN: &str = "roc_builtins.list.join"; pub const LIST_RANGE: &str = "roc_builtins.list.range"; diff --git a/compiler/builtins/src/std.rs b/compiler/builtins/src/std.rs index c0d8f9b39e..25d057cc81 100644 --- a/compiler/builtins/src/std.rs +++ b/compiler/builtins/src/std.rs @@ -848,6 +848,13 @@ pub fn types() -> MutMap { Box::new(list_type(flex(TVAR1))), ); + // drop : List elem, Nat -> List elem + add_top_level_function_type!( + Symbol::LIST_DROP, + vec![list_type(flex(TVAR1)), nat_type()], + Box::new(list_type(flex(TVAR1))), + ); + // prepend : List elem, elem -> List elem add_top_level_function_type!( Symbol::LIST_PREPEND, diff --git a/compiler/can/src/builtins.rs b/compiler/can/src/builtins.rs index edd6bb49d2..7283316ba9 100644 --- a/compiler/can/src/builtins.rs +++ b/compiler/can/src/builtins.rs @@ -84,6 +84,7 @@ pub fn builtin_defs_map(symbol: Symbol, var_store: &mut VarStore) -> Option LIST_MAP => list_map, LIST_MAP2 => list_map2, LIST_MAP3 => list_map3, + LIST_DROP => list_drop, LIST_MAP_WITH_INDEX => list_map_with_index, LIST_KEEP_IF => list_keep_if, LIST_KEEP_OKS => list_keep_oks, @@ -1881,6 +1882,28 @@ fn list_set(symbol: Symbol, var_store: &mut VarStore) -> Def { list_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(); + let index_var = var_store.fresh(); + + let body = RunLowLevel { + op: LowLevel::ListDrop, + args: vec![ + (list_var, Var(Symbol::ARG_1)), + (index_var, Var(Symbol::ARG_2)), + ], + ret_var: list_var, + }; + + defn( + symbol, + vec![(list_var, Symbol::ARG_1), (index_var, Symbol::ARG_2)], + var_store, + body, + list_var, + ) +} /// List.append : List elem, elem -> List elem fn list_append(symbol: Symbol, var_store: &mut VarStore) -> Def { diff --git a/compiler/gen/src/llvm/build.rs b/compiler/gen/src/llvm/build.rs index 3a6957de5b..4b20416264 100644 --- a/compiler/gen/src/llvm/build.rs +++ b/compiler/gen/src/llvm/build.rs @@ -6,9 +6,9 @@ use crate::llvm::build_dict::{ use crate::llvm::build_hash::generic_hash; use crate::llvm::build_list::{ allocate_list, empty_list, empty_polymorphic_list, list_append, list_concat, list_contains, - list_get_unsafe, list_join, list_keep_errs, list_keep_if, list_keep_oks, list_len, list_map, - list_map2, list_map3, list_map_with_index, list_prepend, list_range, list_repeat, list_reverse, - list_set, list_single, list_sort_with, list_walk_help, + list_drop, list_get_unsafe, list_join, list_keep_errs, list_keep_if, list_keep_oks, list_len, + list_map, list_map2, list_map3, list_map_with_index, list_prepend, list_range, list_repeat, + list_reverse, list_set, list_single, list_sort_with, list_walk_help, }; use crate::llvm::build_str::{ empty_str, str_concat, str_count_graphemes, str_ends_with, str_from_float, str_from_int, @@ -3883,6 +3883,27 @@ fn run_low_level<'a, 'ctx, 'env>( list_append(env, inplace, original_wrapper, elem, elem_layout) } + ListDrop => { + // List.drop : List elem, Nat -> List elem + debug_assert_eq!(args.len(), 2); + + let (list, list_layout) = load_symbol_and_layout(scope, &args[0]); + let original_wrapper = list.into_struct_value(); + + let count = load_symbol(scope, &args[1]); + + match list_layout { + Layout::Builtin(Builtin::EmptyList) => empty_list(env), + Layout::Builtin(Builtin::List(_, element_layout)) => list_drop( + env, + layout_ids, + original_wrapper, + count.into_int_value(), + element_layout, + ), + _ => unreachable!("Invalid layout {:?} in List.drop", list_layout), + } + } ListPrepend => { // List.prepend : List elem, elem -> List elem debug_assert_eq!(args.len(), 2); diff --git a/compiler/gen/src/llvm/build_list.rs b/compiler/gen/src/llvm/build_list.rs index 6634c51862..1637f33c8a 100644 --- a/compiler/gen/src/llvm/build_list.rs +++ b/compiler/gen/src/llvm/build_list.rs @@ -322,6 +322,28 @@ pub fn list_append<'a, 'ctx, 'env>( ) } +/// List.drop : List elem, Nat -> List elem +pub fn list_drop<'a, 'ctx, 'env>( + env: &Env<'a, 'ctx, 'env>, + layout_ids: &mut LayoutIds<'a>, + original_wrapper: StructValue<'ctx>, + count: IntValue<'ctx>, + element_layout: &Layout<'a>, +) -> BasicValueEnum<'ctx> { + let dec_element_fn = build_dec_wrapper(env, layout_ids, &element_layout); + call_bitcode_fn_returns_list( + env, + &[ + pass_list_as_i128(env, original_wrapper.into()), + alignment_intvalue(env, &element_layout), + layout_width(env, &element_layout), + count.into(), + dec_element_fn.as_global_value().as_pointer_value().into(), + ], + &bitcode::LIST_DROP, + ) +} + /// List.set : List elem, Int, elem -> List elem pub fn list_set<'a, 'ctx, 'env>( parent: FunctionValue<'ctx>, diff --git a/compiler/gen_dev/tests/gen_num.rs b/compiler/gen_dev/tests/gen_num.rs index 80324a3f69..68b2dca31c 100644 --- a/compiler/gen_dev/tests/gen_num.rs +++ b/compiler/gen_dev/tests/gen_num.rs @@ -128,9 +128,7 @@ mod gen_num { af = 31 ag = 32 - # This can't be one line because it causes a stack overflow in the frontend :( - tmp = a + b + c + d + e + f + g + h + i + j + k + l + m + n + o + p + q - tmp + r + s + t + u + v + w + x + y + z + aa + ab + ac + ad + ae + af + ag + a + b + c + d + e + f + g + h + i + j + k + l + m + n + o + p + q + r + s + t + u + v + w + x + y + z + aa + ab + ac + ad + ae + af + ag "# ), 528, diff --git a/compiler/module/src/low_level.rs b/compiler/module/src/low_level.rs index cf6db1f394..65c8f93084 100644 --- a/compiler/module/src/low_level.rs +++ b/compiler/module/src/low_level.rs @@ -39,6 +39,7 @@ pub enum LowLevel { ListKeepOks, ListKeepErrs, ListSortWith, + ListDrop, DictSize, DictEmpty, DictInsert, diff --git a/compiler/module/src/symbol.rs b/compiler/module/src/symbol.rs index ef16a3b778..cb4fd13387 100644 --- a/compiler/module/src/symbol.rs +++ b/compiler/module/src/symbol.rs @@ -928,6 +928,7 @@ define_builtins! { 29 LIST_WALK_UNTIL: "walkUntil" 30 LIST_RANGE: "range" 31 LIST_SORT_WITH: "sortWith" + 32 LIST_DROP: "drop" } 5 RESULT: "Result" => { 0 RESULT_RESULT: "Result" imported // the Result.Result type alias diff --git a/compiler/mono/src/borrow.rs b/compiler/mono/src/borrow.rs index d4cdea36f9..de3a7ca7f5 100644 --- a/compiler/mono/src/borrow.rs +++ b/compiler/mono/src/borrow.rs @@ -665,6 +665,7 @@ pub fn lowlevel_borrow_signature(arena: &Bump, op: LowLevel) -> &[bool] { // TODO when we have lists with capacity (if ever) // List.append should own its first argument ListAppend => arena.alloc_slice_copy(&[owned, owned]), + ListDrop => arena.alloc_slice_copy(&[owned, irrelevant]), Eq | NotEq => arena.alloc_slice_copy(&[borrowed, borrowed]), diff --git a/compiler/test_gen/src/gen_list.rs b/compiler/test_gen/src/gen_list.rs index 124f719a6b..5e9d02b031 100644 --- a/compiler/test_gen/src/gen_list.rs +++ b/compiler/test_gen/src/gen_list.rs @@ -147,6 +147,17 @@ fn list_append() { ); } +#[test] +fn list_drop() { + assert_evals_to!( + "List.drop [1,2,3] 2", + RocList::from_slice(&[3]), + RocList + ); + assert_evals_to!("List.drop [] 1", RocList::from_slice(&[]), RocList); + assert_evals_to!("List.drop [1,2] 5", RocList::from_slice(&[]), RocList); +} + #[test] fn list_append_to_empty_list() { assert_evals_to!("List.append [] 3", RocList::from_slice(&[3]), RocList);