Merge branch 'main' into symbol_tags

crates/compiler/mono/src/borrow.rs

@@ -1031,7 +1031,7 @@ pub fn lowlevel_borrow_signature(arena: &Bump, op: LowLevel) -> &[Ownership] {
             unreachable!("These lowlevel operations are turned into mono Expr's")
         }
 
-        PtrCast | RefCountInc | RefCountDec => {
+        PtrCast | PtrWrite | RefCountInc | RefCountDec => {
             unreachable!("Only inserted *after* borrow checking: {:?}", op);
         }
     }

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

@@ -5,8 +5,8 @@ use roc_module::symbol::{IdentIds, ModuleId, Symbol};
 use roc_target::TargetInfo;
 
 use crate::ir::{
-    Call, CallSpecId, CallType, Expr, HostExposedLayouts, JoinPointId, ModifyRc, Proc, ProcLayout,
-    SelfRecursive, Stmt, UpdateModeId,
+    Call, CallSpecId, CallType, Expr, HostExposedLayouts, JoinPointId, ModifyRc, PassedFunction,
+    Proc, ProcLayout, SelfRecursive, Stmt, UpdateModeId,
 };
 use crate::layout::{
     Builtin, InLayout, LambdaName, Layout, LayoutInterner, Niche, STLayoutInterner, UnionLayout,
@@ -20,6 +20,7 @@ const LAYOUT_UNIT: InLayout = Layout::UNIT;
 
 const ARG_1: Symbol = Symbol::ARG_1;
 const ARG_2: Symbol = Symbol::ARG_2;
+const ARG_3: Symbol = Symbol::ARG_3;
 
 /// "Infinite" reference count, for static values
 /// Ref counts are encoded as negative numbers where isize::MIN represents 1
@@ -580,6 +581,175 @@ impl<'a> CodeGenHelp<'a> {
     }
 }
 
