preparations for specializing closures

compiler/constrain/src/expr.rs

@@ -1277,8 +1277,7 @@ fn constrain_closure_size(
         ));
     }
 
-    let tag_name_string = format!("Closure_{:?}_{}", name, closure_var.index());
+    let tag_name = roc_module::ident::TagName::Closure(name);
-    let tag_name = roc_module::ident::TagName::Global(tag_name_string.into());
     let closure_type = Type::TagUnion(
         vec![(tag_name, tag_arguments)],
         Box::new(Type::Variable(closure_ext_var)),

compiler/gen/src/llvm/convert.rs

@@ -111,8 +111,9 @@ pub fn basic_type_from_layout<'ctx>(
             basic_type_from_function_layout(arena, context, args, None, ret_layout, ptr_bytes)
         }
         Closure(args, closure_layout, ret_layout) => {
+            let closure_data_layout = closure_layout.as_block_of_memory_layout();
             let closure_data =
-                basic_type_from_layout(arena, context, &closure_layout.as_layout(), ptr_bytes);
+                basic_type_from_layout(arena, context, &closure_data_layout, ptr_bytes);
 
             let function_pointer = basic_type_from_function_layout(
                 arena,

compiler/module/src/ident.rs

@@ -37,6 +37,9 @@ pub enum TagName {
     /// Private tags are associated with a specific module, and as such use a
     /// Symbol just like all other module-specific identifiers.
     Private(Symbol),
+
+    /// Used to connect the closure size to the function it corresponds to
+    Closure(Symbol),
 }
 
 impl TagName {
@@ -44,6 +47,7 @@ impl TagName {
         match self {
             TagName::Global(uppercase) => uppercase.as_inline_str().clone(),
             TagName::Private(symbol) => symbol.fully_qualified(interns, home),
+            TagName::Closure(symbol) => symbol.fully_qualified(interns, home),
         }
     }
 }

compiler/mono/src/ir.rs

@@ -893,6 +893,7 @@ impl<'a> Expr<'a> {
                 let doc_tag = match tag_name {
                     TagName::Global(s) => alloc.text(s.as_str()),
                     TagName::Private(s) => alloc.text(format!("{}", s)),
+                    TagName::Closure(s) => alloc.text(format!("Closure({})", s)),
                 };
 
                 let it = arguments.iter().map(|s| symbol_to_doc(alloc, *s));
@@ -910,6 +911,7 @@ impl<'a> Expr<'a> {
                 let doc_tag = match tag_name {
                     TagName::Global(s) => alloc.text(s.as_str()),
                     TagName::Private(s) => alloc.text(format!("{}", s)),
+                    TagName::Closure(s) => alloc.text(format!("Closure({})", s)),
                 };
 
                 let it = arguments.iter().map(|s| symbol_to_doc(alloc, *s));
@@ -1434,7 +1436,7 @@ fn specialize_external<'a>(
     }
 
     let (proc_args, ret_layout) =
-        build_specialized_proc_from_var(env, layout_cache, pattern_symbols, fn_var)?;
+        build_specialized_proc_from_var(env, layout_cache, proc_name, pattern_symbols, fn_var)?;
 
     // reset subs, so we don't get type errors when specializing for a different signature
     layout_cache.rollback_to(cache_snapshot);
@@ -1461,6 +1463,7 @@ fn specialize_external<'a>(
 fn build_specialized_proc_from_var<'a>(
     env: &mut Env<'a, '_>,
     layout_cache: &mut LayoutCache<'a>,
+    proc_name: Symbol,
     pattern_symbols: &[Symbol],
     fn_var: Variable,
 ) -> Result<(&'a [(Layout<'a>, Symbol)], Layout<'a>), LayoutProblem> {
@@ -1471,6 +1474,7 @@ fn build_specialized_proc_from_var<'a>(
 
             build_specialized_proc(
                 env.arena,
+                proc_name,
                 pattern_symbols,
                 pattern_layouts_vec,
                 None,
@@ -1483,6 +1487,7 @@ fn build_specialized_proc_from_var<'a>(
 
             build_specialized_proc(
                 env.arena,
+                proc_name,
                 pattern_symbols,
                 pattern_layouts_vec,
                 Some(closure_layout),
@@ -1496,6 +1501,7 @@ fn build_specialized_proc_from_var<'a>(
                     build_specialized_proc_adapter(
                         env,
                         layout_cache,
+                        proc_name,
                         pattern_symbols,
                         &pattern_vars,
                         closure_layout,
@@ -1505,16 +1511,27 @@ fn build_specialized_proc_from_var<'a>(
                 Content::Structure(FlatType::Apply(Symbol::ATTR_ATTR, args))
                     if !pattern_symbols.is_empty() =>
                 {
-                    build_specialized_proc_from_var(env, layout_cache, pattern_symbols, args[1])
+                    build_specialized_proc_from_var(
-                }
+                        env,
-                Content::Alias(_, _, actual) => {
+                        layout_cache,
-                    build_specialized_proc_from_var(env, layout_cache, pattern_symbols, actual)
+                        proc_name,
+                        pattern_symbols,
+                        args[1],
+                    )
                 }
+                Content::Alias(_, _, actual) => build_specialized_proc_from_var(
+                    env,
+                    layout_cache,
+                    proc_name,
+                    pattern_symbols,
+                    actual,
+                ),
                 _ => {
                     // a top-level constant 0-argument thunk
                     build_specialized_proc_adapter(
                         env,
                         layout_cache,
+                        proc_name,
                         pattern_symbols,
                         &[],
                         None,
@@ -1529,6 +1546,7 @@ fn build_specialized_proc_from_var<'a>(
 fn build_specialized_proc_adapter<'a>(
     env: &mut Env<'a, '_>,
     layout_cache: &mut LayoutCache<'a>,
+    proc_name: Symbol,
     pattern_symbols: &[Symbol],
     pattern_vars: &[Variable],
     opt_closure_layout: Option<ClosureLayout<'a>>,
@@ -1548,6 +1566,7 @@ fn build_specialized_proc_adapter<'a>(
 
     build_specialized_proc(
         env.arena,
+        proc_name,
         pattern_symbols,
         arg_layouts,
         opt_closure_layout,
@@ -1558,6 +1577,7 @@ fn build_specialized_proc_adapter<'a>(
 #[allow(clippy::type_complexity)]
 fn build_specialized_proc<'a>(
     arena: &'a Bump,
+    proc_name: Symbol,
     pattern_symbols: &[Symbol],
     pattern_layouts: Vec<Layout<'a>>,
     opt_closure_layout: Option<ClosureLayout<'a>>,
@@ -1592,7 +1612,7 @@ fn build_specialized_proc<'a>(
         Some(layout) if pattern_symbols.last() == Some(&Symbol::ARG_CLOSURE) => {
             // here we define the lifted (now top-level) f function. Its final argument is `Symbol::ARG_CLOSURE`,
             // it stores the closure structure (just an integer in this case)
-            proc_args.push((layout.as_layout(), Symbol::ARG_CLOSURE));
+            proc_args.push((layout.as_named_layout(proc_name), Symbol::ARG_CLOSURE));
 
             debug_assert_eq!(
                 pattern_layouts_len + 1,
@@ -1604,7 +1624,7 @@ fn build_specialized_proc<'a>(
             // else if there is a closure layout, we're building the `f_closure` value
             // that means we're really creating a ( function_ptr, closure_data ) pair
 
-            let closure_data_layout = layout.as_layout();
+            let closure_data_layout = layout.as_block_of_memory_layout();
             let function_ptr_layout = Layout::FunctionPointer(
                 arena.alloc([Layout::Struct(&[]), closure_data_layout.clone()]),
                 arena.alloc(ret_layout),

compiler/mono/src/layout.rs

@@ -33,7 +33,7 @@ pub enum Layout<'a> {
 #[derive(Clone, Debug, PartialEq, Eq, Hash)]
 pub struct ClosureLayout<'a> {
     /// the layout that this specific closure captures
-    captured: &'a [Layout<'a>],
+    captured: &'a [(TagName, &'a [Layout<'a>])],
 
     /// the layout that represents the maximum size the closure layout can have
     max_size: &'a [Layout<'a>],
@@ -44,7 +44,7 @@ impl<'a> ClosureLayout<'a> {
         let layout = Layout::Builtin(Builtin::Int1);
         let layouts = arena.alloc([layout]);
         ClosureLayout {
-            captured: layouts,
+            captured: &[],
             max_size: layouts,
         }
     }
@@ -52,18 +52,45 @@ impl<'a> ClosureLayout<'a> {
         let layout = Layout::Builtin(Builtin::Int8);
         let layouts = arena.alloc([layout]);
         ClosureLayout {
-            captured: layouts,
+            captured: &[],
             max_size: layouts,
         }
     }
     fn from_unwrapped(layouts: &'a [Layout<'a>]) -> Self {
         debug_assert!(!layouts.is_empty());
         ClosureLayout {
-            captured: layouts,
+            captured: &[],
             max_size: layouts,
         }
     }
 
+    fn from_wrapped(tags: &'a [(TagName, &'a [Layout<'a>])]) -> Self {
+        // NOTE we fabricate a pointer size here.
+        // That's fine because we don't care about the exact size, just the biggest one
+        let pointer_size = 8;
+
+        let mut largest_size = 0;
+        let mut largest = None;
+
+        for (_, tag_args) in tags.iter() {
+            let size = tag_args.iter().map(|l| l.stack_size(pointer_size)).sum();
+
+            // >= because some of our layouts have 0 size, but are still valid layouts
+            if size >= largest_size {
+                largest_size = size;
+                largest = Some(tag_args);
+            }
+        }
+
+        match largest {
+            None => unreachable!("A tag union layout must always contain 2 or more tags"),
+            Some(max_size) => ClosureLayout {
+                captured: tags,
+                max_size,
+            },
+        }
+    }
+
     pub fn from_var(
         arena: &'a Bump,
         subs: &Subs,
@@ -97,9 +124,10 @@ impl<'a> ClosureLayout<'a> {
 
                         Ok(Some(closure_layout))
                     }
-                    Wrapped(_tags) => {
+                    Wrapped(tags) => {
                         // Wrapped(Vec<'a, (TagName, &'a [Layout<'a>])>),
-                        todo!("can't specialize multi-size closures yet")
+                        let closure_layout = ClosureLayout::from_wrapped(tags.into_bump_slice());
+                        Ok(Some(closure_layout))
                     }
                 }
             }
@@ -136,17 +164,52 @@ impl<'a> ClosureLayout<'a> {
             .sum()
     }
     pub fn contains_refcounted(&self) -> bool {
-        self.captured.iter().any(|l| l.contains_refcounted())
+        self.captured
+            .iter()
+            .map(|t| t.1)
+            .flatten()
+            .any(|l| l.contains_refcounted())
     }
     pub fn safe_to_memcpy(&self) -> bool {
-        self.captured.iter().all(|l| l.safe_to_memcpy())
+        self.captured
+            .iter()
+            .map(|t| t.1)
+            .flatten()
+            .all(|l| l.safe_to_memcpy())
+    }
+
+    pub fn as_named_layout(&self, symbol: Symbol) -> Layout<'a> {
+        let layouts = if self.captured.is_empty() {
+            self.max_size
+        } else if let Some((_, tag_args)) = self
+            .captured
+            .iter()
+            .find(|(tn, _)| *tn == TagName::Closure(symbol))
+        {
+            tag_args
+        } else {
+            unreachable!(
+                "invariant broken, TagName::Closure({:?}) is not in {:?}",
+                symbol, &self.captured
+            );
+        };
+
+        if layouts.len() == 1 {
+            layouts[0].clone()
+        } else {
+            Layout::Struct(layouts)
+        }
     }
 
     pub fn as_layout(&self) -> Layout<'a> {
-        if self.captured.len() == 1 {
+        if self.captured.is_empty() {
-            self.captured[0].clone()
+            if self.max_size.len() == 1 {
+                self.max_size[0].clone()
+            } else {
+                Layout::Struct(self.max_size)
+            }
         } else {
-            Layout::Struct(self.captured)
+            panic!();
         }
     }
 

compiler/reporting/src/error/type.rs

@@ -924,6 +924,10 @@ fn add_category<'b>(
             alloc.private_tag_name(*name),
             alloc.text(" private tag application has the type:"),
         ]),
+        TagApply {
+            tag_name: TagName::Closure(_name),
+            args_count: _,
+        } => unreachable!("closure tags are for internal use only"),
 
         Record => alloc.concat(vec![this_is, alloc.text(" a record of type:")]),
 

compiler/reporting/src/report.rs

@@ -248,6 +248,7 @@ impl<'a> RocDocAllocator<'a> {
         match tn {
             TagName::Global(uppercase) => self.global_tag_name(uppercase),
             TagName::Private(symbol) => self.private_tag_name(symbol),
+            TagName::Closure(_symbol) => unreachable!("closure tags are internal only"),
         }
     }