generate callers based on layout traversal

crates/compiler/mono/src/ir.rs

@@ -34,7 +34,7 @@ use roc_std::RocDec;
 use roc_target::TargetInfo;
 use roc_types::subs::{
     instantiate_rigids, storage_copy_var_to, Content, ExhaustiveMark, FlatType, RedundantMark,
-    StorageSubs, Subs, Variable, VariableSubsSlice,
+    StorageSubs, Subs,  Variable, VariableSubsSlice,
 };
 use std::collections::HashMap;
 use ven_pretty::{BoxAllocator, DocAllocator, DocBuilder};
@@ -305,12 +305,21 @@ pub struct Proc<'a> {
     pub host_exposed_layouts: HostExposedLayouts<'a>,
 }
 
+#[derive(Clone, Debug, PartialEq, Eq)]
+pub struct HostExposedLambdaSet<'a> {
+    pub id: LambdaSetId,
+    /// Symbol of the exposed function
+    pub symbol: Symbol,
+    pub proc_layout: ProcLayout<'a>,
+    pub raw_function_layout: RawFunctionLayout<'a>,
+}
+
 #[derive(Clone, Debug, PartialEq, Eq)]
 pub enum HostExposedLayouts<'a> {
     NotHostExposed,
     HostExposed {
         rigids: BumpMap<Lowercase, InLayout<'a>>,
-        aliases: BumpMap<Symbol, (Symbol, ProcLayout<'a>, RawFunctionLayout<'a>)>,
+        aliases: BumpMap<Symbol, HostExposedLambdaSet<'a>>,
     },
 }
 
