make closure layout more robust

compiler/gen/src/llvm/convert.rs

@@ -52,6 +52,7 @@ fn basic_type_from_function_layout<'ctx>(
     arena: &Bump,
     context: &'ctx Context,
     args: &[Layout<'_>],
+    closure_type: Option<BasicTypeEnum<'ctx>>,
     ret_layout: &Layout<'_>,
     ptr_bytes: u32,
 ) -> BasicTypeEnum<'ctx> {
@@ -64,6 +65,10 @@ fn basic_type_from_function_layout<'ctx>(
         ));
     }
 
+    if let Some(closure) = closure_type {
+        arg_basic_types.push(closure);
+    }
+
     let fn_type = get_fn_type(&ret_type, arg_basic_types.into_bump_slice());
     let ptr_type = fn_type.ptr_type(AddressSpace::Generic);
 
@@ -103,19 +108,27 @@ pub fn basic_type_from_layout<'ctx>(
 
     match layout {
         FunctionPointer(args, ret_layout) => {
-            basic_type_from_function_layout(arena, context, args, ret_layout, ptr_bytes)
+            basic_type_from_function_layout(arena, context, args, None, ret_layout, ptr_bytes)
         }
         Closure(args, closure_layout, ret_layout) => {
-            let args = {
+            //            let closure_data = block_of_memory(
-                let mut temp = Vec::from_iter_in(args.iter().cloned(), arena);
+            //                context,
-                temp.push(Layout::Struct(closure_layout));
+            //                // &closure_layout.into_block_of_memory_layout(),
-                temp.into_bump_slice()
+            //                &closure_layout.into_layout(),
-            };
+            //                ptr_bytes,
+            //            );
 
-            let function_pointer =
+            let closure_data =
-                basic_type_from_function_layout(arena, context, args, ret_layout, ptr_bytes);
+                basic_type_from_layout(arena, context, &closure_layout.into_layout(), ptr_bytes);
 
-            let closure_data = basic_type_from_record(arena, context, closure_layout, ptr_bytes);
+            let function_pointer = basic_type_from_function_layout(
+                arena,
+                context,
+                args,
+                Some(closure_data),
+                ret_layout,
+                ptr_bytes,
+            );
 
             context
                 .struct_type(&[function_pointer, closure_data], false)

compiler/gen/src/llvm/refcounting.rs

@@ -64,7 +64,7 @@ pub fn decrement_refcount_layout<'a, 'ctx, 'env>(
             decrement_refcount_builtin(env, parent, layout_ids, value, layout, builtin)
         }
         Closure(_, closure_layout, _) => {
-            if closure_layout.iter().any(|f| f.contains_refcounted()) {
+            if closure_layout.contains_refcounted() {
                 let wrapper_struct = value.into_struct_value();
 
                 let field_ptr = env
@@ -72,7 +72,13 @@ pub fn decrement_refcount_layout<'a, 'ctx, 'env>(
                     .build_extract_value(wrapper_struct, 1, "decrement_closure_data")
                     .unwrap();
 
-                decrement_refcount_struct(env, parent, layout_ids, field_ptr, closure_layout)
+                decrement_refcount_layout(
+                    env,
+                    parent,
+                    layout_ids,
+                    field_ptr,
+                    &closure_layout.into_layout(),
+                )
             }
         }
         Struct(layouts) => {

compiler/gen/tests/gen_primitives.rs

@@ -902,7 +902,6 @@ mod gen_primitives {
     }
 
     #[test]
-    #[ignore]
     fn closure() {
         assert_evals_to!(
             indoc!(
@@ -945,4 +944,32 @@ mod gen_primitives {
             i64
         );
     }
+
+    #[test]
+    #[ignore]
+    fn specialize_closure() {
+        assert_evals_to!(
+            indoc!(
+                r#"
+                app Test provides [ main ] imports []
+
+                foo = \{} ->
+                    x = 41
+                    y = 1
+
+                    f = \{} -> x
+                    g = \{} -> x + y
+
+                    [ f, g ]
+
+                main = 
+                    items = foo {}
+
+                    List.len items
+                "#
+            ),
+            2,
+            i64
+        );
+    }
 }

compiler/mono/src/ir.rs

@@ -2596,7 +2596,7 @@ pub fn with_hole<'a>(
                                 let closure_symbol = function_symbol;
 
                                 // layout of the closure record
-                                let closure_record_layout = Layout::Struct(closure_fields);
+                                let closure_record_layout = closure_fields.into_layout();
 
                                 let arg_symbols = {
                                     let mut temp =
@@ -2921,7 +2921,7 @@ pub fn from_can<'a>(
                                     return_type,
                                 );
 
-                                let closure_data_layout = Layout::Struct(closure_fields);
+                                let closure_data_layout = closure_fields.into_layout();
                                 // define the function pointer
                                 let function_ptr_layout = {
                                     let mut temp = Vec::from_iter_in(

compiler/mono/src/layout.rs

@@ -26,10 +26,54 @@ pub enum Layout<'a> {
     RecursivePointer,
     /// A function. The types of its arguments, then the type of its return value.
     FunctionPointer(&'a [Layout<'a>], &'a Layout<'a>),
-    Closure(&'a [Layout<'a>], &'a [Layout<'a>], &'a Layout<'a>),
+    Closure(&'a [Layout<'a>], ClosureLayout<'a>, &'a Layout<'a>),
     Pointer(&'a Layout<'a>),
 }
 
+#[derive(Clone, Debug, PartialEq, Eq, Hash)]
+pub struct ClosureLayout<'a> {
+    /// the layout that this specific closure captures
+    captured: &'a [Layout<'a>],
+
+    /// the layout that represents the maximum size the closure layout can have
+    max_size: &'a [Layout<'a>],
+}
+
+impl<'a> ClosureLayout<'a> {
+    fn from_unwrapped(layouts: &'a [Layout<'a>]) -> Self {
+        debug_assert!(layouts.len() > 0);
+        ClosureLayout {
+            captured: layouts,
+            max_size: layouts,
+        }
+    }
+
+    pub fn stack_size(&self, pointer_size: u32) -> u32 {
+        self.max_size
+            .iter()
+            .map(|l| l.stack_size(pointer_size))
+            .sum()
+    }
+    pub fn contains_refcounted(&self) -> bool {
+        self.captured.iter().any(|l| l.contains_refcounted())
+    }
+    pub fn safe_to_memcpy(&self) -> bool {
+        self.captured.iter().all(|l| l.safe_to_memcpy())
+    }
+
+    pub fn into_layout(&self) -> Layout<'a> {
+        if self.captured.len() == 1 {
+            self.captured[0].clone()
+        } else {
+            Layout::Struct(self.captured)
+        }
+    }
+
+    pub fn into_block_of_memory_layout(&self) -> Layout<'a> {
+        Layout::Struct(self.max_size)
+    }
+}
+
 #[derive(Clone, Debug, PartialEq, Eq, Hash, Copy)]
 pub enum MemoryMode {
     Unique,
@@ -140,9 +184,7 @@ impl<'a> Layout<'a> {
                 // Function pointers are immutable and can always be safely copied
                 true
             }
-            Closure(_, closure_layout, _) => {
+            Closure(_, closure_layout, _) => closure_layout.safe_to_memcpy(),
-                closure_layout.iter().all(|field| field.safe_to_memcpy())
-            }
             Pointer(_) => {
                 // We cannot memcpy pointers, because then we would have the same pointer in multiple places!
                 false
@@ -195,13 +237,7 @@ impl<'a> Layout<'a> {
                 })
                 .max()
                 .unwrap_or_default(),
-            Closure(_, closure_layout, _) => {
+            Closure(_, closure_layout, _) => pointer_size + closure_layout.stack_size(pointer_size),
-                pointer_size
-                    + closure_layout
-                        .iter()
-                        .map(|x| x.stack_size(pointer_size))
-                        .sum::<u32>()
-            }
             FunctionPointer(_, _) => pointer_size,
             RecursivePointer => pointer_size,
             Pointer(_) => pointer_size,
@@ -231,7 +267,7 @@ impl<'a> Layout<'a> {
                 .flatten()
                 .any(|f| f.is_refcounted()),
             RecursiveUnion(_) => true,
-            Closure(_, closure_layout, _) => closure_layout.iter().any(|f| f.contains_refcounted()),
+            Closure(_, closure_layout, _) => closure_layout.contains_refcounted(),
             FunctionPointer(_, _) | RecursivePointer | Pointer(_) => false,
         }
     }
