Code gen toplevel module functions

cli/src/main.rs

@@ -16,7 +16,7 @@ use roc_gen::llvm::build::{
 use roc_gen::llvm::convert::basic_type_from_layout;
 use roc_load::file::{LoadedModule, LoadingProblem};
 use roc_module::symbol::Symbol;
-use roc_mono::expr::{Expr, PartialProc, Procs};
+use roc_mono::expr::{Env, Expr, PartialProc, Procs};
 use roc_mono::layout::Layout;
 use std::time::SystemTime;
 
@@ -354,6 +354,18 @@ fn gen(
         ptr_bytes,
     };
     let mut procs = Procs::default();
+    let mut mono_problems = std::vec::Vec::new();
+    let mut ident_ids = env.interns.all_ident_ids.remove(&home).unwrap();
+    let mut mono_env = Env {
+        arena,
+        subs: &mut subs,
+        problems: &mut mono_problems,
+        home,
+        ident_ids: &mut ident_ids,
+        pointer_size: ptr_bytes,
+        symbol_counter: 0,
+        jump_counter: arena.alloc(0),
+    };
 
     // TODO remove this
     decls_by_id.insert(loaded.module_id, home_decls);
@@ -368,20 +380,31 @@ fn gen(
             match decl {
                 Declare(def) => match def.loc_pattern.value {
                     Identifier(symbol) => {
-                        let proc = match def.loc_expr.value {
+                        match def.loc_expr.value {
-                            Closure(_, _, _, _, _) => {
+                            Closure(annotation, _, _, loc_args, boxed_body) => {
-                                todo!("TODO register PartialProc for Closure");
+                                let (loc_body, ret_var) = *boxed_body;
-                            }
-                            body => PartialProc {
-                                annotation: def.expr_var,
-                                // This is a 0-arity thunk, so it has no arguments.
-                                patterns: bumpalo::collections::Vec::new_in(arena),
-                                body,
-                            },
-                        };
 
-                        procs.user_defined.insert(symbol, proc);
+                                procs.insert_closure(
-                        procs.module_decls.insert(symbol);
+                                    &mut mono_env,
+                                    Some(symbol),
+                                    annotation,
+                                    loc_args,
+                                    loc_body,
+                                    ret_var,
+                                );
+                            }
+                            body => {
+                                let proc = PartialProc {
+                                    annotation: def.expr_var,
+                                    // This is a 0-arity thunk, so it has no arguments.
+                                    patterns: bumpalo::collections::Vec::new_in(arena),
+                                    body,
+                                };
+
+                                procs.user_defined.insert(symbol, proc);
+                                procs.module_thunks.insert(symbol);
+                            }
+                        };
                     }
                     other => {
                         todo!("TODO gracefully handle Declare({:?})", other);
@@ -397,20 +420,8 @@ fn gen(
         }
     }
 
-    let mut ident_ids = env.interns.all_ident_ids.remove(&home).unwrap();
-
     // Populate Procs further and get the low-level Expr from the canonical Expr
-    let mut mono_problems = std::vec::Vec::new();
+    let main_body = Expr::new(&mut mono_env, loc_expr.value, &mut procs);
-    let main_body = Expr::new(
-        &arena,
-        &mut subs,
-        &mut mono_problems,
-        loc_expr.value,
-        &mut procs,
-        home,
-        &mut ident_ids,
-        ptr_bytes,
-    );
 
     // Put this module's ident_ids back in the interns, so we can use them in env.
     env.interns.all_ident_ids.insert(home, ident_ids);

cli/src/repl.rs

@@ -246,16 +246,17 @@ pub fn gen(src: &str, target: Triple, opt_level: OptLevel) -> Result<(String, St
 
     // Populate Procs and get the low-level Expr from the canonical Expr
     let mut mono_problems = Vec::new();
-    let main_body = roc_mono::expr::Expr::new(
+    let mut mono_env = roc_mono::expr::Env {
-        &arena,
+        arena: &arena,
-        &mut subs,
+        subs: &mut subs,
-        &mut mono_problems,
+        problems: &mut mono_problems,
-        loc_expr.value,
-        &mut procs,
         home,
-        &mut ident_ids,
+        ident_ids: &mut ident_ids,
-        ptr_bytes,
+        pointer_size: ptr_bytes,
-    );
+        symbol_counter: 0,
+        jump_counter: arena.alloc(0),
+    };
+    let main_body = roc_mono::expr::Expr::new(&mut mono_env, loc_expr.value, &mut procs);
 
     // Put this module's ident_ids back in the interns, so we can use them in Env.
     env.interns.all_ident_ids.insert(home, ident_ids);

compiler/gen/tests/helpers/eval.rs

@@ -63,7 +63,17 @@ macro_rules! assert_llvm_evals_to {
 
         // Populate Procs and get the low-level Expr from the canonical Expr
         let mut mono_problems = Vec::new();
-        let main_body = Expr::new(&arena, &mut subs, &mut mono_problems, loc_expr.value, &mut procs, home, &mut ident_ids, ptr_bytes);
+        let mut mono_env = roc_mono::expr::Env {
+            arena: &arena,
+            subs: &mut subs,
+            problems: &mut mono_problems,
+            home,
+            ident_ids: &mut ident_ids,
+            pointer_size: ptr_bytes,
+            symbol_counter: 0,
+            jump_counter: arena.alloc(0),
+        };
+        let main_body = Expr::new(&mut mono_env, loc_expr.value, &mut procs);
 
         // Put this module's ident_ids back in the interns, so we can use them in Env.
         env.interns.all_ident_ids.insert(home, ident_ids);
@@ -215,7 +225,17 @@ macro_rules! assert_opt_evals_to {
 
         // Populate Procs and get the low-level Expr from the canonical Expr
         let mut mono_problems = Vec::new();
-        let main_body = Expr::new(&arena, &mut subs, &mut mono_problems, loc_expr.value, &mut procs, home, &mut ident_ids, ptr_bytes);
+        let mut mono_env = roc_mono::expr::Env {
+            arena: &arena,
+            subs: &mut subs,
+            problems: &mut mono_problems,
+            home,
+            ident_ids: &mut ident_ids,
+            pointer_size: ptr_bytes,
+            symbol_counter: 0,
+            jump_counter: arena.alloc(0),
+        };
+        let main_body = Expr::new(&mut mono_env, loc_expr.value, &mut procs);
 
         // Put this module's ident_ids back in the interns, so we can use them in Env.
         env.interns.all_ident_ids.insert(home, ident_ids);
@@ -361,7 +381,17 @@ macro_rules! emit_expr {
         let mut ident_ids = env.interns.all_ident_ids.remove(&home).unwrap();
 
         // Populate Procs and get the low-level Expr from the canonical Expr
-        let main_body = Expr::new(&arena, &mut subs, loc_expr.value, &mut procs, home, &mut ident_ids, $crate::helpers::eval::POINTER_SIZE);
+        let mut mono_env = roc_mono::expr::Env {
+            arena: &arena,
+            subs: &mut subs,
+            problems: &mut mono_problems,
+            home,
+            ident_ids: &mut ident_ids,
+            pointer_size: ptr_bytes,
+            symbol_counter: 0,
+            jump_counter: arena.alloc(0),
+        };
+        let main_body = Expr::new(&mut mono_env, loc_expr.value, &mut procs);
 
         // Put this module's ident_ids back in the interns, so we can use them in Env.
         env.interns.all_ident_ids.insert(home, ident_ids);

compiler/mono/src/expr.rs

@@ -28,7 +28,7 @@ pub struct Proc<'a> {
 #[derive(Clone, Debug, PartialEq, Default)]
 pub struct Procs<'a> {
     pub user_defined: MutMap<Symbol, PartialProc<'a>>,
-    pub module_decls: MutSet<Symbol>,
+    pub module_thunks: MutSet<Symbol>,
     anonymous: MutMap<Symbol, Option<Proc<'a>>>,
     specializations: MutMap<ContentHash, (Symbol, Option<Proc<'a>>)>,
     builtin: MutSet<Symbol>,
@@ -43,6 +43,66 @@ impl<'a> Procs<'a> {
         self.anonymous.insert(symbol, proc);
     }
 
+    pub fn insert_closure(
+        &mut self,
+        env: &mut Env<'a, '_>,
+        name: Option<Symbol>,
+        annotation: Variable,
+        loc_args: std::vec::Vec<(Variable, Located<roc_can::pattern::Pattern>)>,
+        loc_body: Located<roc_can::expr::Expr>,
+        ret_var: Variable,
+    ) -> Symbol {
+        // turn record/tag patterns into a when expression, e.g.
+        //
+        // foo = \{ x } -> body
+        //
+        // becomes
+        //
+        // foo = \r -> when r is { x } -> body
+        //
+        // conversion of one-pattern when expressions will do the most optimal thing
+
+        let (arg_vars, arg_symbols, body) = patterns_to_when(env, loc_args, ret_var, loc_body);
+
+        match name {
+            Some(symbol) => {
+                // a named closure
+                self.insert_user_defined(
+                    symbol,
+                    PartialProc {
+                        annotation,
+                        patterns: arg_symbols,
+                        body: body.value,
+                    },
+                );
+
+                symbol
+            }
+            None => {
+                // an anonymous closure. These will always be specialized already
+                // by the surrounding context
+                let symbol = env.fresh_symbol();
+
+                let opt_proc = specialize_proc_body(
+                    env,
+                    self,
+                    annotation,
+                    ret_var,
+                    symbol,
+                    &arg_vars,
+                    &arg_symbols,
+                    annotation,
+                    body.value,
+                )
+                .ok();
+
+                self.insert_anonymous(symbol, opt_proc);
+
+                symbol
+            }
+        }
+    }
+
     fn insert_specialization(
         &mut self,
         hash: ContentHash,
@@ -97,8 +157,8 @@ pub struct Env<'a, 'i> {
     pub home: ModuleId,
     pub ident_ids: &'i mut IdentIds,
     pub pointer_size: u32,
-    symbol_counter: usize,
     pub jump_counter: &'a mut u64,
+    pub symbol_counter: usize,
 }
 
 impl<'a, 'i> Env<'a, 'i> {
@@ -204,27 +264,11 @@ pub enum MonoProblem {
 #[allow(clippy::too_many_arguments)]
 impl<'a> Expr<'a> {
     pub fn new(
-        arena: &'a Bump,
+        env: &mut Env<'a, '_>,
-        subs: &'a mut Subs,
-        problems: &mut std::vec::Vec<MonoProblem>,
         can_expr: roc_can::expr::Expr,
         procs: &mut Procs<'a>,
-        home: ModuleId,
-        ident_ids: &mut IdentIds,
-        pointer_size: u32,
     ) -> Self {
-        let mut env = Env {
+        from_can(env, can_expr, procs, None)
-            arena,
-            subs,
-            problems,
-            home,
-            ident_ids,
-            pointer_size,
-            symbol_counter: 0,
-            jump_counter: arena.alloc(0),
-        };
-
-        from_can(&mut env, can_expr, procs, None)
     }
 }
 
@@ -428,7 +472,7 @@ fn from_can<'a>(
         Float(_, num) => Expr::Float(num),
         Str(string) | BlockStr(string) => Expr::Str(env.arena.alloc(string)),
         Var(symbol) => {
-            if procs.module_decls.contains(&symbol) {
+            if procs.module_thunks.contains(&symbol) {
                 let partial_proc = procs.get_user_defined(symbol).unwrap();
                 let fn_var = partial_proc.annotation;
                 let ret_var = partial_proc.annotation;
@@ -444,55 +488,7 @@ fn from_can<'a>(
 
         Closure(annotation, _, _, loc_args, boxed_body) => {
             let (loc_body, ret_var) = *boxed_body;
-
+            let symbol = procs.insert_closure(env, name, annotation, loc_args, loc_body, ret_var);
-            // turn record/tag patterns into a when expression, e.g.
-            //
-            // foo = \{ x } -> body
-            //
-            // becomes
-            //
-            // foo = \r -> when r is { x } -> body
-            //
-            // conversion of one-pattern when expressions will do the most optimal thing
-
-            let (arg_vars, arg_symbols, body) = patterns_to_when(env, loc_args, ret_var, loc_body);
-
-            let symbol = match name {
-                Some(symbol) => {
-                    // a named closure
-                    procs.insert_user_defined(
-                        symbol,
-                        PartialProc {
-                            annotation,
-                            patterns: arg_symbols,
-                            body: body.value,
-                        },
-                    );
-                    symbol
-                }
-                None => {
-                    // an anonymous closure. These will always be specialized already
-                    // by the surrounding context
-                    let symbol = env.fresh_symbol();
-
-                    let opt_proc = specialize_proc_body(
-                        env,
-                        procs,
-                        annotation,
-                        ret_var,
-                        symbol,
-                        &arg_vars,
-                        &arg_symbols,
-                        annotation,
-                        body.value,
-                    )
-                    .ok();
-
-                    procs.insert_anonymous(symbol, opt_proc);
-
-                    symbol
-                }
-            };
 
             Expr::FunctionPointer(symbol)
         }

compiler/mono/tests/test_mono.rs

@@ -53,16 +53,17 @@ mod test_mono {
 
         // Populate Procs and Subs, and get the low-level Expr from the canonical Expr
         let mut mono_problems = Vec::new();
-        let mono_expr = Expr::new(
+        let mut mono_env = roc_mono::expr::Env {
-            &arena,
+            arena: &arena,
-            &mut subs,
+            subs: &mut subs,
-            &mut mono_problems,
+            problems: &mut mono_problems,
-            loc_expr.value,
-            &mut procs,
             home,
-            &mut ident_ids,
+            ident_ids: &mut ident_ids,
             pointer_size,
-        );
+            symbol_counter: 0,
+            jump_counter: arena.alloc(0),
+        };
+        let mono_expr = Expr::new(&mut mono_env, loc_expr.value, &mut procs);
 
         // Put this module's ident_ids back in the interns
         interns.all_ident_ids.insert(home, ident_ids);

compiler/mono/tests/test_opt.rs

@@ -42,16 +42,17 @@ mod test_opt {
 
         // Populate Procs and Subs, and get the low-level Expr from the canonical Expr
         let mut mono_problems = Vec::new();
-        let mono_expr = Expr::new(
+        let mut mono_env = roc_mono::expr::Env {
-            &arena,
+            arena: &arena,
-            &mut subs,
+            subs: &mut subs,
-            &mut mono_problems,
+            problems: &mut mono_problems,
-            loc_expr.value,
-            &mut procs,
             home,
-            &mut ident_ids,
+            ident_ids: &mut ident_ids,
             pointer_size,
-        );
+            symbol_counter: 0,
+            jump_counter: arena.alloc(0),
+        };
+        let mono_expr = Expr::new(&mut mono_env, loc_expr.value, &mut procs);
 
         let unexpected_calls = extract_named_calls(&mono_expr, &mut calls);
         let expected = CallProblems::default();
@@ -193,16 +194,17 @@ mod test_opt {
 
         // Populate Procs and Subs, and get the low-level Expr from the canonical Expr
         let mut mono_problems = Vec::new();
-        let mono_expr = Expr::new(
+        let mut mono_env = roc_mono::expr::Env {
-            &arena,
+            arena: &arena,
-            &mut subs,
+            subs: &mut subs,
-            &mut mono_problems,
+            problems: &mut mono_problems,
-            loc_expr.value,
-            &mut procs,
             home,
-            &mut ident_ids,
+            ident_ids: &mut ident_ids,
             pointer_size,
-        );
+            symbol_counter: 0,
+            jump_counter: arena.alloc(0),
+        };
+        let mono_expr = Expr::new(&mut mono_env, loc_expr.value, &mut procs);
 
         assert_eq!(mono_expr, expected);
     }

compiler/reporting/tests/test_reporting.rs

@@ -90,16 +90,17 @@ mod test_reporting {
             let pointer_size = std::mem::size_of::<u64>() as u32;
 
             // Populate Procs and Subs, and get the low-level Expr from the canonical Expr
-            let _mono_expr = Expr::new(
+            let mut mono_env = roc_mono::expr::Env {
-                &arena,
+                arena: &arena,
-                &mut subs,
+                subs: &mut subs,
-                &mut mono_problems,
+                problems: &mut mono_problems,
-                loc_expr.value,
-                &mut procs,
                 home,
-                &mut ident_ids,
+                ident_ids: &mut ident_ids,
                 pointer_size,
-            );
+                symbol_counter: 0,
+                jump_counter: arena.alloc(0),
+            };
+            let _mono_expr = Expr::new(&mut mono_env, loc_expr.value, &mut procs);
         }
 
         Ok((unify_problems, can_problems, mono_problems, home, interns))