+pub struct CallerProc<'a> {
+    pub proc_symbol: Symbol,
+    pub proc_layout: ProcLayout<'a>,
+    pub proc: Proc<'a>,
+}
+
+impl<'a> CallerProc<'a> {
+    fn create_symbol(home: ModuleId, ident_ids: &mut IdentIds, debug_name: &str) -> Symbol {
+        let ident_id = ident_ids.add_str(debug_name);
+        Symbol::new(home, ident_id)
+    }
+
+    fn create_caller_proc_symbol(
+        home: ModuleId,
+        ident_ids: &mut IdentIds,
+        operation: &str,
+        wrapped_function: Symbol,
+    ) -> Symbol {
+        let debug_name = format!("#help_{}_{}_{:?}", "caller", operation, wrapped_function,);
+
+        Self::create_symbol(home, ident_ids, &debug_name)
+    }
+
+    pub fn new(
+        arena: &'a Bump,
+        home: ModuleId,
+        ident_ids: &mut IdentIds,
+        layout_interner: &mut STLayoutInterner<'a>,
+        passed_function: &PassedFunction<'a>,
+        capture_layout: Option<InLayout<'a>>,
+    ) -> Self {
+        let mut ctx = Context {
+            new_linker_data: Vec::new_in(arena),
+            recursive_union: None,
+            op: HelperOp::Eq,
+        };
+
+        let box_capture_layout = if let Some(capture_layout) = capture_layout {
+            layout_interner.insert(Layout::Boxed(capture_layout))
+        } else {
+            layout_interner.insert(Layout::Boxed(Layout::UNIT))
+        };
+
+        let box_argument_layout =
+            layout_interner.insert(Layout::Boxed(passed_function.argument_layouts[0]));
+
+        let box_return_layout =
+            layout_interner.insert(Layout::Boxed(passed_function.return_layout));
+
+        let proc_layout = ProcLayout {
+            arguments: arena.alloc([box_capture_layout, box_argument_layout, box_return_layout]),
+            result: Layout::UNIT,
+            niche: Niche::NONE,
+        };
+
+        let proc_symbol =
+            Self::create_caller_proc_symbol(home, ident_ids, "map", passed_function.name.name());
+
+        ctx.new_linker_data.push((proc_symbol, proc_layout));
+
+        let unbox_capture = Expr::ExprUnbox {
+            symbol: Symbol::ARG_1,
+        };
+
+        let unbox_argument = Expr::ExprUnbox {
+            symbol: Symbol::ARG_2,
+        };
+
+        let unboxed_capture = Self::create_symbol(home, ident_ids, "unboxed_capture");
+        let unboxed_argument = Self::create_symbol(home, ident_ids, "unboxed_argument");
+        let call_result = Self::create_symbol(home, ident_ids, "call_result");
+        let unit_symbol = Self::create_symbol(home, ident_ids, "unit_symbol");
+        let ignored = Self::create_symbol(home, ident_ids, "ignored");
+
+        let call = Expr::Call(Call {
+            call_type: CallType::ByName {
+                name: passed_function.name,
+                ret_layout: passed_function.return_layout,
+                arg_layouts: passed_function.argument_layouts,
+                specialization_id: passed_function.specialization_id,
+            },
+            arguments: if capture_layout.is_some() {
+                arena.alloc([unboxed_argument, unboxed_capture])
+            } else {
+                arena.alloc([unboxed_argument])
+            },
+        });
+
+        let ptr_write = Expr::Call(Call {
+            call_type: CallType::LowLevel {
+                op: LowLevel::PtrWrite,
+                update_mode: UpdateModeId::BACKEND_DUMMY,
+            },
+            arguments: arena.alloc([Symbol::ARG_3, call_result]),
+        });
+
+        let mut body = Stmt::Let(
+            unboxed_argument,
+            unbox_argument,
+            passed_function.argument_layouts[0],
+            arena.alloc(Stmt::Let(
+                call_result,
+                call,
+                passed_function.return_layout,
+                arena.alloc(Stmt::Let(
+                    ignored,
+                    ptr_write,
+                    box_return_layout,
+                    arena.alloc(Stmt::Let(
+                        unit_symbol,
+                        Expr::Struct(&[]),
+                        Layout::UNIT,
+                        arena.alloc(Stmt::Ret(unit_symbol)),
+                    )),
+                )),
+            )),
+        );
+
+        if let Some(capture_layout) = capture_layout {
+            body = Stmt::Let(
+                unboxed_capture,
+                unbox_capture,
+                capture_layout,
+                arena.alloc(body),
+            );
+        }
+
+        let args: &'a [(InLayout<'a>, Symbol)] = {
+            arena.alloc([
+                (box_capture_layout, ARG_1),
+                (box_argument_layout, ARG_2),
+                (box_return_layout, ARG_3),
+            ])
+        };
+
+        let proc = Proc {
+            name: LambdaName::no_niche(proc_symbol),
+            args,
+            body,
+            closure_data_layout: None,
+            ret_layout: Layout::UNIT,
+            is_self_recursive: SelfRecursive::NotSelfRecursive,
+            host_exposed_layouts: HostExposedLayouts::NotHostExposed,
+        };
+
+        if false {
+            let allocator = ven_pretty::BoxAllocator;
+            let doc = proc
+                .to_doc::<_, (), _>(
+                    &allocator,
+                    layout_interner,
+                    true,
+                    crate::ir::Parens::NotNeeded,
+                )
+                .1
+                .pretty(80)
+                .to_string();
+
+            println!("{}", doc);
+        }
+
+        Self {
+            proc_symbol,
+            proc_layout,
+            proc,
+        }
+    }
+}
+
 fn let_lowlevel<'a>(
     arena: &'a Bump,
     result_layout: InLayout<'a>,

crates/compiler/mono/src/ir.rs

