List.keepIf in zig

compiler/builtins/bitcode/src/list.zig

@@ -46,7 +46,7 @@ pub const RocList = extern struct {
         };
     }
 
-    pub fn makeUnique(self: RocDict, allocator: *Allocator, alignment: Alignment, key_width: usize, value_width: usize) RocDict {
+    pub fn makeUnique(self: RocList, allocator: *Allocator, alignment: usize, element_width: usize) RocList {
         if (self.isEmpty()) {
             return self;
         }
@@ -56,28 +56,28 @@ pub const RocList = extern struct {
         }
 
         // unfortunately, we have to clone
-        var new_dict = RocDict.allocate(allocator, self.number_of_levels, self.dict_entries_len, alignment, key_width, value_width);
+        var new_list = RocList.allocate(allocator, self.length, alignment, element_width);
 
-        var old_bytes: [*]u8 = @ptrCast([*]u8, self.dict_bytes);
+        var old_bytes: [*]u8 = @ptrCast([*]u8, self.bytes);
-        var new_bytes: [*]u8 = @ptrCast([*]u8, new_dict.dict_bytes);
+        var new_bytes: [*]u8 = @ptrCast([*]u8, new_list.bytes);
 
-        const number_of_bytes = self.capacity() * (@sizeOf(Slot) + key_width + value_width);
+        const number_of_bytes = self.len() * element_width;
         @memcpy(new_bytes, old_bytes, number_of_bytes);
 
         // NOTE we fuse an increment of all keys/values with a decrement of the input dict
-        const data_bytes = self.capacity() * slotSize(key_width, value_width);
+        const data_bytes = self.len() * element_width;
-        decref(allocator, alignment, self.dict_bytes, data_bytes);
+        utils.decref(allocator, alignment, self.bytes, data_bytes);
 
-        return new_dict;
+        return new_list;
     }
 
     pub fn reallocate(
-        self: RocDict,
+        self: RocList,
         allocator: *Allocator,
-        alignment: Alignment,
+        alignment: usize,
         new_length: usize,
         element_width: usize,
-    ) RocDict {
+    ) RocList {
         const old_length = self.length;
         const delta_length = new_length - old_length;
 
@@ -106,8 +106,8 @@ pub const RocList = extern struct {
     }
 };
 
-const ListMapCaller = fn (?[*]u8, ?[*]u8, ?[*]u8) callconv(.C) void;
+const Caller1 = fn (?[*]u8, ?[*]u8, ?[*]u8) callconv(.C) void;
-pub fn listMap(list: RocList, transform: Opaque, caller: ListMapCaller, alignment: usize, old_element_width: usize, new_element_width: usize) callconv(.C) RocList {
+pub fn listMap(list: RocList, transform: Opaque, caller: Caller1, alignment: usize, old_element_width: usize, new_element_width: usize) callconv(.C) RocList {
     if (list.bytes) |source_ptr| {
         const size = list.len();
         var i: usize = 0;
@@ -125,3 +125,35 @@ pub fn listMap(list: RocList, transform: Opaque, caller: ListMapCaller, alignmen
         return RocList.empty();
     }
 }
+
+pub fn listKeepIf(list: RocList, transform: Opaque, caller: Caller1, alignment: usize, element_width: usize) callconv(.C) RocList {
+    if (list.bytes) |source_ptr| {
+        const size = list.len();
+        var i: usize = 0;
+        var output = list.makeUnique(std.heap.c_allocator, alignment, list.len() * element_width);
+        const target_ptr = output.bytes orelse unreachable;
+
+        var kept: usize = 0;
+        while (i < size) : (i += 1) {
+            var keep = false;
+            const element = source_ptr + (i * element_width);
+            caller(transform, element, @ptrCast(?[*]u8, &keep));
+
+            if (keep) {
+                @memcpy(target_ptr + (kept * element_width), element, element_width);
+
+                kept += 1;
+            } else {
+                // TODO decrement the value?
+            }
+        }
+
+        output.length = kept;
+
+        // utils.decref(std.heap.c_allocator, alignment, list.bytes, size * old_element_width);
+
+        return output;
+    } else {
+        return RocList.empty();
+    }
+}

compiler/builtins/bitcode/src/main.zig

@@ -7,6 +7,7 @@ const list = @import("list.zig");
 
 comptime {
     exportListFn(list.listMap, "map");
+    exportListFn(list.listKeepIf, "keep_if");
 }
 
 // Dict Module

compiler/builtins/src/bitcode.rs

@@ -60,3 +60,4 @@ pub const DICT_WALK: &str = "roc_builtins.dict.walk";
 pub const SET_FROM_LIST: &str = "roc_builtins.dict.set_from_list";
 
 pub const LIST_MAP: &str = "roc_builtins.list.map";
+pub const LIST_KEEP_IF: &str = "roc_builtins.list.keep_if";

compiler/gen/src/llvm/build.rs

@@ -3632,8 +3632,7 @@ fn run_low_level<'a, 'ctx, 'env>(
 
             match list_layout {
                 Layout::Builtin(Builtin::EmptyList) => {
-                    // no elements, so `key` is not in here
+                    return empty_list(env);
-                    panic!("key type unknown")
                 }
                 Layout::Builtin(Builtin::List(_, element_layout)) => {
                     list_map(env, layout_ids, func, func_layout, list, element_layout)
@@ -3651,16 +3650,15 @@ fn run_low_level<'a, 'ctx, 'env>(
 
             let inplace = get_inplace_from_layout(layout);
 
-            list_keep_if(
+            match list_layout {
-                env,
+                Layout::Builtin(Builtin::EmptyList) => {
-                layout_ids,
+                    return empty_list(env);
-                inplace,
+                }
-                parent,
+                Layout::Builtin(Builtin::List(_, element_layout)) => {
-                func,
+                    list_keep_if(env, layout_ids, func, func_layout, list, element_layout)
-                func_layout,
+                }
-                list,
+                _ => unreachable!("invalid list layout"),
-                list_layout,
+            }
-            )
         }
         ListContains => {
             // List.contains : List elem, elem -> Bool

compiler/gen/src/llvm/build_list.rs

@@ -1210,216 +1210,51 @@ pub fn list_contains_help<'a, 'ctx, 'env>(
 pub fn list_keep_if<'a, 'ctx, 'env>(
     env: &Env<'a, 'ctx, 'env>,
     layout_ids: &mut LayoutIds<'a>,
-    output_inplace: InPlace,
+    transform: BasicValueEnum<'ctx>,
-    parent: FunctionValue<'ctx>,
+    transform_layout: &Layout<'a>,
-    func: BasicValueEnum<'ctx>,
-    func_layout: &Layout<'a>,
     list: BasicValueEnum<'ctx>,
-    list_layout: &Layout<'a>,
+    element_layout: &Layout<'a>,
 ) -> BasicValueEnum<'ctx> {
     let builder = env.builder;
-    let ctx = env.context;
 
-    let wrapper_struct = list.into_struct_value();
+    let u8_ptr = env.context.i8_type().ptr_type(AddressSpace::Generic);
-    let (input_inplace, element_layout) = match list_layout.clone() {
-        Layout::Builtin(Builtin::EmptyList) => (
-            InPlace::InPlace,
-            // this pointer will never actually be dereferenced
-            Layout::Builtin(Builtin::Int64),
-        ),
-        Layout::Builtin(Builtin::List(memory_mode, elem_layout)) => (
-            match memory_mode {
-                MemoryMode::Unique => InPlace::InPlace,
-                MemoryMode::Refcounted => InPlace::Clone,
-            },
-            elem_layout.clone(),
-        ),
 
-        _ => unreachable!("Invalid layout {:?} in List.keepIf", list_layout),
+    let list_i128 = complex_bitcast(env.builder, list, env.context.i128_type().into(), "to_i128");
-    };
 
-    let list_type = basic_type_from_layout(env.arena, env.context, &list_layout, env.ptr_bytes);
+    let transform_ptr = builder.build_alloca(transform.get_type(), "transform_ptr");
-    let elem_type = basic_type_from_layout(env.arena, env.context, &element_layout, env.ptr_bytes);
+    env.builder.build_store(transform_ptr, transform);
-    let ptr_type = elem_type.ptr_type(AddressSpace::Generic);
 
-    let list_ptr = load_list_ptr(builder, wrapper_struct, ptr_type);
+    let stepper_caller =
-    let length = list_len(builder, list.into_struct_value());
+        build_transform_caller(env, layout_ids, transform_layout, &[element_layout.clone()])
+            .as_global_value()
+            .as_pointer_value();
 
-    let zero = ctx.i64_type().const_zero();
+    let element_width = env
+        .ptr_int()
+        .const_int(element_layout.stack_size(env.ptr_bytes) as u64, false);
 
-    match input_inplace {
+    let alignment = element_layout.alignment_bytes(env.ptr_bytes);
-        InPlace::InPlace => {
+    let alignment_iv = env.ptr_int().const_int(alignment as u64, false);
-            let new_length = list_keep_if_help(
+
+    let output = call_bitcode_fn(
         env,
-                input_inplace,
+        &[
-                parent,
+            list_i128.into(),
-                length,
+            env.builder
-                list_ptr,
+                .build_bitcast(transform_ptr, u8_ptr, "to_opaque"),
-                list_ptr,
+            stepper_caller.into(),
-                func,
+            alignment_iv.into(),
-                func_layout,
+            element_width.into(),
+        ],
+        &bitcode::LIST_KEEP_IF,
     );
 
-            store_list(env, list_ptr, new_length)
+    complex_bitcast(
-        }
+        env.builder,
-        InPlace::Clone => {
+        output,
-            let len_0_block = ctx.append_basic_block(parent, "len_0_block");
+        collection(env.context, env.ptr_bytes).into(),
-            let len_n_block = ctx.append_basic_block(parent, "len_n_block");
+        "from_i128",
-            let cont_block = ctx.append_basic_block(parent, "cont_block");
+    )
-
-            let result = builder.build_alloca(list_type, "result");
-
-            builder.build_switch(length, len_n_block, &[(zero, len_0_block)]);
-
-            // build block for length 0
-            {
-                builder.position_at_end(len_0_block);
-
-                let new_list = store_list(env, ptr_type.const_zero(), zero);
-
-                builder.build_store(result, new_list);
-                builder.build_unconditional_branch(cont_block);
-            }
-
-            // build block for length > 0
-            {
-                builder.position_at_end(len_n_block);
-
-                let new_list_ptr = allocate_list(env, output_inplace, &element_layout, length);
-
-                let new_length = list_keep_if_help(
-                    env,
-                    InPlace::Clone,
-                    parent,
-                    length,
-                    list_ptr,
-                    new_list_ptr,
-                    func,
-                    func_layout,
-                );
-
-                // store new list pointer there
-                let new_list = store_list(env, new_list_ptr, new_length);
-
-                builder.build_store(result, new_list);
-                builder.build_unconditional_branch(cont_block);
-            }
-
-            builder.position_at_end(cont_block);
-
-            // consume the input list
-            decrement_refcount_layout(env, parent, layout_ids, list, list_layout);
-
-            builder.build_load(result, "load_result")
-        }
-    }
-}
-
-#[allow(clippy::too_many_arguments)]
-pub fn list_keep_if_help<'a, 'ctx, 'env>(
-    env: &Env<'a, 'ctx, 'env>,
-    _inplace: InPlace,
-    parent: FunctionValue<'ctx>,
-    length: IntValue<'ctx>,
-    source_ptr: PointerValue<'ctx>,
-    dest_ptr: PointerValue<'ctx>,
-    func: BasicValueEnum<'ctx>,
-    func_layout: &Layout<'a>,
-) -> IntValue<'ctx> {
-    match (func, func_layout) {
-        (
-            BasicValueEnum::PointerValue(func_ptr),
-            Layout::FunctionPointer(_, Layout::Builtin(Builtin::Int1)),
-        ) => {
-            let builder = env.builder;
-            let ctx = env.context;
-
-            let index_alloca = builder.build_alloca(ctx.i64_type(), "index_alloca");
-            let next_free_index_alloca =
-                builder.build_alloca(ctx.i64_type(), "next_free_index_alloca");
-
-            builder.build_store(index_alloca, ctx.i64_type().const_zero());
-            builder.build_store(next_free_index_alloca, ctx.i64_type().const_zero());
-
-            // while (length > next_index)
-            let condition_bb = ctx.append_basic_block(parent, "condition");
-            builder.build_unconditional_branch(condition_bb);
-            builder.position_at_end(condition_bb);
-
-            let index = builder.build_load(index_alloca, "index").into_int_value();
-
-            let condition = builder.build_int_compare(IntPredicate::SGT, length, index, "loopcond");
-
-            let body_bb = ctx.append_basic_block(parent, "body");
-            let cont_bb = ctx.append_basic_block(parent, "cont");
-            builder.build_conditional_branch(condition, body_bb, cont_bb);
-
-            // loop body
-            builder.position_at_end(body_bb);
-
-            let elem_ptr = unsafe { builder.build_in_bounds_gep(source_ptr, &[index], "elem_ptr") };
-
-            let elem = builder.build_load(elem_ptr, "load_elem");
-
-            let call_site_value =
-                builder.build_call(func_ptr, env.arena.alloc([elem]), "#keep_if_insert_func");
-
-            // set the calling convention explicitly for this call
-            call_site_value.set_call_convention(crate::llvm::build::FAST_CALL_CONV);
-
-            let should_keep = call_site_value
-                .try_as_basic_value()
-                .left()
-                .unwrap_or_else(|| panic!("LLVM error: Invalid call by pointer."))
-                .into_int_value();
-
-            let filter_pass_bb = ctx.append_basic_block(parent, "loop");
-            let after_filter_pass_bb = ctx.append_basic_block(parent, "after_loop");
-
-            let one = ctx.i64_type().const_int(1, false);
-
-            builder.build_conditional_branch(should_keep, filter_pass_bb, after_filter_pass_bb);
-            builder.position_at_end(filter_pass_bb);
-
-            let next_free_index = builder
-                .build_load(next_free_index_alloca, "load_next_free")
-                .into_int_value();
-
-            // TODO if next_free_index equals index, and we are mutating in place,
-            // then maybe we should not write this value back into memory
-            let dest_elem_ptr = unsafe {
-                builder.build_in_bounds_gep(dest_ptr, &[next_free_index], "dest_elem_ptr")
-            };
-
-            builder.build_store(dest_elem_ptr, elem);
-
-            builder.build_store(
-                next_free_index_alloca,
-                builder.build_int_add(next_free_index, one, "incremented_next_free_index"),
-            );
-
-            builder.build_unconditional_branch(after_filter_pass_bb);
-            builder.position_at_end(after_filter_pass_bb);
-
-            builder.build_store(
-                index_alloca,
-                builder.build_int_add(index, one, "incremented_index"),
-            );
-
-            builder.build_unconditional_branch(condition_bb);
-
-            // continuation
-            builder.position_at_end(cont_bb);
-
-            builder
-                .build_load(next_free_index_alloca, "new_length")
-                .into_int_value()
-        }
-        _ => unreachable!(
-            "Invalid function basic value enum or layout for List.keepIf : {:?}",
-            (func, func_layout)
-        ),
-    }
 }
 
 /// List.map : List before, (before -> after) -> List after