Collect deeply nested type names in alias analysis

Previously we only collected type names that appeared on the surface of
a type during alias analysis, but certain types that need to be named
(i.e. recursive types) may be observed only when we actually convert the IR
to the morphic IR. Make sure we collect those appropriately.

This is a cherry pick of the relevant changes in #4121.

crates/compiler/alias_analysis/src/lib.rs

@@ -1,5 +1,4 @@
-//! Performs analysis and optimizations to remove unneeded [reference counts](https://en.wikipedia.org/wiki/Reference_counting)
-//! at runtime, and supports in-place mutation.
+use bumpalo::Bump;
 use morphic_lib::TypeContext;
 use morphic_lib::{
     BlockExpr, BlockId, CalleeSpecVar, ConstDefBuilder, ConstName, EntryPointName, ExprContext,
@@ -18,8 +17,6 @@ use roc_mono::layout::{
     Builtin, CapturesNiche, Layout, RawFunctionLayout, STLayoutInterner, UnionLayout,
 };
 
-use roc_error_macros::internal_error;
-
 // just using one module for now
 pub const MOD_APP: ModName = ModName(b"UserApp");
 
@@ -136,7 +133,8 @@ fn bytes_as_ascii(bytes: &[u8]) -> String {
 }
 
 pub fn spec_program<'a, I>(
-    interner: &STLayoutInterner,
+    arena: &'a Bump,
+    interner: &STLayoutInterner<'a>,
     opt_level: OptLevel,
     opt_entry_point: Option<roc_mono::ir::EntryPoint<'a>>,
     procs: I,
@@ -221,7 +219,7 @@ where
                 );
             }
 
-            let (spec, type_names) = proc_spec(interner, proc)?;
+            let (spec, type_names) = proc_spec(arena, interner, proc)?;
 
             type_definitions.extend(type_names);
 
@@ -238,12 +236,18 @@ where
             );
             let roc_main = FuncName(&roc_main_bytes);
 
+            let mut env = Env::new(arena);
+
             let entry_point_function = build_entry_point(
+                &mut env,
                 interner,
                 entry_point.layout,
                 roc_main,
                 &host_exposed_functions,
             )?;
+
+            type_definitions.extend(env.type_names);
+
             let entry_point_name = FuncName(ENTRY_POINT_NAME);
             m.add_func(entry_point_name, entry_point_function)?;
         }
@@ -254,7 +258,12 @@ where
 
             let mut builder = TypeDefBuilder::new();
 
-            let variant_types = recursive_variant_types(&mut builder, interner, &union_layout)?;
+            let mut env = Env::new(arena);
+            let variant_types =
+                recursive_variant_types(&mut env, &mut builder, interner, &union_layout)?;
+
+            // FIXME: dropping additional env.type_names here!
+
             let root_type = if let UnionLayout::NonNullableUnwrapped(_) = union_layout {
                 debug_assert_eq!(variant_types.len(), 1);
                 variant_types[0]
@@ -311,11 +320,12 @@ fn terrible_hack(builder: &mut FuncDefBuilder, block: BlockId, type_id: TypeId)
     builder.add_unwrap_union(block, value, 1)
 }
 