@@ -63,7 +63,6 @@ pub fn pretty_print_ir_symbols() -> bool {
 // if your changes cause this number to go down, great!
 // please change it to the lower number.
 // if it went up, maybe check that the change is really required
-
 roc_error_macros::assert_sizeof_wasm!(Literal, 24);
 roc_error_macros::assert_sizeof_wasm!(Expr, 48);
 roc_error_macros::assert_sizeof_wasm!(Stmt, 64);
@@ -3027,7 +3026,7 @@ fn specialize_external_help<'a>(
                 for in_layout in host_exposed_layouts {
                     let layout = layout_cache.interner.get(in_layout);
 
-                    let all_glue_procs = generate_glue_procs(
+                    let mut all_glue_procs = generate_glue_procs(
                         env.home,
                         env.ident_ids,
                         env.arena,
@@ -3035,6 +3034,17 @@ fn specialize_external_help<'a>(
                         env.arena.alloc(layout),
                     );
 
+                    all_glue_procs.extern_names = {
+                        let mut layout_env = layout::Env::from_components(
+                            layout_cache,
+                            env.subs,
+                            env.arena,
+                            env.target_info,
+                        );
+
+                        find_lambda_sets(&mut layout_env, variable)
+                    };
+
                     // for now, getters are not processed here
                     let GlueProcs {
                         getters,
@@ -3053,23 +3063,10 @@ fn specialize_external_help<'a>(
 
                     let mut aliases = BumpMap::default();
 
-                    for (id, mut raw_function_layout) in extern_names {
+                    for (id, raw_function_layout) in extern_names {
                         let symbol = env.unique_symbol();
                         let lambda_name = LambdaName::no_niche(symbol);
 
-                        // fix the recursion in the rocLovesRust example
-                        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, (top_level, proc)) = generate_host_exposed_function(
                             env,
                             procs,
@@ -9521,6 +9518,126 @@ impl LambdaSetId {
     }
 }
 
+fn find_lambda_sets<'a>(
+    env: &mut crate::layout::Env<'a, '_>,
+    initial: Variable,
+) -> Vec<'a, (LambdaSetId, RawFunctionLayout<'a>)> {
+    let mut stack = bumpalo::collections::Vec::new_in(env.arena);
+
+    // ignore the lambda set of top-level functions
+    match env.subs.get_without_compacting(initial).content {
+        Content::Structure(FlatType::Func(arguments, _, result)) => {
+            let arguments = &env.subs.variables[arguments.indices()];
+
+            stack.extend(arguments.iter().copied());
+            stack.push(result);
+        }
+        _ => {
+            stack.push(initial);
+        }
+    }
+
+    let lambda_set_variables = find_lambda_sets_help(env.subs, stack);
+    let mut answer =
+        bumpalo::collections::Vec::with_capacity_in(lambda_set_variables.len(), env.arena);
+
+    for (variable, lambda_set_id) in lambda_set_variables {
+        let lambda_set = env.subs.get_lambda_set(variable);
+        let raw_function_layout = RawFunctionLayout::from_var(env, lambda_set.ambient_function)
+            .value()
+            .unwrap();
+
+        let key = (lambda_set_id, raw_function_layout);
+        answer.push(key);
+    }
+
+    answer
+}
+
+pub fn find_lambda_sets_help(
+    subs: &Subs,
+    mut stack: Vec<'_, Variable>,
+) -> MutMap<Variable, LambdaSetId> {
+    use roc_types::subs::GetSubsSlice;
+
+    let mut lambda_set_id = LambdaSetId::default();
+
+    let mut result = MutMap::default();
+
+    while let Some(var) = stack.pop() {
+        match subs.get_content_without_compacting(var) {
+            Content::RangedNumber(_)
+            | Content::Error
+            | Content::FlexVar(_)
+            | Content::RigidVar(_)
+            | Content::FlexAbleVar(_, _)
+            | Content::RigidAbleVar(_, _)
+            | Content::RecursionVar { .. } => {}
+            Content::Structure(flat_type) => match flat_type {
+                FlatType::Apply(_, arguments) => {
+                    stack.extend(subs.get_subs_slice(*arguments).iter().rev());
+                }
+                FlatType::Func(arguments, lambda_set_var, ret_var) => {
+                    result.insert(*lambda_set_var, lambda_set_id);
+                    lambda_set_id = lambda_set_id.next();
+
+                    let arguments = &subs.variables[arguments.indices()];
+
+                    stack.extend(arguments.iter().copied());
+                    stack.push(*lambda_set_var);
+                    stack.push(*ret_var);
+                }
+                FlatType::Record(fields, ext) => {
+                    stack.extend(subs.get_subs_slice(fields.variables()).iter().rev());
+                    stack.push(*ext);
+                }
+                FlatType::Tuple(elements, ext) => {
+                    stack.extend(subs.get_subs_slice(elements.variables()).iter().rev());
+                    stack.push(*ext);
+                }
+                FlatType::FunctionOrTagUnion(_, _, ext) => {
+                    // just the ext
+                    match ext {
+                        roc_types::subs::TagExt::Openness(var) => stack.push(*var),
+                        roc_types::subs::TagExt::Any(_) => { /* ignore */ }
+                    }
+                }
+                FlatType::TagUnion(union_tags, ext)
+                | FlatType::RecursiveTagUnion(_, union_tags, ext) => {
+                    for tag in union_tags.variables() {
+                        stack.extend(
+                            subs.get_subs_slice(subs.variable_slices[tag.index as usize])
+                                .iter()
+                                .rev(),
+                        );
+                    }
+
+                    match ext {
+                        roc_types::subs::TagExt::Openness(var) => stack.push(*var),
+                        roc_types::subs::TagExt::Any(_) => { /* ignore */ }
+                    }
+                }
+                FlatType::EmptyRecord => {}
+                FlatType::EmptyTuple => {}
+                FlatType::EmptyTagUnion => {}
+            },
+            Content::Alias(_, _, actual, _) => {
+                stack.push(*actual);
+            }
+            Content::LambdaSet(lambda_set) => {
+                // the lambda set itself should already be caught by Func above, but the
+                // capture can itself contain more lambda sets
+                for index in lambda_set.solved.variables() {
+                    let subs_slice = subs.variable_slices[index.index as usize];
+                    stack.extend(subs.variables[subs_slice.indices()].iter());
+                }
+            }
+        }
+    }
+
+    result
+}
+
 pub fn generate_glue_procs<'a, 'i, I>(
     home: ModuleId,
     ident_ids: &mut IdentIds,

crates/compiler/mono/src/ir/decision_tree.rs

@@ -1193,7 +1193,7 @@ fn extract<'a>(
 
 /// FIND IRRELEVANT BRANCHES
 
-fn is_irrelevant_to(selected_path: &[PathInstruction], branch: &Branch<'_>) -> bool {
+fn is_irrelevant_to(selected_path: &[PathInstruction], branch: &Branch) -> bool {
     match branch
         .patterns
         .iter()
@@ -1358,7 +1358,7 @@ fn small_branching_factor(branches: &[Branch], path: &[PathInstruction]) -> usiz
         branches.iter().any(|b| is_irrelevant_to(path, b))
     };
 
-    relevant_tests.len() + (if !fallbacks { 0 } else { 1 })
+    relevant_tests.len() + usize::from(fallbacks)
 }
 
 #[derive(Debug, PartialEq)]

crates/compiler/mono/src/layout.rs

@@ -634,7 +634,10 @@ impl<'a> RawFunctionLayout<'a> {
     /// Returns Err(()) if given an error, or Ok(Layout) if given a non-erroneous Structure.
     /// Panics if given a FlexVar or RigidVar, since those should have been
     /// monomorphized away already!
-    fn from_var(env: &mut Env<'a, '_>, var: Variable) -> Cacheable<RawFunctionLayoutResult<'a>> {
+    pub(crate) fn from_var(
+        env: &mut Env<'a, '_>,
+        var: Variable,
+    ) -> Cacheable<RawFunctionLayoutResult<'a>> {
         env.cached_raw_function_or(var, |env| {
             if env.is_seen(var) {
                 unreachable!("The initial variable of a signature cannot be seen already")
@@ -2175,9 +2178,9 @@ macro_rules! list_element_layout {
 
 pub struct Env<'a, 'b> {
     target_info: TargetInfo,
-    arena: &'a Bump,
+    pub(crate) arena: &'a Bump,
     seen: Vec<'a, Variable>,
-    subs: &'b Subs,
+    pub(crate) subs: &'b Subs,
     cache: &'b mut LayoutCache<'a>,
 }
 

crates/compiler/mono/src/layout/intern.rs

@@ -433,6 +433,22 @@ impl<'a> InLayout<'a> {
     pub fn index(&self) -> usize {
         self.0
     }
+
+    pub fn try_int_width(self) -> Option<IntWidth> {
+        match self {
+            Layout::U8 => Some(IntWidth::U8),
+            Layout::U16 => Some(IntWidth::U16),
+            Layout::U32 => Some(IntWidth::U32),
+            Layout::U64 => Some(IntWidth::U64),
+            Layout::U128 => Some(IntWidth::U128),
+            Layout::I8 => Some(IntWidth::I8),
+            Layout::I16 => Some(IntWidth::I16),
+            Layout::I32 => Some(IntWidth::I32),
+            Layout::I64 => Some(IntWidth::I64),
+            Layout::I128 => Some(IntWidth::I128),
+            _ => None,
+        }
+    }
 }
 
 /// A concurrent interner, suitable for usage between threads.