fix: Lazily elaborate functions in comptime interpreter

compiler/noirc_frontend/src/elaborator/comptime.rs

@@ -1,5 +1,3 @@
-use std::mem::replace;
-
 use crate::{
     hir_def::expr::HirIdent,
     node_interner::{DependencyId, FuncId},
@@ -11,42 +9,40 @@ impl<'context> Elaborator<'context> {
     /// Elaborate an expression from the middle of a comptime scope.
     /// When this happens we require additional information to know
     /// what variables should be in scope.
-    pub fn elaborate_item_from_comptime<T>(
-        &mut self,
+    pub fn elaborate_item_from_comptime<'a, T>(
+        &'a mut self,
         current_function: Option<FuncId>,
-        f: impl FnOnce(&mut Self) -> T,
+        f: impl FnOnce(&mut Elaborator<'a>) -> T,
     ) -> T {
-        self.function_context.push(FunctionContext::default());
-        let old_scope = self.scopes.end_function();
-        self.scopes.start_function();
-        let function_id = current_function.map(DependencyId::Function);
-        let old_item = replace(&mut self.current_item, function_id);
+        // Create a fresh elaborator to ensure no state is changed from
+        // this elaborator
+        let mut elaborator = Elaborator::new(
+            self.interner,
+            self.def_maps,
+            self.crate_id,
+            self.debug_comptime_in_file,
+        );
 
-        // Note: recover_generics isn't good enough here because any existing generics
-        // should not be in scope of this new function
-        let old_generics = std::mem::take(&mut self.generics);
+        elaborator.function_context.push(FunctionContext::default());
+        elaborator.scopes.start_function();
 
-        let old_crate_and_module = current_function.map(|function| {
-            let meta = self.interner.function_meta(&function);
-            let old_crate = replace(&mut self.crate_id, meta.source_crate);
-            let old_module = replace(&mut self.local_module, meta.source_module);
-            self.introduce_generics_into_scope(meta.all_generics.clone());
-            (old_crate, old_module)
-        });
+        if let Some(function) = current_function {
+            let meta = elaborator.interner.function_meta(&function);
+            elaborator.current_item = Some(DependencyId::Function(function));
+            elaborator.crate_id = meta.source_crate;
+            elaborator.local_module = meta.source_module;
+            elaborator.file = meta.source_file;
+            elaborator.introduce_generics_into_scope(meta.all_generics.clone());
+        }
 
-        self.populate_scope_from_comptime_scopes();
-        let result = f(self);
+        elaborator.comptime_scopes = std::mem::take(&mut self.comptime_scopes);
+        elaborator.populate_scope_from_comptime_scopes();
 
-        if let Some((old_crate, old_module)) = old_crate_and_module {
-            self.crate_id = old_crate;
-            self.local_module = old_module;
-        }
+        let result = f(&mut elaborator);
+        elaborator.check_and_pop_function_context();
 
-        self.generics = old_generics;
-        self.current_item = old_item;
-        self.scopes.end_function();
-        self.scopes.0.push(old_scope);
-        self.check_and_pop_function_context();
+        self.comptime_scopes = elaborator.comptime_scopes;
+        self.errors.append(&mut elaborator.errors);
         result
     }
 

compiler/noirc_frontend/src/elaborator/mod.rs

@@ -15,6 +15,7 @@ use crate::{
             },
             dc_mod,
         },
+        def_map::DefMaps,
         resolution::{errors::ResolverError, path_resolver::PathResolver},
         scope::ScopeForest as GenericScopeForest,
         type_check::TypeCheckError,
@@ -54,7 +55,7 @@ use crate::{
 use crate::{
     hir::{
         def_collector::dc_crate::{UnresolvedFunctions, UnresolvedTraitImpl},
-        def_map::{CrateDefMap, ModuleData},
+        def_map::ModuleData,
     },
     hir_def::traits::TraitImpl,
     macros_api::ItemVisibility,
@@ -102,7 +103,7 @@ pub struct Elaborator<'context> {
 
     pub(crate) interner: &'context mut NodeInterner,
 
-    def_maps: &'context mut BTreeMap<CrateId, CrateDefMap>,
+    def_maps: &'context mut DefMaps,
 
     file: FileId,
 
@@ -130,8 +131,6 @@ pub struct Elaborator<'context> {
     /// to the corresponding trait impl ID.
     current_trait_impl: Option<TraitImplId>,
 
-    trait_id: Option<TraitId>,
-
     /// In-resolution names
     ///
     /// This needs to be a set because we can have multiple in-resolution
@@ -195,22 +194,22 @@ struct FunctionContext {
 
 impl<'context> Elaborator<'context> {
     pub fn new(
-        context: &'context mut Context,
+        interner: &'context mut NodeInterner,
+        def_maps: &'context mut DefMaps,
         crate_id: CrateId,
         debug_comptime_in_file: Option<FileId>,
     ) -> Self {
         Self {
             scopes: ScopeForest::default(),
             errors: Vec::new(),
-            interner: &mut context.def_interner,
-            def_maps: &mut context.def_maps,
+            interner,
+            def_maps,
             file: FileId::dummy(),
             nested_loops: 0,
             generics: Vec::new(),
             lambda_stack: Vec::new(),
             self_type: None,
             current_item: None,
-            trait_id: None,
             local_module: LocalModuleId::dummy_id(),
             crate_id,
             resolving_ids: BTreeSet::new(),
@@ -223,6 +222,19 @@ impl<'context> Elaborator<'context> {
         }
     }
 
+    pub fn from_context(
+        context: &'context mut Context,
+        crate_id: CrateId,
+        debug_comptime_in_file: Option<FileId>,
+    ) -> Self {
+        Self::new(
+            &mut context.def_interner,
+            &mut context.def_maps,
+            crate_id,
+            debug_comptime_in_file,
+        )
+    }
+
     pub fn elaborate(
         context: &'context mut Context,
         crate_id: CrateId,
@@ -238,7 +250,7 @@ impl<'context> Elaborator<'context> {
         items: CollectedItems,
         debug_comptime_in_file: Option<FileId>,
     ) -> Self {
-        let mut this = Self::new(context, crate_id, debug_comptime_in_file);
+        let mut this = Self::from_context(context, crate_id, debug_comptime_in_file);
 
         // Filter out comptime items to execute their functions first if needed.
         // This step is why comptime items can only refer to other comptime items
@@ -337,17 +349,12 @@ impl<'context> Elaborator<'context> {
     }
 
     fn elaborate_functions(&mut self, functions: UnresolvedFunctions) {
-        self.file = functions.file_id;
-        self.trait_id = functions.trait_id; // TODO: Resolve?
-        self.self_type = functions.self_type;
-
-        for (local_module, id, _) in functions.functions {
-            self.local_module = local_module;
-            self.recover_generics(|this| this.elaborate_function(id));
+        for (_, id, _) in functions.functions {
+            self.elaborate_function(id);
         }
 
+        self.generics.clear();
         self.self_type = None;
-        self.trait_id = None;
     }
 
     fn introduce_generics_into_scope(&mut self, all_generics: Vec<ResolvedGeneric>) {
@@ -365,7 +372,7 @@ impl<'context> Elaborator<'context> {
         self.generics = all_generics;
     }
 
-    fn elaborate_function(&mut self, id: FuncId) {
+    pub(crate) fn elaborate_function(&mut self, id: FuncId) {
         let func_meta = self.interner.func_meta.get_mut(&id);
         let func_meta =
             func_meta.expect("FuncMetas should be declared before a function is elaborated");
@@ -378,11 +385,21 @@ impl<'context> Elaborator<'context> {
             FunctionBody::Resolving => return,
         };
 
+        let func_meta = func_meta.clone();
+
+        assert_eq!(
+            self.crate_id, func_meta.source_crate,
+            "Functions in other crates should be already elaborated"
+        );
+
+        self.local_module = func_meta.source_module;
+        self.file = func_meta.source_file;
+        self.self_type = func_meta.self_type.clone();
+        self.current_trait_impl = func_meta.trait_impl;
+
         self.scopes.start_function();
         let old_item = std::mem::replace(&mut self.current_item, Some(DependencyId::Function(id)));
 
-        let func_meta = func_meta.clone();
-
         self.trait_bounds = func_meta.trait_constraints.clone();
         self.function_context.push(FunctionContext::default());
 
@@ -775,6 +792,8 @@ impl<'context> Elaborator<'context> {
             source_crate: self.crate_id,
             source_module: self.local_module,
             function_body: FunctionBody::Unresolved(func.kind, body, func.def.span),
+            self_type: self.self_type.clone(),
+            source_file: self.file,
         };
 
         self.interner.push_fn_meta(meta, func_id);
@@ -1003,10 +1022,7 @@ impl<'context> Elaborator<'context> {
         self.self_type = None;
     }
 
-    fn get_module_mut(
-        def_maps: &mut BTreeMap<CrateId, CrateDefMap>,
-        module: ModuleId,
-    ) -> &mut ModuleData {
+    fn get_module_mut(def_maps: &mut DefMaps, module: ModuleId) -> &mut ModuleData {
         let message = "A crate should always be present for a given crate id";
         &mut def_maps.get_mut(&module.krate).expect(message).modules[module.local_id.0]
     }

compiler/noirc_frontend/src/elaborator/types.rs

@@ -412,7 +412,8 @@ impl<'context> Elaborator<'context> {
         &mut self,
         path: &Path,
     ) -> Option<(TraitMethodId, TraitConstraint, bool)> {
-        let trait_id = self.trait_id?;
+        let trait_impl = self.current_trait_impl?;
+        let trait_id = self.interner.try_get_trait_implementation(trait_impl)?.borrow().trait_id;
 
         if path.kind == PathKind::Plain && path.segments.len() == 2 {
             let name = &path.segments[0].0.contents;

compiler/noirc_frontend/src/hir/comptime/interpreter.rs

@@ -11,6 +11,7 @@
 use crate::elaborator::Elaborator;
 use crate::graph::CrateId;
 use crate::hir_def::expr::ImplKind;
+use crate::hir_def::function::FunctionBody;
 use crate::macros_api::UnaryOp;
 use crate::monomorphization::{
     perform_impl_bindings, perform_instantiation_bindings, resolve_trait_method,
@@ -135,18 +136,35 @@
             self.define_pattern(parameter, typ, argument, arg_location)?;
         }
 
-        let function_body =
-            self.elaborator.interner.function(&function).try_as_expr().ok_or_else(|| {
-                let function = self.elaborator.interner.function_name(&function).to_owned();
-                InterpreterError::NonComptimeFnCallInSameCrate { function, location }
-            })?;
-
+        let function_body = self.get_function_body(function, location)?;
         let result = self.evaluate(function_body)?;
-
         self.exit_function(previous_state);
         Ok(result)
     }
 
+    /// Try to retrieve a function's body.
+    /// If the function has not yet been resolved this will attempt to lazily resolve it.
+    /// Afterwards, if the function's body is still not known or the function is still
+    /// in a Resolving state we issue an error.
+    fn get_function_body(&mut self, function: FuncId, location: Location) -> IResult<ExprId> {
+        let meta = self.elaborator.interner.function_meta(&function);
+        match self.elaborator.interner.function(&function).try_as_expr() {
+            Some(body) => Ok(body),
+            None => {
+                if matches!(&meta.function_body, FunctionBody::Unresolved(..)) {
+                    self.elaborator.elaborate_item_from_comptime(None, |elaborator| {
+                        elaborator.elaborate_function(function);
+                    });
+
+                    self.get_function_body(function, location)
+                } else {
+                    let function = self.elaborator.interner.function_name(&function).to_owned();
+                    Err(InterpreterError::NonComptimeFnCallInSameCrate { function, location })
+                }
+            }
+        }
+    }
+
     fn call_special(
         &mut self,
         function: FuncId,
@@ -173,7 +191,7 @@
             }
         } else {
             let name = self.elaborator.interner.function_name(&function);
             unreachable!("Non-builtin, lowlevel or oracle builtin fn '{name}'")
         }
     }
 
@@ -1167,8 +1185,8 @@
                     let expr = result.into_expression(self.elaborator.interner, location)?;
                     let expr = self
                         .elaborator
                         .elaborate_item_from_comptime(self.current_function, |elab| {
                             elab.elaborate_expression(expr).0
                         });
                     result = self.evaluate(expr)?;
                 }

compiler/noirc_frontend/src/hir/def_map/mod.rs

@@ -48,11 +48,13 @@ impl ModuleId {
 }
 
 impl ModuleId {
-    pub fn module(self, def_maps: &BTreeMap<CrateId, CrateDefMap>) -> &ModuleData {
+    pub fn module(self, def_maps: &DefMaps) -> &ModuleData {
         &def_maps[&self.krate].modules()[self.local_id.0]
     }
 }
 
+pub type DefMaps = BTreeMap<CrateId, CrateDefMap>;
+
 /// Map of all modules and scopes defined within a crate.
 ///
 /// The definitions of the crate are accessible indirectly via the scopes of each module.

compiler/noirc_frontend/src/hir_def/function.rs

@@ -1,3 +1,4 @@
+use fm::FileId;
 use iter_extended::vecmap;
 use noirc_errors::{Location, Span};
 
@@ -156,6 +157,13 @@ pub struct FuncMeta {
 
     /// The module this function was defined in
     pub source_module: LocalModuleId,
+
+    /// THe file this function was defined in
+    pub source_file: FileId,
+
+    /// If this function is from an impl (trait or regular impl), this
+    /// is the object type of the impl. Otherwise this is None.
+    pub self_type: Option<Type>,
 }
 
 #[derive(Debug, Clone)]