-fn build_entry_point(
-    interner: &STLayoutInterner,
-    layout: roc_mono::ir::ProcLayout,
+fn build_entry_point<'a>(
+    env: &mut Env<'a>,
+    interner: &STLayoutInterner<'a>,
+    layout: roc_mono::ir::ProcLayout<'a>,
     func_name: FuncName,
-    host_exposed_functions: &[([u8; SIZE], &[Layout])],
+    host_exposed_functions: &[([u8; SIZE], &'a [Layout<'a>])],
 ) -> Result<FuncDef> {
     let mut builder = FuncDefBuilder::new();
     let outer_block = builder.add_block();
@@ -327,6 +337,7 @@ fn build_entry_point(
 
         // to the modelling language, the arguments appear out of thin air
         let argument_type = build_tuple_type(
+            env,
             &mut builder,
             interner,
             layout.arguments,
@@ -361,9 +372,10 @@ fn build_entry_point(
         let block = builder.add_block();
 
         let type_id = layout_spec(
+            env,
             &mut builder,
             interner,
-            &Layout::struct_no_name_order(layouts),
+            env.arena.alloc(Layout::struct_no_name_order(layouts)),
             &WhenRecursive::Unreachable,
         )?;
 
@@ -389,16 +401,17 @@ fn build_entry_point(
 }
 
 fn proc_spec<'a>(
+    arena: &'a Bump,
     interner: &STLayoutInterner<'a>,
     proc: &Proc<'a>,
 ) -> Result<(FuncDef, MutSet<UnionLayout<'a>>)> {
     let mut builder = FuncDefBuilder::new();
-    let mut env = Env::default();
+    let mut env = Env::new(arena);
 
     let block = builder.add_block();
 
     // introduce the arguments
-    let mut argument_layouts = Vec::new();
+    let mut argument_layouts = bumpalo::collections::Vec::with_capacity_in(proc.args.len(), arena);
     for (i, (layout, symbol)) in proc.args.iter().enumerate() {
         let value_id = builder.add_get_tuple_field(block, builder.get_argument(), i as u32)?;
         env.symbols.insert(*symbol, value_id);
@@ -417,12 +430,16 @@ fn proc_spec<'a>(
 
     let root = BlockExpr(block, value_id);
     let arg_type_id = layout_spec(
+        &mut env,
         &mut builder,
         interner,
-        &Layout::struct_no_name_order(&argument_layouts),
+        arena.alloc(Layout::struct_no_name_order(
+            argument_layouts.into_bump_slice(),
+        )),
         &WhenRecursive::Unreachable,
     )?;
     let ret_type_id = layout_spec(
+        &mut env,
         &mut builder,
         interner,
         &proc.ret_layout,
@@ -434,41 +451,22 @@ fn proc_spec<'a>(
     Ok((spec, env.type_names))
 }
 
-#[derive(Default)]
 struct Env<'a> {
+    arena: &'a Bump,
     symbols: MutMap<Symbol, ValueId>,
     join_points: MutMap<roc_mono::ir::JoinPointId, morphic_lib::ContinuationId>,
     type_names: MutSet<UnionLayout<'a>>,
 }
 
-fn apply_refcount_operation<'a>(
-    builder: &mut FuncDefBuilder,
-    env: &mut Env<'a>,
-    block: BlockId,
-    modify_rc: &ModifyRc,
-) -> Result<()> {
-    match modify_rc {
-        ModifyRc::Inc(symbol, _) => {
-            let argument = env.symbols[symbol];
-
-            // a recursive touch is never worse for optimizations than a normal touch
-            // and a bit more permissive in its type
-            builder.add_recursive_touch(block, argument)?;
+impl<'a> Env<'a> {
+    fn new(arena: &'a Bump) -> Self {
+        Self {
+            arena,
+            symbols: Default::default(),
+            join_points: Default::default(),
+            type_names: Default::default(),
         }
-
-        ModifyRc::Dec(symbol) => {
-            let argument = env.symbols[symbol];
-
-            builder.add_recursive_touch(block, argument)?;
-        }
-        ModifyRc::DecRef(symbol) => {
-            let argument = env.symbols[symbol];
-
-            builder.add_recursive_touch(block, argument)?;
-        }
-    };
-
-    Ok(())
+    }
 }
 
 fn stmt_spec<'a>(
@@ -476,7 +474,7 @@ fn stmt_spec<'a>(
     interner: &STLayoutInterner<'a>,
     env: &mut Env<'a>,
     block: BlockId,
-    layout: &Layout,
+    layout: &Layout<'a>,
     stmt: &Stmt<'a>,
 ) -> Result<ValueId> {
     use Stmt::*;
@@ -488,25 +486,12 @@ fn stmt_spec<'a>(
 
             let mut queue = vec![symbol];
 
-            loop {
-                match continuation {
-                    Let(symbol, expr, expr_layout, c) => {
-                        let value_id = expr_spec(builder, interner, env, block, expr_layout, expr)?;
-                        env.symbols.insert(*symbol, value_id);
+            while let Let(symbol, expr, expr_layout, c) = continuation {
+                let value_id = expr_spec(builder, interner, env, block, expr_layout, expr)?;
+                env.symbols.insert(*symbol, value_id);
 
-                        queue.push(symbol);
-                        continuation = c;
-                    }
-                    Refcounting(modify_rc, c) => {
-                        // in practice it is common to see a chain of `Let`s interspersed with
-                        // Inc/Dec. For e.g. the False interpreter, this caused stack overflows.
-                        // so we handle RC operations here to limit recursion depth
-                        apply_refcount_operation(builder, env, block, modify_rc)?;
-
-                        continuation = c;
-                    }
-                    _ => break,
-                }
+                queue.push(symbol);
+                continuation = c;
             }
 
             let result = stmt_spec(builder, interner, env, block, layout, continuation)?;
@@ -542,11 +527,32 @@ fn stmt_spec<'a>(
         Expect { remainder, .. } => stmt_spec(builder, interner, env, block, layout, remainder),
         ExpectFx { remainder, .. } => stmt_spec(builder, interner, env, block, layout, remainder),
         Ret(symbol) => Ok(env.symbols[symbol]),
-        Refcounting(modify_rc, continuation) => {
-            apply_refcount_operation(builder, env, block, modify_rc)?;
+        Refcounting(modify_rc, continuation) => match modify_rc {
+            ModifyRc::Inc(symbol, _) => {
+                let argument = env.symbols[symbol];
 
-            stmt_spec(builder, interner, env, block, layout, continuation)
-        }
+                // a recursive touch is never worse for optimizations than a normal touch
+                // and a bit more permissive in its type
+                builder.add_recursive_touch(block, argument)?;
+
+                stmt_spec(builder, interner, env, block, layout, continuation)
+            }
+
+            ModifyRc::Dec(symbol) => {
+                let argument = env.symbols[symbol];
+
+                builder.add_recursive_touch(block, argument)?;
+
+                stmt_spec(builder, interner, env, block, layout, continuation)
+            }
+            ModifyRc::DecRef(symbol) => {
+                let argument = env.symbols[symbol];
+
+                builder.add_recursive_touch(block, argument)?;
+
+                stmt_spec(builder, interner, env, block, layout, continuation)
+            }
+        },
         Join {
             id,
             parameters,
@@ -557,6 +563,7 @@ fn stmt_spec<'a>(
 
             for p in parameters.iter() {
                 type_ids.push(layout_spec(
+                    env,
                     builder,
                     interner,
                     &p.layout,
@@ -564,7 +571,8 @@ fn stmt_spec<'a>(
                 )?);
             }
 
-            let ret_type_id = layout_spec(builder, interner, layout, &WhenRecursive::Unreachable)?;
+            let ret_type_id =
+                layout_spec(env, builder, interner, layout, &WhenRecursive::Unreachable)?;
 
             let jp_arg_type_id = builder.add_tuple_type(&type_ids)?;
 
@@ -605,14 +613,15 @@ fn stmt_spec<'a>(
             builder.add_sub_block(block, BlockExpr(cont_block, cont_value_id))
         }
         Jump(id, symbols) => {
-            let ret_type_id = layout_spec(builder, interner, layout, &WhenRecursive::Unreachable)?;
+            let ret_type_id =
+                layout_spec(env, builder, interner, layout, &WhenRecursive::Unreachable)?;
             let argument = build_tuple_value(builder, env, block, symbols)?;
 
             let jpid = env.join_points[id];
             builder.add_jump(block, jpid, argument, ret_type_id)
         }
         RuntimeError(_) => {
-            let type_id = layout_spec(builder, interner, layout, &WhenRecursive::Unreachable)?;
+            let type_id = layout_spec(env, builder, interner, layout, &WhenRecursive::Unreachable)?;
 
             builder.add_terminate(block, type_id)
         }
@@ -629,10 +638,9 @@ fn build_tuple_value(
 
     for field in symbols.iter() {
         let value_id = match env.symbols.get(field) {
-            None => internal_error!(
+            None => panic!(
                 "Symbol {:?} is not defined in environment {:?}",
-                field,
-                &env.symbols
+                field, &env.symbols
             ),
             Some(x) => *x,
         };
@@ -648,32 +656,34 @@ enum WhenRecursive<'a> {
     Loop(UnionLayout<'a>),
 }
 
-fn build_recursive_tuple_type(
+fn build_recursive_tuple_type<'a>(
+    env: &mut Env<'a>,
     builder: &mut impl TypeContext,
-    interner: &STLayoutInterner,
-    layouts: &[Layout],
+    interner: &STLayoutInterner<'a>,
+    layouts: &[Layout<'a>],
     when_recursive: &WhenRecursive,
 ) -> Result<TypeId> {
     let mut field_types = Vec::new();
 
     for field in layouts.iter() {
-        let type_id = layout_spec_help(builder, interner, field, when_recursive)?;
+        let type_id = layout_spec_help(env, builder, interner, field, when_recursive)?;
         field_types.push(type_id);
     }
 
     builder.add_tuple_type(&field_types)
 }
 
-fn build_tuple_type(
+fn build_tuple_type<'a>(
+    env: &mut Env<'a>,
     builder: &mut impl TypeContext,
-    interner: &STLayoutInterner,
-    layouts: &[Layout],
+    interner: &STLayoutInterner<'a>,
+    layouts: &[Layout<'a>],
     when_recursive: &WhenRecursive,
 ) -> Result<TypeId> {
     let mut field_types = Vec::new();
 
     for field in layouts.iter() {
-        field_types.push(layout_spec(builder, interner, field, when_recursive)?);
+        field_types.push(layout_spec(env, builder, interner, field, when_recursive)?);
     }
 
     builder.add_tuple_type(&field_types)
@@ -705,13 +715,13 @@ fn add_loop(
     builder.add_sub_block(block, BlockExpr(sub_block, unreachable))
 }
 
-fn call_spec(
+fn call_spec<'a>(
     builder: &mut FuncDefBuilder,
-    interner: &STLayoutInterner,
-    env: &Env,
+    interner: &STLayoutInterner<'a>,
+    env: &mut Env<'a>,
     block: BlockId,
-    layout: &Layout,
-    call: &Call,
+    layout: &Layout<'a>,
+    call: &Call<'a>,
 ) -> Result<ValueId> {
     use CallType::*;
 
@@ -743,8 +753,13 @@ fn call_spec(
                 .map(|symbol| env.symbols[symbol])
                 .collect();
 
-            let result_type =
-                layout_spec(builder, interner, ret_layout, &WhenRecursive::Unreachable)?;
+            let result_type = layout_spec(
+                env,
+                builder,
+                interner,
+                ret_layout,
+                &WhenRecursive::Unreachable,
+            )?;
 
             builder.add_unknown_with(block, &arguments, result_type)
         }
@@ -816,6 +831,7 @@ fn call_spec(
                     };
 
                     let output_element_type = layout_spec(
+                        env,
                         builder,
                         interner,
                         return_layout,
@@ -824,6 +840,7 @@ fn call_spec(
 
                     let state_layout = Layout::Builtin(Builtin::List(return_layout));
                     let state_type = layout_spec(
+                        env,
                         builder,
                         interner,
                         &state_layout,
@@ -854,6 +871,7 @@ fn call_spec(
 
                     let state_layout = Layout::Builtin(Builtin::List(&argument_layouts[0]));
                     let state_type = layout_spec(
+                        env,
                         builder,
                         interner,
                         &state_layout,
@@ -883,6 +901,7 @@ fn call_spec(
                     };
 
                     let output_element_type = layout_spec(
+                        env,
                         builder,
                         interner,
                         return_layout,
@@ -891,6 +910,7 @@ fn call_spec(
 
                     let state_layout = Layout::Builtin(Builtin::List(return_layout));
                     let state_type = layout_spec(
+                        env,
                         builder,
                         interner,
                         &state_layout,
@@ -926,6 +946,7 @@ fn call_spec(
                     };
 
                     let output_element_type = layout_spec(
+                        env,
                         builder,
                         interner,
                         return_layout,
@@ -934,6 +955,7 @@ fn call_spec(
 
                     let state_layout = Layout::Builtin(Builtin::List(return_layout));
                     let state_type = layout_spec(
+                        env,
                         builder,
                         interner,
                         &state_layout,
@@ -975,6 +997,7 @@ fn call_spec(
                     };
 
                     let output_element_type = layout_spec(
+                        env,
                         builder,
                         interner,
                         return_layout,
@@ -983,6 +1006,7 @@ fn call_spec(
 
                     let state_layout = Layout::Builtin(Builtin::List(return_layout));
                     let state_type = layout_spec(
+                        env,
                         builder,
                         interner,
                         &state_layout,
@@ -1030,19 +1054,19 @@ fn list_clone(
 }
 
 #[allow(clippy::too_many_arguments)]
-fn lowlevel_spec(
+fn lowlevel_spec<'a>(
     builder: &mut FuncDefBuilder,
-    interner: &STLayoutInterner,
-    env: &Env,
+    interner: &STLayoutInterner<'a>,
+    env: &mut Env<'a>,
     block: BlockId,
-    layout: &Layout,
+    layout: &Layout<'a>,
     op: &LowLevel,
     update_mode: roc_mono::ir::UpdateModeId,
     arguments: &[Symbol],
 ) -> Result<ValueId> {
     use LowLevel::*;
 
-    let type_id = layout_spec(builder, interner, layout, &WhenRecursive::Unreachable)?;
+    let type_id = layout_spec(env, builder, interner, layout, &WhenRecursive::Unreachable)?;
     let mode = update_mode.to_bytes();
     let update_mode_var = UpdateModeVar(&mode);
 
@@ -1151,6 +1175,7 @@ fn lowlevel_spec(
             match layout {
                 Layout::Builtin(Builtin::List(element_layout)) => {
                     let type_id = layout_spec(
+                        env,
                         builder,
                         interner,
                         element_layout,
@@ -1198,28 +1223,31 @@ fn lowlevel_spec(
             // TODO overly pessimstic
             let arguments: Vec<_> = arguments.iter().map(|symbol| env.symbols[symbol]).collect();
 
-            let result_type = layout_spec(builder, interner, layout, &WhenRecursive::Unreachable)?;
+            let result_type =
+                layout_spec(env, builder, interner, layout, &WhenRecursive::Unreachable)?;
 
             builder.add_unknown_with(block, &arguments, result_type)
         }
     }
 }
 
-fn recursive_tag_variant(
+fn recursive_tag_variant<'a>(
+    env: &mut Env<'a>,
     builder: &mut impl TypeContext,
-    interner: &STLayoutInterner,
+    interner: &STLayoutInterner<'a>,
     union_layout: &UnionLayout,
-    fields: &[Layout],
+    fields: &[Layout<'a>],
 ) -> Result<TypeId> {
     let when_recursive = WhenRecursive::Loop(*union_layout);
 
-    build_recursive_tuple_type(builder, interner, fields, &when_recursive)
+    build_recursive_tuple_type(env, builder, interner, fields, &when_recursive)
 }
 
-fn recursive_variant_types(
+fn recursive_variant_types<'a>(
+    env: &mut Env<'a>,
     builder: &mut impl TypeContext,
-    interner: &STLayoutInterner,
-    union_layout: &UnionLayout,
+    interner: &STLayoutInterner<'a>,
+    union_layout: &UnionLayout<'a>,
 ) -> Result<Vec<TypeId>> {
     use UnionLayout::*;
 
@@ -1233,11 +1261,18 @@ fn recursive_variant_types(
             result = Vec::with_capacity(tags.len());
 
             for tag in tags.iter() {
-                result.push(recursive_tag_variant(builder, interner, union_layout, tag)?);
+                result.push(recursive_tag_variant(
+                    env,
+                    builder,
+                    interner,
+                    union_layout,
+                    tag,
+                )?);
             }
         }
         NonNullableUnwrapped(fields) => {
             result = vec![recursive_tag_variant(
+                env,
                 builder,
                 interner,
                 union_layout,
@@ -1253,21 +1288,39 @@ fn recursive_variant_types(
             let cutoff = *nullable_id as usize;
 
             for tag in tags[..cutoff].iter() {
-                result.push(recursive_tag_variant(builder, interner, union_layout, tag)?);
+                result.push(recursive_tag_variant(
+                    env,
+                    builder,
+                    interner,
+                    union_layout,
+                    tag,
+                )?);
             }
 
-            result.push(recursive_tag_variant(builder, interner, union_layout, &[])?);
+            result.push(recursive_tag_variant(
+                env,
+                builder,
+                interner,
+                union_layout,
+                &[],
+            )?);
 
             for tag in tags[cutoff..].iter() {
-                result.push(recursive_tag_variant(builder, interner, union_layout, tag)?);
+                result.push(recursive_tag_variant(
+                    env,
+                    builder,
+                    interner,
+                    union_layout,
+                    tag,
+                )?);
             }
         }
         NullableUnwrapped {
             nullable_id,
             other_fields: fields,
         } => {
-            let unit = recursive_tag_variant(builder, interner, union_layout, &[])?;
-            let other_type = recursive_tag_variant(builder, interner, union_layout, fields)?;
+            let unit = recursive_tag_variant(env, builder, interner, union_layout, &[])?;
+            let other_type = recursive_tag_variant(env, builder, interner, union_layout, fields)?;
 
             if *nullable_id {
                 // nullable_id == 1
@@ -1289,7 +1342,7 @@ fn worst_case_type(context: &mut impl TypeContext) -> Result<TypeId> {
 
 fn expr_spec<'a>(
     builder: &mut FuncDefBuilder,
-    interner: &STLayoutInterner,
+    interner: &STLayoutInterner<'a>,
     env: &mut Env<'a>,
     block: BlockId,
     layout: &Layout<'a>,
@@ -1316,6 +1369,7 @@ fn expr_spec<'a>(
             let value_id = match tag_layout {
                 UnionLayout::NonRecursive(tags) => {
                     let variant_types = non_recursive_variant_types(
+                        env,
                         builder,
                         interner,
                         tags,
@@ -1339,7 +1393,7 @@ fn expr_spec<'a>(
                 UnionLayout::NullableUnwrapped { .. } => data_id,
             };
 
-            let variant_types = recursive_variant_types(builder, interner, tag_layout)?;
+            let variant_types = recursive_variant_types(env, builder, interner, tag_layout)?;
 
             let union_id =
                 builder.add_make_union(block, &variant_types, *tag_id as u32, value_id)?;
@@ -1425,7 +1479,13 @@ fn expr_spec<'a>(
             builder.add_get_tuple_field(block, value_id, *index as u32)
         }
         Array { elem_layout, elems } => {
-            let type_id = layout_spec(builder, interner, elem_layout, &WhenRecursive::Unreachable)?;
+            let type_id = layout_spec(
+                env,
+                builder,
+                interner,
+                elem_layout,
+                &WhenRecursive::Unreachable,
+            )?;
 
             let list = new_list(builder, block, type_id)?;
 
@@ -1453,6 +1513,7 @@ fn expr_spec<'a>(
         EmptyArray => match layout {
             Layout::Builtin(Builtin::List(element_layout)) => {
                 let type_id = layout_spec(
+                    env,
                     builder,
                     interner,
                     element_layout,
@@ -1463,13 +1524,13 @@ fn expr_spec<'a>(
             _ => unreachable!("empty array does not have a list layout"),
         },
         Reset { symbol, .. } => {
-            let type_id = layout_spec(builder, interner, layout, &WhenRecursive::Unreachable)?;
+            let type_id = layout_spec(env, builder, interner, layout, &WhenRecursive::Unreachable)?;
             let value_id = env.symbols[symbol];
 
             builder.add_unknown_with(block, &[value_id], type_id)
         }
         RuntimeErrorFunction(_) => {
-            let type_id = layout_spec(builder, interner, layout, &WhenRecursive::Unreachable)?;
+            let type_id = layout_spec(env, builder, interner, layout, &WhenRecursive::Unreachable)?;
 
             builder.add_terminate(block, type_id)
         }
@@ -1496,45 +1557,55 @@ fn literal_spec(
     }
 }
 
-fn layout_spec(
+fn layout_spec<'a>(
+    env: &mut Env<'a>,
     builder: &mut impl TypeContext,
-    interner: &STLayoutInterner,
-    layout: &Layout,
+    interner: &STLayoutInterner<'a>,
+    layout: &Layout<'a>,
     when_recursive: &WhenRecursive,
 ) -> Result<TypeId> {
-    layout_spec_help(builder, interner, layout, when_recursive)
+    layout_spec_help(env, builder, interner, layout, when_recursive)
 }
 
-fn non_recursive_variant_types(
+fn non_recursive_variant_types<'a>(
+    env: &mut Env<'a>,
     builder: &mut impl TypeContext,
-    interner: &STLayoutInterner,
-    tags: &[&[Layout]],
+    interner: &STLayoutInterner<'a>,
+    tags: &[&[Layout<'a>]],
     // If there is a recursive pointer latent within this layout, coming from a containing layout.
     when_recursive: &WhenRecursive,
 ) -> Result<Vec<TypeId>> {
     let mut result = Vec::with_capacity(tags.len());
 
     for tag in tags.iter() {
-        result.push(build_tuple_type(builder, interner, tag, when_recursive)?);
+        result.push(build_tuple_type(
+            env,
+            builder,
+            interner,
+            tag,
+            when_recursive,
+        )?);
     }
 
     Ok(result)
 }
 
-fn layout_spec_help(
+fn layout_spec_help<'a>(
+    env: &mut Env<'a>,
     builder: &mut impl TypeContext,
-    interner: &STLayoutInterner,
-    layout: &Layout,
+    interner: &STLayoutInterner<'a>,
+    layout: &Layout<'a>,
     when_recursive: &WhenRecursive,
 ) -> Result<TypeId> {
     use Layout::*;
 
     match layout {
-        Builtin(builtin) => builtin_spec(builder, interner, builtin, when_recursive),
+        Builtin(builtin) => builtin_spec(env, builder, interner, builtin, when_recursive),
         Struct { field_layouts, .. } => {
-            build_recursive_tuple_type(builder, interner, field_layouts, when_recursive)
+            build_recursive_tuple_type(env, builder, interner, field_layouts, when_recursive)
         }
         LambdaSet(lambda_set) => layout_spec_help(
+            env,
             builder,
             interner,
             &lambda_set.runtime_representation(interner),
@@ -1550,7 +1621,7 @@ fn layout_spec_help(
                 }
                 UnionLayout::NonRecursive(tags) => {
                     let variant_types =
-                        non_recursive_variant_types(builder, interner, tags, when_recursive)?;
+                        non_recursive_variant_types(env, builder, interner, tags, when_recursive)?;
                     builder.add_union_type(&variant_types)
                 }
                 UnionLayout::Recursive(_)
@@ -1560,13 +1631,16 @@ fn layout_spec_help(
                     let type_name_bytes = recursive_tag_union_name_bytes(union_layout).as_bytes();
                     let type_name = TypeName(&type_name_bytes);
 
+                    env.type_names.insert(*union_layout);
+
                     Ok(builder.add_named_type(MOD_APP, type_name))
                 }
             }
         }
 
         Boxed(inner_layout) => {
-            let inner_type = layout_spec_help(builder, interner, inner_layout, when_recursive)?;
+            let inner_type =
+                layout_spec_help(env, builder, interner, inner_layout, when_recursive)?;
             let cell_type = builder.add_heap_cell_type();
 
             builder.add_tuple_type(&[cell_type, inner_type])
@@ -1591,10 +1665,11 @@ fn layout_spec_help(
     }
 }
 
-fn builtin_spec(
+fn builtin_spec<'a>(
+    env: &mut Env<'a>,
     builder: &mut impl TypeContext,
-    interner: &STLayoutInterner,
-    builtin: &Builtin,
+    interner: &STLayoutInterner<'a>,
+    builtin: &Builtin<'a>,
     when_recursive: &WhenRecursive,
 ) -> Result<TypeId> {
     use Builtin::*;
@@ -1604,7 +1679,8 @@ fn builtin_spec(
         Decimal | Float(_) => builder.add_tuple_type(&[]),
         Str => str_type(builder),
         List(element_layout) => {
-            let element_type = layout_spec_help(builder, interner, element_layout, when_recursive)?;
+            let element_type =
+                layout_spec_help(env, builder, interner, element_layout, when_recursive)?;
 
             let cell = builder.add_heap_cell_type();
             let bag = builder.add_bag_type(element_type)?;