Lift Niche from just captures-niche to generic niche

crates/compiler/mono/src/borrow.rs

@@ -513,7 +513,7 @@ impl<'a> BorrowInfState<'a> {
                 ..
             } => {
                 let top_level =
-                    ProcLayout::new(self.arena, arg_layouts, name.captures_niche(), **ret_layout);
+                    ProcLayout::new(self.arena, arg_layouts, name.niche(), **ret_layout);
 
                 // get the borrow signature of the applied function
                 let ps = param_map
@@ -556,7 +556,7 @@ impl<'a> BorrowInfState<'a> {
                 let closure_layout = ProcLayout {
                     arguments: passed_function.argument_layouts,
                     result: passed_function.return_layout,
-                    captures_niche: passed_function.name.captures_niche(),
+                    niche: passed_function.name.niche(),
                 };
 
                 let function_ps =
@@ -739,8 +739,7 @@ impl<'a> BorrowInfState<'a> {
             Stmt::Ret(z),
         ) = (v, b)
         {
-            let top_level =
-                ProcLayout::new(self.arena, arg_layouts, g.captures_niche(), **ret_layout);
+            let top_level = ProcLayout::new(self.arena, arg_layouts, g.niche(), **ret_layout);
 
             if self.current_proc == g.name() && x == *z {
                 // anonymous functions (for which the ps may not be known)

crates/compiler/mono/src/code_gen_help/mod.rs

@@ -9,7 +9,9 @@ use crate::ir::{
     Call, CallSpecId, CallType, Expr, HostExposedLayouts, JoinPointId, ModifyRc, Proc, ProcLayout,
     SelfRecursive, Stmt, UpdateModeId,
 };
-use crate::layout::{Builtin, CapturesNiche, LambdaName, Layout, STLayoutInterner, UnionLayout};
+use crate::layout::{
+    Builtin, CapturesNiche, LambdaName, Layout, Niche, STLayoutInterner, UnionLayout,
+};
 
 mod equality;
 mod refcount;
@@ -404,23 +406,23 @@ impl<'a> CodeGenHelp<'a> {
             HelperOp::Inc => ProcLayout {
                 arguments: self.arena.alloc([*layout, self.layout_isize]),
                 result: LAYOUT_UNIT,
-                captures_niche: CapturesNiche::no_niche(),
+                niche: Niche::NONE,
             },
             HelperOp::Dec => ProcLayout {
                 arguments: self.arena.alloc([*layout]),
                 result: LAYOUT_UNIT,
-                captures_niche: CapturesNiche::no_niche(),
+                niche: Niche::NONE,
             },
             HelperOp::Reset => ProcLayout {
                 arguments: self.arena.alloc([*layout]),
                 result: *layout,
-                captures_niche: CapturesNiche::no_niche(),
+                niche: Niche::NONE,
             },
             HelperOp::DecRef(_) => unreachable!("No generated Proc for DecRef"),
             HelperOp::Eq => ProcLayout {
                 arguments: self.arena.alloc([*layout, *layout]),
                 result: LAYOUT_BOOL,
-                captures_niche: CapturesNiche::no_niche(),
+                niche: Niche::NONE,
             },
         };
 

crates/compiler/mono/src/debug/checker.rs

@@ -556,7 +556,7 @@ impl<'a, 'r> Ctx<'a, 'r> {
                 let proc_layout = ProcLayout {
                     arguments: arg_layouts,
                     result: **ret_layout,
-                    captures_niche: name.captures_niche(),
+                    niche: name.niche(),
                 };
                 if !self.procs.contains_key(&(name.name(), proc_layout)) {
                     let similar = self

crates/compiler/mono/src/debug/report.rs

@@ -6,7 +6,7 @@ use ven_pretty::{Arena, DocAllocator, DocBuilder};
 
 use crate::{
     ir::{Parens, ProcLayout},
-    layout::{CapturesNiche, Layout},
+    layout::{Layout, Niche},
 };
 
 use super::{
@@ -465,7 +465,7 @@ where
     let ProcLayout {
         arguments,
         result,
-        captures_niche,
+        niche: captures_niche,
     } = proc_layout;
     let args = f.intersperse(
         arguments
@@ -478,20 +478,19 @@ where
         f.reflow(" -> "),
         result.to_doc(f, interner, Parens::NotNeeded),
     ]);
-    let niche = if captures_niche == CapturesNiche::no_niche() {
-        f.reflow("(no niche)")
-    } else {
-        f.concat([
+    let niche = match captures_niche {
+        Niche::NONE => f.reflow("(no niche)"),
+        Niche::Captures(captures_niche) => f.concat([
             f.reflow("(niche {"),
             f.intersperse(
                 captures_niche
-                    .0
+                    .captures()
                     .iter()
                     .map(|&c| interner.get(c).to_doc(f, interner, Parens::NotNeeded)),
                 f.reflow(", "),
             ),
             f.reflow("})"),
-        ])
+        ]),
     };
     f.concat([fun, f.space(), niche])
 }

crates/compiler/mono/src/inc_dec.rs

@@ -595,7 +595,7 @@ impl<'a, 'i> Context<'a, 'i> {
                 ..
             } => {
                 let top_level =
-                    ProcLayout::new(self.arena, arg_layouts, name.captures_niche(), **ret_layout);
+                    ProcLayout::new(self.arena, arg_layouts, name.niche(), **ret_layout);
 
                 // get the borrow signature
                 let ps = self
@@ -645,7 +645,7 @@ impl<'a, 'i> Context<'a, 'i> {
         let function_layout = ProcLayout {
             arguments: passed_function.argument_layouts,
             result: passed_function.return_layout,
-            captures_niche: passed_function.name.captures_niche(),
+            niche: passed_function.name.niche(),
         };
 
         let function_ps = match self

crates/compiler/mono/src/ir.rs

@@ -1,9 +1,9 @@
 #![allow(clippy::manual_map)]
 
 use crate::layout::{
-    self, Builtin, CapturesNiche, ClosureCallOptions, ClosureRepresentation, EnumDispatch,
-    LambdaName, LambdaSet, Layout, LayoutCache, LayoutProblem, RawFunctionLayout, STLayoutInterner,
-    TagIdIntType, UnionLayout, WrappedVariant,
+    self, Builtin, ClosureCallOptions, ClosureRepresentation, EnumDispatch, LambdaName, LambdaSet,
+    Layout, LayoutCache, LayoutProblem, Niche, RawFunctionLayout, STLayoutInterner, TagIdIntType,
+    UnionLayout, WrappedVariant,
 };
 use bumpalo::collections::{CollectIn, Vec};
 use bumpalo::Bump;
@@ -3391,7 +3391,7 @@ fn specialize_proc_help<'a>(
                     let top_level = ProcLayout::new(
                         env.arena,
                         top_level_arguments.into_bump_slice(),
-                        CapturesNiche::no_niche(),
+                        Niche::NONE,
                         *return_layout,
                     );
 
@@ -3424,7 +3424,7 @@ fn specialize_proc_help<'a>(
                         *symbol,
                         (
                             name,
-                            ProcLayout::new(env.arena, &[], CapturesNiche::no_niche(), result),
+                            ProcLayout::new(env.arena, &[], Niche::NONE, result),
                             layout,
                         ),
                     );
@@ -3956,14 +3956,14 @@ fn specialize_variable<'a>(
 pub struct ProcLayout<'a> {
     pub arguments: &'a [Layout<'a>],
     pub result: Layout<'a>,
-    pub captures_niche: CapturesNiche<'a>,
+    pub niche: Niche<'a>,
 }
 
 impl<'a> ProcLayout<'a> {
     pub(crate) fn new(
         arena: &'a Bump,
         old_arguments: &'a [Layout<'a>],
-        old_captures_niche: CapturesNiche<'a>,
+        old_niche: Niche<'a>,
         result: Layout<'a>,
     ) -> Self {
         let mut arguments = Vec::with_capacity_in(old_arguments.len(), arena);
@@ -3978,7 +3978,7 @@ impl<'a> ProcLayout<'a> {
 
         ProcLayout {
             arguments: arguments.into_bump_slice(),
-            captures_niche: old_captures_niche,
+            niche: old_niche,
             result: new_result,
         }
     }
@@ -3992,10 +3992,10 @@ impl<'a> ProcLayout<'a> {
             RawFunctionLayout::Function(arguments, lambda_set, result) => {
                 let arguments =
                     lambda_set.extend_argument_list_for_named(arena, lambda_name, arguments);
-                ProcLayout::new(arena, arguments, lambda_name.captures_niche(), *result)
+                ProcLayout::new(arena, arguments, lambda_name.niche(), *result)
             }
             RawFunctionLayout::ZeroArgumentThunk(result) => {
-                ProcLayout::new(arena, &[], CapturesNiche::no_niche(), result)
+                ProcLayout::new(arena, &[], Niche::NONE, result)
             }
         }
     }
@@ -8382,8 +8382,7 @@ fn specialize_symbol<'a>(
                         // let layout = Layout::Closure(argument_layouts, lambda_set, ret_layout);
                         // panic!("suspicious");
                         let layout = Layout::LambdaSet(lambda_set);
-                        let top_level =
-                            ProcLayout::new(env.arena, &[], CapturesNiche::no_niche(), layout);
+                        let top_level = ProcLayout::new(env.arena, &[], Niche::NONE, layout);
                         procs.insert_passed_by_name(
                             env,
                             arg_var,
@@ -8427,8 +8426,7 @@ fn specialize_symbol<'a>(
                 }
                 RawFunctionLayout::ZeroArgumentThunk(ret_layout) => {
                     // this is a 0-argument thunk
-                    let top_level =
-                        ProcLayout::new(env.arena, &[], CapturesNiche::no_niche(), ret_layout);
+                    let top_level = ProcLayout::new(env.arena, &[], Niche::NONE, ret_layout);
                     procs.insert_passed_by_name(
                         env,
                         arg_var,
@@ -8756,12 +8754,7 @@ fn call_by_name_help<'a>(
     let top_level_layout = {
         let argument_layouts =
             lambda_set.extend_argument_list_for_named(env.arena, proc_name, argument_layouts);
-        ProcLayout::new(
-            env.arena,
-            argument_layouts,
-            proc_name.captures_niche(),
-            *ret_layout,
-        )
+        ProcLayout::new(env.arena, argument_layouts, proc_name.niche(), *ret_layout)
     };
 
     // the variables of the given arguments
@@ -9016,7 +9009,7 @@ fn call_by_name_module_thunk<'a>(
     assigned: Symbol,
     hole: &'a Stmt<'a>,
 ) -> Stmt<'a> {
-    let top_level_layout = ProcLayout::new(env.arena, &[], CapturesNiche::no_niche(), *ret_layout);
+    let top_level_layout = ProcLayout::new(env.arena, &[], Niche::NONE, *ret_layout);
 
     let inner_layout = *ret_layout;
 

crates/compiler/mono/src/layout.rs

@@ -1241,18 +1241,28 @@ impl std::fmt::Debug for LambdaSet<'_> {
 ///
 /// See also https://github.com/roc-lang/roc/issues/3336.
 #[derive(Clone, Copy, PartialEq, Eq, Hash, Debug)]
-pub struct CapturesNiche<'a>(pub(crate) &'a [InLayout<'a>]);
+pub struct CapturesNiche<'a>(&'a [InLayout<'a>]);
 
-impl CapturesNiche<'_> {
-    pub fn no_niche() -> Self {
-        Self(&[])
+impl<'a> CapturesNiche<'a> {
+    pub(crate) fn captures(&self) -> &'a [Layout<'a>] {
+        self.0
     }
 }
 
+#[derive(Debug, Copy, Clone, PartialEq, Eq, Hash)]
+pub enum Niche<'a> {
+    /// A niche for a proc are its captures.
+    Captures(CapturesNiche<'a>),
+}
+
+impl<'a> Niche<'a> {
+    pub const NONE: Niche<'a> = Niche::Captures(CapturesNiche(&[]));
+}
+
 #[derive(Clone, Copy, PartialEq, Eq, Hash, Debug)]
 pub struct LambdaName<'a> {
     name: Symbol,
-    captures_niche: CapturesNiche<'a>,
+    niche: Niche<'a>,
 }
 
 impl<'a> LambdaName<'a> {
@@ -1262,29 +1272,29 @@ impl<'a> LambdaName<'a> {
     }
 
     #[inline(always)]
-    pub fn captures_niche(&self) -> CapturesNiche<'a> {
-        self.captures_niche
+    pub fn niche(&self) -> Niche<'a> {
+        self.niche
     }
 
     #[inline(always)]
-    pub fn no_captures(&self) -> bool {
-        self.captures_niche.0.is_empty()
+    pub(crate) fn no_captures(&self) -> bool {
+        match self.niche {
+            Niche::NONE => true,
+            Niche::Captures(captures) => captures.0.is_empty(),
+        }
     }
 
     #[inline(always)]
     pub fn no_niche(name: Symbol) -> Self {
         Self {
             name,
-            captures_niche: CapturesNiche::no_niche(),
+            niche: Niche::NONE,
         }
     }
 
     #[inline(always)]
-    pub fn replace_name(&self, name: Symbol) -> Self {
-        Self {
-            name,
-            captures_niche: self.captures_niche,
-        }
+    pub(crate) fn replace_name(&self, name: Symbol) -> Self {
+        Self { name, ..*self }
     }
 }
 
@@ -1377,9 +1387,12 @@ impl<'a> LambdaSet<'a> {
     }
 
     pub fn iter_set(&self) -> impl ExactSizeIterator<Item = LambdaName<'a>> {
-        self.set.iter().map(|(name, captures_layouts)| LambdaName {
-            name: *name,
-            captures_niche: CapturesNiche(captures_layouts),
+        self.set.iter().map(|(name, captures_layouts)| {
+            let niche = match captures_layouts {
+                [] => Niche::NONE,
+                _ => Niche::Captures(CapturesNiche(captures_layouts)),
+            };
+            LambdaName { name: *name, niche }
         })
     }
 
@@ -1403,12 +1416,17 @@ impl<'a> LambdaSet<'a> {
     {
         debug_assert!(self.contains(lambda_name.name));
 
+        let captures = match lambda_name.niche {
+            Niche::Captures(captures) => captures.0,
+            Niche::NONE => &[],
+        };
+
         let comparator = |other_name: Symbol, other_captures_layouts: &[InLayout]| {
             other_name == lambda_name.name
                 // Make sure all captures are equal
                 && other_captures_layouts
                     .iter()
-                    .eq(lambda_name.captures_niche.0)
+                    .eq(captures)
         };
 
         self.layout_for_member(interner, comparator)
@@ -1455,7 +1473,7 @@ impl<'a> LambdaSet<'a> {
 
         LambdaName {
             name: *name,
-            captures_niche: CapturesNiche(layouts),
+            niche: Niche::Captures(CapturesNiche(layouts)),
         }
     }
 
@@ -1658,6 +1676,7 @@ impl<'a> LambdaSet<'a> {
         lambda_name: LambdaName<'a>,
         argument_layouts: &'a [Layout<'a>],
     ) -> &'a [Layout<'a>] {
+        let Niche::Captures(CapturesNiche(captures)) = lambda_name.niche;
         // TODO(https://github.com/roc-lang/roc/issues/4831): we should turn on this debug-assert;
         // however, currently it causes false-positives, because host-exposed functions that are
         // function pointers to platform-exposed functions are compiled as if they are proper
@@ -1674,7 +1693,7 @@ impl<'a> LambdaSet<'a> {
         // );
 
         // If we don't capture, there is nothing to extend.
-        if lambda_name.captures_niche.0.is_empty() {
+        if captures.is_empty() {
             argument_layouts
         } else {
             let mut arguments = Vec::with_capacity_in(argument_layouts.len() + 1, arena);