better

lambdaclass · Feb 5, 2024 · 36ae964 · 36ae964
1 parent 9031247
commit 36ae964
Showing 1 changed file with 58 additions and 62 deletions.
diff --git a/crates/concrete_ir/src/lowering/prepass.rs b/crates/concrete_ir/src/lowering/prepass.rs
@@ -22,70 +22,66 @@ pub fn prepass_module(mut ctx: BuildCtx, mod_def: &ast::modules::Module) -> Buil
         },
     );
 
-    for ct in &mod_def.contents {
-        match ct {
-            ast::modules::ModuleDefItem::Constant(info) => {
-                let m = ctx
-                    .body
-                    .modules
-                    .get_mut(&module_id)
-                    .expect("module should exist");
-                let next_id = ctx.gen.next_defid();
-                m.symbols
-                    .constants
-                    .insert(info.decl.name.name.clone(), next_id);
-            }
-            ast::modules::ModuleDefItem::Function(info) => {
-                let m = ctx
-                    .body
-                    .modules
-                    .get_mut(&module_id)
-                    .expect("module should exist");
-                let next_id = ctx.gen.next_defid();
-                m.symbols
-                    .functions
-                    .insert(info.decl.name.name.clone(), next_id);
-            }
-            ast::modules::ModuleDefItem::Struct(info) => {
-                let m = ctx
-                    .body
-                    .modules
-                    .get_mut(&module_id)
-                    .expect("module should exist");
-                let next_id = ctx.gen.next_defid();
-                m.symbols.structs.insert(info.name.name.clone(), next_id);
-            }
-            ast::modules::ModuleDefItem::Type(info) => {
-                let m = ctx
-                    .body
-                    .modules
-                    .get_mut(&module_id)
-                    .expect("module should exist");
-                let next_id = ctx.gen.next_defid();
-                m.symbols.types.insert(info.name.name.clone(), next_id);
-            }
-            ast::modules::ModuleDefItem::Module(info) => {
-                let m = ctx
-                    .body
-                    .modules
-                    .get_mut(&module_id)
-                    .expect("module should exist");
-                let next_id = ctx.gen.next_defid();
-                m.symbols.modules.insert(info.name.name.clone(), next_id);
-                m.modules.insert(
-                    next_id,
-                    ModuleBody {
-                        id: next_id,
-                        parent_id: Some(module_id),
-                        symbols: Default::default(),
-                        functions: Default::default(),
-                        modules: Default::default(),
-                        function_signatures: Default::default(),
-                    },
-                );
-                ctx = prepass_sub_module(ctx, &[module_id], next_id, info);
+    {
+        let mut gen = ctx.gen;
+        let current_module = ctx
+            .body
+            .modules
+            .get_mut(&module_id)
+            .expect("module should exist");
+
+        for ct in &mod_def.contents {
+            match ct {
+                ast::modules::ModuleDefItem::Constant(info) => {
+                    let next_id = gen.next_defid();
+                    current_module
+                        .symbols
+                        .constants
+                        .insert(info.decl.name.name.clone(), next_id);
+                }
+                ast::modules::ModuleDefItem::Function(info) => {
+                    let next_id = gen.next_defid();
+                    current_module
+                        .symbols
+                        .functions
+                        .insert(info.decl.name.name.clone(), next_id);
+                }
+                ast::modules::ModuleDefItem::Struct(info) => {
+                    let next_id = gen.next_defid();
+                    current_module
+                        .symbols
+                        .structs
+                        .insert(info.name.name.clone(), next_id);
+                }
+                ast::modules::ModuleDefItem::Type(info) => {
+                    let next_id = gen.next_defid();
+                    current_module
+                        .symbols
+                        .types
+                        .insert(info.name.name.clone(), next_id);
+                }
+                ast::modules::ModuleDefItem::Module(_) => {}
             }
         }
+
+        ctx.gen = gen;
+    }
+
+    for ct in &mod_def.contents {
+        if let ast::modules::ModuleDefItem::Module(info) = ct {
+            let next_id = ctx.gen.next_defid();
+            let current_module = ctx
+                .body
+                .modules
+                .get_mut(&module_id)
+                .expect("module should exist");
+
+            current_module
+                .symbols
+                .modules
+                .insert(info.name.name.clone(), next_id);
+            ctx = prepass_sub_module(ctx, &[module_id], next_id, mod_def);
+        }
     }
 
     ctx