refactor ClosureLayout

compiler/constrain/src/expr.rs

@@ -1256,8 +1256,8 @@ fn constrain_closure_size(
     closure_ext_var: Variable,
     variables: &mut Vec<Variable>,
 ) -> Constraint {
-    debug_assert!(variables.iter().find(|s| **s == closure_var).is_some());
+    debug_assert!(variables.iter().any(|s| *s == closure_var));
-    debug_assert!(variables.iter().find(|s| **s == closure_ext_var).is_some());
+    debug_assert!(variables.iter().any(|s| *s == closure_ext_var));
 
     let mut tag_arguments = Vec::with_capacity(captured_symbols.len());
     let mut captured_symbols_constraints = Vec::with_capacity(captured_symbols.len());

compiler/gen/tests/gen_tags.rs

@@ -436,7 +436,6 @@ mod gen_tags {
     }
 
     #[test]
-    #[ignore]
     fn maybe_is_just_nested() {
         assert_evals_to!(
             indoc!(

compiler/mono/src/ir.rs

@@ -1661,7 +1661,7 @@ fn build_specialized_proc<'a>(
             let closure_layout =
                 Layout::Struct(arena.alloc([function_ptr_layout, closure_data_layout]));
 
-            return Ok((&[], None, closure_layout));
+            Ok((&[], None, closure_layout))
         }
         None => {
             // else we're making a normal function, no closure problems to worry about

compiler/mono/src/layout.rs

@@ -33,9 +33,11 @@ pub enum Layout<'a> {
 #[derive(Clone, Debug, PartialEq, Eq, Hash)]
 pub struct ClosureLayout<'a> {
     /// the layout that this specific closure captures
+    /// uses a Vec instead of a MutMap because it's Hash
+    /// the vec is likely to be small, so linear search is fine
     captured: &'a [(TagName, &'a [Layout<'a>])],
 
-    max_size: &'a Layout<'a>,
+    layout: &'a Layout<'a>,
 }
 
 impl<'a> ClosureLayout<'a> {
@@ -44,7 +46,7 @@ impl<'a> ClosureLayout<'a> {
         let layouts = arena.alloc(layout);
         ClosureLayout {
             captured: &[],
-            max_size: layouts,
+            layout: layouts,
         }
     }
     fn from_byte(arena: &'a Bump) -> Self {
@@ -52,49 +54,42 @@ impl<'a> ClosureLayout<'a> {
         let layouts = arena.alloc(layout);
         ClosureLayout {
             captured: &[],
-            max_size: layouts,
+            layout: layouts,
         }
     }
     fn from_unwrapped(arena: &'a Bump, layouts: &'a [Layout<'a>]) -> Self {
         debug_assert!(!layouts.is_empty());
-        if layouts.len() == 1 {
+        let layout = if layouts.len() == 1 {
-            ClosureLayout {
+            &layouts[0]
-                captured: &[],
-                max_size: &layouts[0],
-            }
         } else {
+            arena.alloc(Layout::Struct(layouts))
+        };
+
         ClosureLayout {
             captured: &[],
-                max_size: arena.alloc(Layout::Struct(layouts)),
+            layout,
-            }
         }
     }
 
     fn from_tag_union(arena: &'a Bump, tags: &'a [(TagName, &'a [Layout<'a>])]) -> Self {
-        debug_assert!(!tags.is_empty());
+        debug_assert!(tags.len() > 1);
 
         let mut tag_arguments = Vec::with_capacity_in(tags.len(), arena);
-        let mut foobar = Vec::with_capacity_in(tags.len(), arena);
 
-        for (name, tag_args_with_discr) in tags.iter() {
+        for (_, tag_args) in tags.iter() {
-            // drop the discriminant, essentially creating an untagged union
+            tag_arguments.push(&tag_args[0..]);
-            let tag_args = &tag_args_with_discr[0..];
-
-            tag_arguments.push((name.clone(), tag_args));
-            foobar.push(tag_args);
         }
 
-        let arguments = tag_arguments.into_bump_slice();
         ClosureLayout {
-            captured: arguments,
+            captured: tags,
-            max_size: arena.alloc(Layout::Union(foobar.into_bump_slice())),
+            layout: arena.alloc(Layout::Union(tag_arguments.into_bump_slice())),
         }
     }
 
     pub fn get_wrapped(&self) -> crate::ir::Wrapped {
         use crate::ir::Wrapped;
 
-        match self.max_size {
+        match self.layout {
             Layout::Struct(_) => Wrapped::RecordOrSingleTagUnion,
             Layout::Union(_) => Wrapped::MultiTagUnion,
             _ => Wrapped::SingleElementRecord,
@@ -157,7 +152,7 @@ impl<'a> ClosureLayout<'a> {
         closure_layout: Self,
         ret_layout: &'a Layout<'a>,
     ) -> Layout<'a> {
-        let closure_data_layout = closure_layout.max_size;
+        let closure_data_layout = closure_layout.layout;
 
         // define the function pointer
         let function_ptr_layout = {
@@ -170,18 +165,18 @@ impl<'a> ClosureLayout<'a> {
     }
 
     pub fn stack_size(&self, pointer_size: u32) -> u32 {
-        self.max_size.stack_size(pointer_size)
+        self.layout.stack_size(pointer_size)
     }
     pub fn contains_refcounted(&self) -> bool {
-        self.max_size.contains_refcounted()
+        self.layout.contains_refcounted()
     }
     pub fn safe_to_memcpy(&self) -> bool {
-        self.max_size.safe_to_memcpy()
+        self.layout.safe_to_memcpy()
     }
 
     pub fn as_named_layout(&self, symbol: Symbol) -> Layout<'a> {
         let layouts = if self.captured.is_empty() {
-            self.max_size.clone()
+            self.layout.clone()
         } else if let Some((_, tag_args)) = self
             .captured
             .iter()
@@ -203,7 +198,7 @@ impl<'a> ClosureLayout<'a> {
     }
 
     pub fn as_block_of_memory_layout(&self) -> Layout<'a> {
-        self.max_size.clone()
+        self.layout.clone()
     }
 
     pub fn build_closure_data(
@@ -213,13 +208,18 @@ impl<'a> ClosureLayout<'a> {
     ) -> Result<crate::ir::Expr<'a>, Symbol> {
         use crate::ir::Expr;
 
-        match self.max_size {
+        match self.layout {
-            Layout::Struct(fields) if fields.len() > 1 => Ok(Expr::Struct(symbols)),
+            Layout::Struct(fields) => {
+                debug_assert!(fields.len() > 1);
+                debug_assert_eq!(fields.len(), symbols.len());
+
+                Ok(Expr::Struct(symbols))
+            }
             Layout::Union(tags) => {
                 let expr = Expr::Tag {
                     tag_layout: Layout::Union(tags),
                     tag_name: TagName::Closure(original),
-                    tag_id: 0 as u8,
+                    tag_id: 0,
                     union_size: tags.len() as u8,
                     arguments: symbols,
                 };