@@ -498,36 +534,44 @@ fn layout_from_flat_type<'a>(
 
             let ret = Layout::from_var(env, ret_var)?;
 
-            match Layout::from_var(env, closure_var) {
+            let mut tags = std::vec::Vec::new();
-                Ok(Layout::Builtin(builtin)) => Ok(Layout::Closure(
+            match roc_types::pretty_print::chase_ext_tag_union(env.subs, closure_var, &mut tags) {
-                    fn_args.into_bump_slice(),
+                Ok(()) | Err((_, Content::FlexVar(_))) if !tags.is_empty() => {
-                    arena.alloc([Layout::Builtin(builtin.clone())]),
+                    // this is a closure
-                    arena.alloc(ret),
+                    let variant = union_sorted_tags_help(env.arena, tags, None, env.subs);
-                )),
+
-                Ok(Layout::Struct(&[])) => {
+                    let fn_args = fn_args.into_bump_slice();
-                    //  TODO check for stack size of 0, rather than empty record specifically
+                    let ret = arena.alloc(ret);
+
+                    use UnionVariant::*;
+                    match variant {
+                        Never | Unit => {
+                            // a max closure size of 0 means this is a standart top-level function
+                            Ok(Layout::FunctionPointer(fn_args, ret))
+                        }
+                        BoolUnion {
+                            ttrue: _,
+                            ffalse: _,
+                        } => todo!(),
+                        ByteUnion(_tagnames) => todo!(),
+                        Unwrapped(layouts) => {
+                            let closure_layout =
+                                ClosureLayout::from_unwrapped(layouts.into_bump_slice());
+
+                            Ok(Layout::Closure(fn_args, closure_layout, ret))
+                        }
+                        Wrapped(_tags) => todo!(),
+                    }
+                }
+
+                Ok(()) | Err((_, Content::FlexVar(_))) => {
+                    // a max closure size of 0 means this is a standart top-level function
                     Ok(Layout::FunctionPointer(
                         fn_args.into_bump_slice(),
                         arena.alloc(ret),
                     ))
                 }
-                Ok(Layout::Struct(closure_layouts)) => Ok(Layout::Closure(
+                Err(_) => todo!(),
-                    fn_args.into_bump_slice(),
-                    closure_layouts,
-                    arena.alloc(ret),
-                )),
-                Ok(closure_layout) => {
-                    // the closure parameter can be a tag union if there are multiple sizes
-                    // we must make sure we can distinguish between that tag union,
-                    // and the closure containing just one element, that happens to be a tag union.
-                    todo!("TODO closure layout {:?}", &closure_layout)
-                }
-                Err(LayoutProblem::UnresolvedTypeVar) => Ok(Layout::FunctionPointer(
-                    fn_args.into_bump_slice(),
-                    arena.alloc(ret),
-                )),
-
-                error => error,
             }
         }
         Record(fields, ext_var) => {