@@ -3068,25 +3077,126 @@ fn specialize_external_help<'a>(
     );
 
     match specialization_result {
-        Ok((proc, layout)) => {
+        Ok((mut proc, layout)) => {
             let top_level = ProcLayout::from_raw_named(env.arena, name, layout);
 
             if procs.is_module_thunk(name.name()) {
                 debug_assert!(top_level.arguments.is_empty());
             }
 
-            // layouts that are (transitively) used in the type of `mainForHost`.
-            let host_exposed_layouts = top_level
-                .arguments
-                .iter()
-                .copied()
-                .chain([top_level.result]);
+            if !host_exposed_aliases.is_empty() {
+                // layouts that are (transitively) used in the type of `mainForHost`.
+                let mut host_exposed_layouts: Vec<_> = top_level
+                    .arguments
+                    .iter()
+                    .copied()
+                    .chain([top_level.result])
+                    .collect_in(env.arena);
 
-            // TODO: In the future, we will generate glue procs here
-            for in_layout in host_exposed_layouts {
-                let layout = layout_cache.interner.get(in_layout);
+                // it is very likely we see the same types across functions, or in multiple arguments
+                host_exposed_layouts.sort();
+                host_exposed_layouts.dedup();
 
-                let _ = layout;
+                // TODO: In the future, we will generate glue procs here
+                for in_layout in host_exposed_layouts {
+                    let layout = layout_cache.interner.get(in_layout);
+
+                    let all_glue_procs = generate_glue_procs(
+                        env.home,
+                        env.ident_ids,
+                        env.arena,
+                        &mut layout_cache.interner,
+                        env.arena.alloc(layout),
+                    );
+
+                    // for now, getters are not processed here
+                    let GlueProcs {
+                        getters: _,
+                        extern_names,
+                    } = all_glue_procs;
+
+                    let mut aliases = BumpMap::default();
+
+                    for (id, mut raw_function_layout) in extern_names {
+                        let symbol = env.unique_symbol();
+                        let lambda_name = LambdaName::no_niche(symbol);
+
+                        if false {
+                            raw_function_layout = match raw_function_layout {
+                                RawFunctionLayout::Function(a, mut lambda_set, _) => {
+                                    lambda_set.ret = in_layout;
+                                    RawFunctionLayout::Function(a, lambda_set, in_layout)
+                                }
+                                RawFunctionLayout::ZeroArgumentThunk(x) => {
+                                    RawFunctionLayout::ZeroArgumentThunk(x)
+                                }
+                            };
+                        }
+
+                        let (key, (_name, top_level, proc)) = generate_host_exposed_function(
+                            env,
+                            procs,
+                            layout_cache,
+                            lambda_name,
+                            raw_function_layout,
+                        );
+
+                        procs
+                            .specialized
+                            .insert_specialized(symbol, top_level, proc);
+
+                        let hels = HostExposedLambdaSet {
+                            id,
+                            symbol,
+                            proc_layout: top_level,
+                            raw_function_layout,
+                        };
+
+                        aliases.insert(key, hels);
+                    }
+
+                    // pre-glue: generate named callers for as-exposed aliases
+                    for (alias_name, variable) in host_exposed_aliases {
+                        let raw_function_layout = layout_cache
+                            .raw_from_var(env.arena, *variable, env.subs)
+                            .unwrap();
+
+                        let symbol = env.unique_symbol();
+                        let lambda_name = LambdaName::no_niche(symbol);
+
+                        let (_lambda_name, (proc_name, proc_layout, proc)) =
+                            generate_host_exposed_function(
+                                env,
+                                procs,
+                                layout_cache,
+                                lambda_name,
+                                raw_function_layout,
+                            );
+
+                        procs
+                            .specialized
+                            .insert_specialized(proc_name, proc_layout, proc);
+
+                        let hels = HostExposedLambdaSet {
+                            id: LambdaSetId::default(),
+                            symbol: proc_name,
+                            proc_layout,
+                            raw_function_layout,
+                        };
+
+                        aliases.insert(*alias_name, hels);
+                    }
+
+                    match &mut proc.host_exposed_layouts {
+                        HostExposedLayouts::HostExposed { aliases: old, .. } => old.extend(aliases),
+                        hep @ HostExposedLayouts::NotHostExposed => {
+                            *hep = HostExposedLayouts::HostExposed {
+                                aliases,
+                                rigids: Default::default(),
+                            };
+                        }
+                    }
+                }
             }
 
             procs
@@ -3221,7 +3331,7 @@ fn generate_host_exposed_function<'a>(
     layout: RawFunctionLayout<'a>,
 ) -> (Symbol, (Symbol, ProcLayout<'a>, Proc<'a>)) {
     let symbol = lambda_name.name();
-    let function_name = env.unique_symbol();
+    let function_name = symbol;
 
     match layout {
         RawFunctionLayout::Function(_, lambda_set, _) => {
@@ -3379,7 +3489,7 @@ fn specialize_proc_help<'a>(
     let body_var = partial_proc.body_var;
 
     // determine the layout of aliases/rigids exposed to the host
-    let host_exposed_layouts = if host_exposed_variables.is_empty() {
+    let host_exposed_layouts = if true || host_exposed_variables.is_empty() {
         HostExposedLayouts::NotHostExposed
     } else {
         let mut aliases = BumpMap::new_in(env.arena);
@@ -3400,12 +3510,20 @@ fn specialize_proc_help<'a>(
             procs
                 .specialized
                 .insert_specialized(proc_name, proc_layout, proc);
-            aliases.insert(*alias_name, (proc_name, proc_layout, raw_function_layout));
+
+            let hels = HostExposedLambdaSet {
+                id: LambdaSetId::default(),
+                symbol: proc_name,
+                proc_layout,
+                raw_function_layout,
+            };
+
+            aliases.insert(*alias_name, hels);
         }
 
         HostExposedLayouts::HostExposed {
             rigids: BumpMap::new_in(env.arena),
-            aliases,
+            aliases: Default::default(),
         }
     };
 
@@ -11129,3 +11247,344 @@ where
         remainder: env.arena.alloc(switch),
     }
 }
+
+pub struct GlueLayouts<'a> {
+    pub getters: std::vec::Vec<(Symbol, ProcLayout<'a>)>,
+}
+
+type GlueProcId = u16;
+
+pub struct GlueProc<'a> {
+    pub name: Symbol,
+    pub proc_layout: ProcLayout<'a>,
+    pub proc: Proc<'a>,
+}
+
+pub struct GlueProcs<'a> {
+    pub getters: Vec<'a, (Layout<'a>, Vec<'a, GlueProc<'a>>)>,
+    pub extern_names: Vec<'a, (LambdaSetId, RawFunctionLayout<'a>)>,
+}
+
+#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Default)]
+pub struct LambdaSetId(pub u32);
+
+impl LambdaSetId {
+    #[must_use]
+    pub fn next(self) -> Self {
+        debug_assert!(self.0 < u32::MAX);
+        Self(self.0 + 1)
+    }
+
+    pub const fn invalid() -> Self {
+        Self(u32::MAX)
+    }
+}
+
+pub fn generate_glue_procs<'a, 'i, I>(
+    home: ModuleId,
+    ident_ids: &mut IdentIds,
+    arena: &'a Bump,
+    layout_interner: &'i mut I,
+    layout: &'a Layout<'a>,
+) -> GlueProcs<'a>
+where
+    I: LayoutInterner<'a>,
+{
+    let mut answer = GlueProcs {
+        getters: Vec::new_in(arena),
+        extern_names: Vec::new_in(arena),
+    };
+
+    let mut lambda_set_id = LambdaSetId(0);
+
+    let mut stack: Vec<'a, Layout<'a>> = Vec::from_iter_in([*layout], arena);
+    let mut next_unique_id = 0;
+
+    macro_rules! handle_struct_field_layouts {
+        ($field_layouts: expr) => {{
+            if $field_layouts.iter().any(|l| {
+                layout_interner
+                    .get(*l)
+                    .has_varying_stack_size(layout_interner, arena)
+            }) {
+                let procs = generate_glue_procs_for_struct_fields(
+                    layout_interner,
+                    home,
+                    &mut next_unique_id,
+                    ident_ids,
+                    arena,
+                    layout,
+                    $field_layouts,
+                );
+
+                answer.getters.push((*layout, procs));
+            }
+
+            for in_layout in $field_layouts.iter().rev() {
+                stack.push(layout_interner.get(*in_layout));
+            }
+        }};
+    }
+
+    macro_rules! handle_tag_field_layouts {
+        ($tag_id:expr, $union_layout:expr, $field_layouts: expr) => {{
+            if $field_layouts.iter().any(|l| {
+                layout_interner
+                    .get(*l)
+                    .has_varying_stack_size(layout_interner, arena)
+            }) {
+                let procs = generate_glue_procs_for_tag_fields(
+                    layout_interner,
+                    home,
+                    &mut next_unique_id,
+                    ident_ids,
+                    arena,
+                    $tag_id,
+                    layout,
+                    $union_layout,
+                    $field_layouts,
+                );
+
+                answer.getters.push((*layout, procs));
+            }
+
+            for in_layout in $field_layouts.iter().rev() {
+                stack.push(layout_interner.get(*in_layout));
+            }
+        }};
+    }
+
+    while let Some(layout) = stack.pop() {
+        match layout {
+            Layout::Builtin(builtin) => match builtin {
+                Builtin::Int(_)
+                | Builtin::Float(_)
+                | Builtin::Bool
+                | Builtin::Decimal
+                | Builtin::Str => { /* do nothing */ }
+                Builtin::List(element) => stack.push(layout_interner.get(element)),
+            },
+            Layout::Struct { field_layouts, .. } => {
+                handle_struct_field_layouts!(field_layouts);
+            }
+            Layout::Boxed(boxed) => {
+                stack.push(layout_interner.get(boxed));
+            }
+            Layout::Union(union_layout) => match union_layout {
+                UnionLayout::NonRecursive(tags) => {
+                    for in_layout in tags.iter().flat_map(|e| e.iter()) {
+                        stack.push(layout_interner.get(*in_layout));
+                    }
+                }
+                UnionLayout::Recursive(tags) => {
+                    for in_layout in tags.iter().flat_map(|e| e.iter()) {
+                        stack.push(layout_interner.get(*in_layout));
+                    }
+                }
+                UnionLayout::NonNullableUnwrapped(field_layouts) => {
+                    handle_tag_field_layouts!(0, union_layout, field_layouts);
+                }
+                UnionLayout::NullableWrapped {
+                    other_tags,
+                    nullable_id,
+                } => {
+                    let tag_ids =
+                        (0..nullable_id).chain(nullable_id + 1..other_tags.len() as u16 + 1);
+                    for (i, field_layouts) in tag_ids.zip(other_tags) {
+                        handle_tag_field_layouts!(i, union_layout, *field_layouts);
+                    }
+                }
+                UnionLayout::NullableUnwrapped { other_fields, .. } => {
+                    for in_layout in other_fields.iter().rev() {
+                        stack.push(layout_interner.get(*in_layout));
+                    }
+                }
+            },
+            Layout::LambdaSet(lambda_set) => {
+                let raw_function_layout =
+                    RawFunctionLayout::Function(*lambda_set.args, lambda_set, lambda_set.ret);
+
+                let key = (lambda_set_id, raw_function_layout);
+                answer.extern_names.push(key);
+
+                // this id is used, increment for the next one
+                lambda_set_id = lambda_set_id.next();
+
+                stack.push(layout_interner.get(lambda_set.runtime_representation()));
+            }
+            Layout::RecursivePointer(_) => {
+                /* do nothing, we've already generated for this type through the Union(_) */
+            }
+        }
+    }
+
+    answer
+}
+
+fn generate_glue_procs_for_struct_fields<'a, 'i, I>(
+    layout_interner: &'i mut I,
+    home: ModuleId,
+    next_unique_id: &mut GlueProcId,
+    ident_ids: &mut IdentIds,
+    arena: &'a Bump,
+    unboxed_struct_layout: &'a Layout<'a>,
+    field_layouts: &'a [InLayout<'a>],
+) -> Vec<'a, GlueProc<'a>>
+where
+    I: LayoutInterner<'a>,
+{
+    let interned_unboxed_struct_layout = layout_interner.insert(*unboxed_struct_layout);
+    let boxed_struct_layout = Layout::Boxed(interned_unboxed_struct_layout);
+    let boxed_struct_layout = layout_interner.insert(boxed_struct_layout);
+    let mut answer = bumpalo::collections::Vec::with_capacity_in(field_layouts.len(), arena);
+
+    for (index, field) in field_layouts.iter().enumerate() {
+        let proc_layout = ProcLayout {
+            arguments: arena.alloc([boxed_struct_layout]),
+            result: *field,
+            niche: Niche::NONE,
+        };
+
+        let symbol = unique_glue_symbol(arena, next_unique_id, home, ident_ids);
+
+        let argument = Symbol::new(home, ident_ids.gen_unique());
+        let unboxed = Symbol::new(home, ident_ids.gen_unique());
+        let result = Symbol::new(home, ident_ids.gen_unique());
+
+        home.register_debug_idents(ident_ids);
+
+        let ret_stmt = arena.alloc(Stmt::Ret(result));
+
+        let field_get_expr = Expr::StructAtIndex {
+            index: index as u64,
+            field_layouts,
+            structure: unboxed,
+        };
+
+        let field_get_stmt = Stmt::Let(result, field_get_expr, *field, ret_stmt);
+
+        let unbox_expr = Expr::ExprUnbox { symbol: argument };
+
+        let unbox_stmt = Stmt::Let(
+            unboxed,
+            unbox_expr,
+            interned_unboxed_struct_layout,
+            arena.alloc(field_get_stmt),
+        );
+
+        let proc = Proc {
+            name: LambdaName::no_niche(symbol),
+            args: arena.alloc([(boxed_struct_layout, argument)]),
+            body: unbox_stmt,
+            closure_data_layout: None,
+            ret_layout: *field,
+            is_self_recursive: SelfRecursive::NotSelfRecursive,
+            must_own_arguments: false,
+            host_exposed_layouts: HostExposedLayouts::NotHostExposed,
+        };
+
+        answer.push(GlueProc {
+            name: symbol,
+            proc_layout,
+            proc,
+        });
+    }
+
+    answer
+}
+
+fn unique_glue_symbol(
+    arena: &Bump,
+    next_unique_id: &mut GlueProcId,
+    home: ModuleId,
+    ident_ids: &mut IdentIds,
+) -> Symbol {
+    let unique_id = *next_unique_id;
+
+    *next_unique_id = unique_id + 1;
+
+    // Turn unique_id into a Symbol without doing a heap allocation.
+    use std::fmt::Write;
+    let mut string = bumpalo::collections::String::with_capacity_in(32, arena);
+    let _result = write!(
+        &mut string,
+        "roc__getter_{}_{}",
+        "", // then name of the platform `main.roc` is the empty string
+        unique_id
+    );
+    debug_assert_eq!(_result, Ok(())); // This should never fail, but doesn't hurt to debug-check!
+
+    let ident_id = ident_ids.get_or_insert(string.into_bump_str());
+
+    Symbol::new(home, ident_id)
+}
+
+#[allow(clippy::too_many_arguments)]
+fn generate_glue_procs_for_tag_fields<'a, 'i, I>(
+    layout_interner: &'i mut I,
+    home: ModuleId,
+    next_unique_id: &mut GlueProcId,
+    ident_ids: &mut IdentIds,
+    arena: &'a Bump,
+    tag_id: TagIdIntType,
+    unboxed_struct_layout: &'a Layout<'a>,
+    union_layout: UnionLayout<'a>,
+    field_layouts: &'a [InLayout<'a>],
+) -> Vec<'a, GlueProc<'a>>
+where
+    I: LayoutInterner<'a>,
+{
+    let interned = layout_interner.insert(*unboxed_struct_layout);
+    let boxed_struct_layout = Layout::Boxed(interned);
+    let boxed_struct_layout = layout_interner.insert(boxed_struct_layout);
+    let mut answer = bumpalo::collections::Vec::with_capacity_in(field_layouts.len(), arena);
+
+    for (index, field) in field_layouts.iter().enumerate() {
+        let proc_layout = ProcLayout {
+            arguments: arena.alloc([boxed_struct_layout]),
+            result: *field,
+            niche: Niche::NONE,
+        };
+        let symbol = unique_glue_symbol(arena, next_unique_id, home, ident_ids);
+
+        let argument = Symbol::new(home, ident_ids.gen_unique());
+        let unboxed = Symbol::new(home, ident_ids.gen_unique());
+        let result = Symbol::new(home, ident_ids.gen_unique());
+
+        home.register_debug_idents(ident_ids);
+
+        let ret_stmt = arena.alloc(Stmt::Ret(result));
+
+        let field_get_expr = Expr::UnionAtIndex {
+            structure: unboxed,
+            tag_id,
+            union_layout,
+            index: index as u64,
+        };
+
+        let field_get_stmt = Stmt::Let(result, field_get_expr, *field, ret_stmt);
+
+        let unbox_expr = Expr::ExprUnbox { symbol: argument };
+
+        let unbox_stmt = Stmt::Let(unboxed, unbox_expr, interned, arena.alloc(field_get_stmt));
+
+        let proc = Proc {
+            name: LambdaName::no_niche(symbol),
+            args: arena.alloc([(boxed_struct_layout, argument)]),
+            body: unbox_stmt,
+            closure_data_layout: None,
+            ret_layout: *field,
+            is_self_recursive: SelfRecursive::NotSelfRecursive,
+            must_own_arguments: false,
+            host_exposed_layouts: HostExposedLayouts::NotHostExposed,
+        };
+
+        answer.push(GlueProc {
+            name: symbol,
+            proc_layout,
+            proc,
+        });
+    }
+
+    answer
+}