diff --git a/compiler/builtins/bitcode/src/list.zig b/compiler/builtins/bitcode/src/list.zig index 52cf7aaf42..ffdbfcba2e 100644 --- a/compiler/builtins/bitcode/src/list.zig +++ b/compiler/builtins/bitcode/src/list.zig @@ -888,6 +888,34 @@ pub fn listTakeFirst( } } +pub fn listTakeLast( + list: RocList, + alignment: u32, + element_width: usize, + take_count: usize, + dec: Dec, +) callconv(.C) RocList { + if (take_count == 0) { + return RocList.empty(); + } + if (list.bytes) |source_ptr| { + const size = list.len(); + if (size <= take_count) { + return list; + } + const drop_count = size - take_count; + return listDrop( + list, + alignment, + element_width, + drop_count, + dec, + ); + } else { + return RocList.empty(); + } +} + pub fn listDrop( list: RocList, alignment: u32, diff --git a/compiler/builtins/bitcode/src/main.zig b/compiler/builtins/bitcode/src/main.zig index e7b8fb4814..ba063b0702 100644 --- a/compiler/builtins/bitcode/src/main.zig +++ b/compiler/builtins/bitcode/src/main.zig @@ -46,6 +46,7 @@ comptime { exportListFn(list.listSortWith, "sort_with"); exportListFn(list.listConcat, "concat"); exportListFn(list.listTakeFirst, "take_first"); + exportListFn(list.listTakeLast, "take_last"); exportListFn(list.listDrop, "drop"); exportListFn(list.listDropAt, "drop_at"); exportListFn(list.listSet, "set"); diff --git a/compiler/builtins/src/bitcode.rs b/compiler/builtins/src/bitcode.rs index 9a40655452..1ee6c97133 100644 --- a/compiler/builtins/src/bitcode.rs +++ b/compiler/builtins/src/bitcode.rs @@ -178,6 +178,7 @@ pub const LIST_REPEAT: &str = "roc_builtins.list.repeat"; pub const LIST_APPEND: &str = "roc_builtins.list.append"; pub const LIST_PREPEND: &str = "roc_builtins.list.prepend"; pub const LIST_TAKE_FIRST: &str = "roc_builtins.list.take_first"; +pub const LIST_TAKE_LAST: &str = "roc_builtins.list.take_last"; pub const LIST_DROP: &str = "roc_builtins.list.drop"; pub const LIST_DROP_AT: &str = "roc_builtins.list.drop_at"; pub const LIST_SWAP: &str = "roc_builtins.list.swap"; diff --git a/compiler/builtins/src/std.rs b/compiler/builtins/src/std.rs index 46b47ede2b..700609d936 100644 --- a/compiler/builtins/src/std.rs +++ b/compiler/builtins/src/std.rs @@ -978,6 +978,13 @@ pub fn types() -> MutMap { Box::new(list_type(flex(TVAR1))), ); + // takeLast : List elem, Nat -> List elem + add_top_level_function_type!( + Symbol::LIST_TAKE_LAST, + vec![list_type(flex(TVAR1)), nat_type()], + Box::new(list_type(flex(TVAR1))), + ); + // 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 b59f57f26e..2a3cebae1e 100644 --- a/compiler/can/src/builtins.rs +++ b/compiler/can/src/builtins.rs @@ -92,6 +92,7 @@ pub fn builtin_defs_map(symbol: Symbol, var_store: &mut VarStore) -> Option LIST_MAP3 => list_map3, LIST_MAP4 => list_map4, LIST_TAKE_FIRST => list_take_first, + LIST_TAKE_LAST => list_take_last, LIST_DROP => list_drop, LIST_DROP_AT => list_drop_at, LIST_DROP_FIRST => list_drop_first, @@ -2029,6 +2030,29 @@ fn list_take_first(symbol: Symbol, var_store: &mut VarStore) -> Def { ) } +/// List.takeLast : List elem, Nat -> List elem +fn list_take_last(symbol: Symbol, var_store: &mut VarStore) -> Def { + let list_var = var_store.fresh(); + let len_var = var_store.fresh(); + + let body = RunLowLevel { + op: LowLevel::ListTakeLast, + args: vec![ + (list_var, Var(Symbol::ARG_1)), + (len_var, Var(Symbol::ARG_2)), + ], + ret_var: list_var, + }; + + defn( + symbol, + vec![(list_var, Symbol::ARG_1), (len_var, Symbol::ARG_2)], + var_store, + body, + list_var, + ) +} + /// List.drop : List elem, Nat -> List elem fn list_drop(symbol: Symbol, var_store: &mut VarStore) -> Def { let list_var = var_store.fresh(); diff --git a/compiler/gen_llvm/src/llvm/build.rs b/compiler/gen_llvm/src/llvm/build.rs index 535bed30c3..c5221f2890 100644 --- a/compiler/gen_llvm/src/llvm/build.rs +++ b/compiler/gen_llvm/src/llvm/build.rs @@ -12,7 +12,7 @@ use crate::llvm::build_list::{ list_contains, list_drop, list_drop_at, list_get_unsafe, list_join, list_keep_errs, list_keep_if, list_keep_oks, list_len, list_map, list_map2, list_map3, list_map4, list_map_with_index, list_prepend, list_range, list_repeat, list_reverse, list_set, - list_single, list_sort_with, list_swap, list_take_first, + list_single, list_sort_with, list_swap, list_take_first, list_take_last, }; use crate::llvm::build_str::{ empty_str, str_concat, str_count_graphemes, str_ends_with, str_from_float, str_from_int, @@ -5172,6 +5172,27 @@ fn run_low_level<'a, 'ctx, 'env>( _ => unreachable!("Invalid layout {:?} in List.takeFirst", list_layout), } } + ListTakeLast => { + // List.takeLast : 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_take_last( + env, + layout_ids, + original_wrapper, + count.into_int_value(), + element_layout, + ), + _ => unreachable!("Invalid layout {:?} in List.takeLast", list_layout), + } + } ListDrop => { // List.drop : List elem, Nat -> List elem debug_assert_eq!(args.len(), 2); diff --git a/compiler/gen_llvm/src/llvm/build_list.rs b/compiler/gen_llvm/src/llvm/build_list.rs index b5d2cd2288..e3f0b00ff4 100644 --- a/compiler/gen_llvm/src/llvm/build_list.rs +++ b/compiler/gen_llvm/src/llvm/build_list.rs @@ -316,6 +316,28 @@ pub fn list_take_first<'a, 'ctx, 'env>( ) } +/// List.takeLast : List elem, Nat -> List elem +pub fn list_take_last<'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_cc(env, original_wrapper.into()), + env.alignment_intvalue(element_layout), + layout_width(env, element_layout), + count.into(), + dec_element_fn.as_global_value().as_pointer_value().into(), + ], + bitcode::LIST_TAKE_LAST, + ) +} + /// List.drop : List elem, Nat -> List elem pub fn list_drop<'a, 'ctx, 'env>( env: &Env<'a, 'ctx, 'env>, diff --git a/compiler/module/src/low_level.rs b/compiler/module/src/low_level.rs index ed6856fbf6..0c0be62ed9 100644 --- a/compiler/module/src/low_level.rs +++ b/compiler/module/src/low_level.rs @@ -43,6 +43,7 @@ pub enum LowLevel { ListKeepErrs, ListSortWith, ListTakeFirst, + ListTakeLast, ListDrop, ListDropAt, ListSwap, @@ -133,6 +134,7 @@ macro_rules! first_order { | ListGetUnsafe | ListSet | ListTakeFirst + | ListTakeLast | ListDrop | ListDropAt | ListSingle diff --git a/compiler/module/src/symbol.rs b/compiler/module/src/symbol.rs index 8c012a249a..5306e59b93 100644 --- a/compiler/module/src/symbol.rs +++ b/compiler/module/src/symbol.rs @@ -1066,6 +1066,7 @@ define_builtins! { 43 LIST_JOIN_MAP_CONCAT: "#joinMapConcat" 44 LIST_ANY: "any" 45 LIST_TAKE_FIRST: "takeFirst" + 46 LIST_TAKE_LAST: "takeLast" } 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 2f9cabd61f..ab6fccb595 100644 --- a/compiler/mono/src/borrow.rs +++ b/compiler/mono/src/borrow.rs @@ -964,6 +964,7 @@ pub fn lowlevel_borrow_signature(arena: &Bump, op: LowLevel) -> &[bool] { // List.append should own its first argument ListAppend => arena.alloc_slice_copy(&[owned, owned]), ListTakeFirst => arena.alloc_slice_copy(&[owned, irrelevant]), + ListTakeLast => arena.alloc_slice_copy(&[owned, irrelevant]), ListDrop => arena.alloc_slice_copy(&[owned, irrelevant]), ListDropAt => arena.alloc_slice_copy(&[owned, irrelevant]), ListSwap => arena.alloc_slice_copy(&[owned, irrelevant, irrelevant]), diff --git a/compiler/mono/src/low_level.rs b/compiler/mono/src/low_level.rs index 8bfbb81ebe..8546ec2c4e 100644 --- a/compiler/mono/src/low_level.rs +++ b/compiler/mono/src/low_level.rs @@ -97,6 +97,7 @@ enum FirstOrder { ListGetUnsafe, ListSet, ListTakeFirst, + ListTakeLast, ListDrop, ListDropAt, ListSingle, diff --git a/compiler/solve/tests/solve_expr.rs b/compiler/solve/tests/solve_expr.rs index c5a8403d5d..fac5c31b4d 100644 --- a/compiler/solve/tests/solve_expr.rs +++ b/compiler/solve/tests/solve_expr.rs @@ -3757,6 +3757,18 @@ mod solve_expr { ); } + #[test] + fn list_take_last() { + infer_eq_without_problem( + indoc!( + r#" + List.takeLast + "# + ), + "List a, Nat -> 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 835343f3ca..4f70651359 100644 --- a/compiler/test_gen/src/gen_list.rs +++ b/compiler/test_gen/src/gen_list.rs @@ -211,6 +211,26 @@ fn list_take_first() { ); } +#[test] +fn list_take_last() { + assert_evals_to!( + "List.takeLast [1, 2, 3] 2", + RocList::from_slice(&[2, 3]), + RocList + ); + assert_evals_to!( + "List.takeLast [1, 2, 3] 0", + RocList::from_slice(&[]), + RocList + ); + assert_evals_to!("List.takeLast [] 1", RocList::from_slice(&[]), RocList); + assert_evals_to!( + "List.takeLast [1,2] 5", + RocList::from_slice(&[1, 2]), + RocList + ); +} + #[test] fn list_drop() { assert_evals_to!(