Intern captures layouts

crates/compiler/mono/src/layout.rs

@@ -5,7 +5,7 @@ use bumpalo::Bump;
 use roc_builtins::bitcode::{FloatWidth, IntWidth};
 use roc_collections::all::{default_hasher, FnvMap, MutMap};
 use roc_error_macros::{internal_error, todo_abilities};
-use roc_intern::{Interned, Interner, SingleThreadedInterner, ThreadLocalInterner};
+use roc_intern::{Interner, SingleThreadedInterner, ThreadLocalInterner};
 use roc_module::ident::{Lowercase, TagName};
 use roc_module::symbol::{Interns, Symbol};
 use roc_problem::can::RuntimeError;
@@ -311,6 +311,10 @@ impl<'a> LayoutCache<'a> {
         }
     }
 
+    pub fn get_in(&self, interned: InLayout<'a>) -> &Layout<'a> {
+        self.interner.get(interned)
+    }
+
     #[cfg(debug_assertions)]
     pub fn statistics(&self) -> (CacheStatistics, CacheStatistics) {
         (self.stats, self.raw_function_stats)
@@ -653,6 +657,10 @@ impl FieldOrderHash {
     }
 }
 
+/// An interned layout.
+// One day I would like to take over `LayoutId` as the name here, but that is currently used elsewhere.
+pub type InLayout<'a> = roc_intern::Interned<Layout<'a>>;
+
 /// Types for code gen must be monomorphic. No type variables allowed!
 #[derive(Clone, Copy, Debug, PartialEq, Eq, Hash, PartialOrd, Ord)]
 pub enum Layout<'a> {
@@ -1149,7 +1157,7 @@ impl Discriminant {
 /// concurrently. The number does not change and will give a reliable output.
 struct SetElement<'a> {
     symbol: Symbol,
-    layout: &'a [Layout<'a>],
+    layout: &'a [InLayout<'a>],
 }
 
 impl std::fmt::Debug for SetElement<'_> {
@@ -1163,7 +1171,7 @@ impl std::fmt::Debug for SetElement<'_> {
 impl std::fmt::Debug for LambdaSet<'_> {
     fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
         struct Helper<'a> {
-            set: &'a [(Symbol, &'a [Layout<'a>])],
+            set: &'a [(Symbol, &'a [InLayout<'a>])],
         }
 
         impl std::fmt::Debug for Helper<'_> {
@@ -1204,7 +1212,7 @@ 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 [Layout<'a>]);
+pub struct CapturesNiche<'a>(pub(crate) &'a [InLayout<'a>]);
 
 impl CapturesNiche<'_> {
     pub fn no_niche() -> Self {
@@ -1254,9 +1262,9 @@ impl<'a> LambdaName<'a> {
 #[derive(Clone, Copy, PartialEq, Eq, Hash, PartialOrd, Ord)]
 pub struct LambdaSet<'a> {
     /// collection of function names and their closure arguments
-    pub(crate) set: &'a [(Symbol, &'a [Layout<'a>])],
+    pub(crate) set: &'a [(Symbol, &'a [InLayout<'a>])],
     /// how the closure will be represented at runtime
-    pub(crate) representation: Interned<Layout<'a>>,
+    pub(crate) representation: InLayout<'a>,
 }
 
 #[derive(Debug)]
@@ -1271,7 +1279,7 @@ pub enum ClosureRepresentation<'a> {
     /// The closure is represented as a union. Includes the tag ID!
     /// Each variant is a different function, and its payloads are the captures.
     Union {
-        alphabetic_order_fields: &'a [Layout<'a>],
+        alphabetic_order_fields: &'a [InLayout<'a>],
         closure_name: Symbol,
         tag_id: TagIdIntType,
         union_layout: UnionLayout<'a>,
@@ -1280,7 +1288,7 @@ pub enum ClosureRepresentation<'a> {
     /// The layouts are sorted alphabetically by the identifier that is captured.
     ///
     /// We MUST sort these according to their stack size before code gen!
-    AlphabeticOrderStruct(&'a [Layout<'a>]),
+    AlphabeticOrderStruct(&'a [InLayout<'a>]),
     /// The closure is one function that captures a single identifier, whose value is unwrapped.
     UnwrappedCapture(Layout<'a>),
     /// The closure dispatches to multiple functions, but none of them capture anything, so this is
@@ -1366,7 +1374,7 @@ impl<'a> LambdaSet<'a> {
     {
         debug_assert!(self.contains(lambda_name.name));
 
-        let comparator = |other_name: Symbol, other_captures_layouts: &[Layout]| {
+        let comparator = |other_name: Symbol, other_captures_layouts: &[InLayout]| {
             other_name == lambda_name.name
                 // Make sure all captures are equal
                 && other_captures_layouts
@@ -1394,7 +1402,7 @@ impl<'a> LambdaSet<'a> {
             self
         );
 
-        let comparator = |other_name: Symbol, other_captures_layouts: &[Layout]| {
+        let comparator = |other_name: Symbol, other_captures_layouts: &[InLayout<'a>]| {
             other_name == function_symbol
                 && other_captures_layouts.iter().zip(captures_layouts).all(
                     |(other_layout, layout)| {
@@ -1424,10 +1432,11 @@ impl<'a> LambdaSet<'a> {
 
     /// Checks if two captured layouts are equivalent under the current lambda set.
     /// Resolves recursive pointers to the layout of the lambda set.
-    fn capture_layouts_eq<I>(&self, interner: &I, left: &Layout, right: &Layout) -> bool
+    fn capture_layouts_eq<I>(&self, interner: &I, left: &InLayout<'a>, right: &Layout) -> bool
     where
         I: Interner<'a, Layout<'a>>,
     {
+        let left = interner.get(*left);
         if left == right {
             return true;
         }
@@ -1460,7 +1469,7 @@ impl<'a> LambdaSet<'a> {
     fn layout_for_member<I, F>(&self, interner: &I, comparator: F) -> ClosureRepresentation<'a>
     where
         I: Interner<'a, Layout<'a>>,
-        F: Fn(Symbol, &[Layout]) -> bool,
+        F: Fn(Symbol, &[InLayout]) -> bool,
     {
         let repr = interner.get(self.representation);
 
@@ -1687,7 +1696,7 @@ impl<'a> LambdaSet<'a> {
                 // sort the tags; make sure ordering stays intact!
                 lambdas.sort_by_key(|(sym, _)| *sym);
 
-                let mut set: Vec<(Symbol, &[Layout])> =
+                let mut set: Vec<(Symbol, &[InLayout])> =
                     Vec::with_capacity_in(lambdas.len(), env.arena);
                 let mut set_with_variables: std::vec::Vec<(&Symbol, &[Variable])> =
                     std::vec::Vec::with_capacity(lambdas.len());
@@ -1708,7 +1717,8 @@ impl<'a> LambdaSet<'a> {
                         // representation, so here the criteria doesn't matter.
                         let mut criteria = CACHEABLE;
                         let arg = cached!(Layout::from_var(env, *var), criteria);
-                        arguments.push(arg);
+                        let arg_in = env.cache.interner.insert(env.arena.alloc(arg));
+                        arguments.push(arg_in);
                     }
 
                     let arguments = arguments.into_bump_slice();