Move Environment from TypeChecker to Context

numbat/src/interpreter/mod.rs

@@ -209,8 +209,9 @@ mod tests {
             .transform(statements)
             .expect("No name resolution errors for inputs in this test suite");
         let mut typechecker = crate::typechecker::TypeChecker::default();
+
         let statements_typechecked = typechecker
-            .check(&statements_transformed)
+            .check(&mut crate::Environment::default(), &statements_transformed)
             .expect("No type check errors for inputs in this test suite");
         BytecodeInterpreter::new().interpret_statements(
             &mut InterpreterSettings::default(),

numbat/src/lib.rs

@@ -66,6 +66,7 @@ use resolver::CodeSource;
 use resolver::Resolver;
 use resolver::ResolverError;
 use thiserror::Error;
+pub use typechecker::Environment;
 use typechecker::{TypeCheckError, TypeChecker};
 
 pub use diagnostic::Diagnostic;
@@ -102,6 +103,7 @@ type Result<T> = std::result::Result<T, Box<NumbatError>>;
 pub struct Context {
     prefix_transformer: Transformer,
     typechecker: TypeChecker,
+    env: Environment,
     interpreter: BytecodeInterpreter,
     resolver: Resolver,
     load_currency_module_on_demand: bool,
@@ -113,6 +115,7 @@ impl Context {
         Context {
             prefix_transformer: Transformer::new(),
             typechecker: TypeChecker::default(),
+            env: Environment::default(),
             interpreter: BytecodeInterpreter::new(),
             resolver: Resolver::new(module_importer),
             load_currency_module_on_demand: false,
@@ -179,7 +182,7 @@ impl Context {
             .iter()
             .filter(|name| !name.starts_with('_'))
             .map(move |name| {
-                let (signature, meta) = self.typechecker.lookup_function(name).unwrap();
+                let (signature, meta) = self.env.lookup_function(name).unwrap();
                 (
                     name.clone(),
                     meta.name.clone(),
@@ -468,7 +471,7 @@ impl Context {
             return help;
         }
 
-        if let Some((fn_signature, fn_metadata)) = self.typechecker.lookup_function(keyword) {
+        if let Some((fn_signature, fn_metadata)) = self.env.lookup_function(keyword) {
             let metadata = fn_metadata.clone();
 
             let mut help = m::text("Function:    ");
@@ -592,10 +595,11 @@ impl Context {
         let transformed_statements = result?;
 
         let typechecker_old = self.typechecker.clone();
+        let env_old = self.env.clone();
 
         let result = self
             .typechecker
-            .check(&transformed_statements)
+            .check(&mut self.env, &transformed_statements)
             .map_err(|err| NumbatError::TypeCheckError(*err));
 
         if result.is_err() {
@@ -611,6 +615,7 @@ impl Context {
             //
             self.prefix_transformer = prefix_transformer_old.clone();
             self.typechecker = typechecker_old.clone();
+            self.env = env_old.clone();
 
             if self.load_currency_module_on_demand {
                 if let Err(NumbatError::TypeCheckError(TypeCheckError::UnknownIdentifier(
@@ -831,6 +836,7 @@ impl Context {
             //
             self.prefix_transformer = prefix_transformer_old;
             self.typechecker = typechecker_old;
+            self.env = env_old;
             self.interpreter = interpreter_old;
         }
 

numbat/src/typechecker/environment.rs

@@ -1,4 +1,4 @@
-use crate::ast::{TypeAnnotation, TypeParameterBound};
+use crate::ast::{self, TypeAnnotation, TypeParameterBound};
 use crate::dimension::DimensionRegistry;
 use crate::pretty_print::PrettyPrint;
 use crate::span::Span;
@@ -166,6 +166,22 @@ impl Environment {
         }
     }
 
+    pub(crate) fn get_proper_function_reference<'a>(
+        &self,
+        expr: &ast::Expression<'a>,
+    ) -> Option<(&'a str, &FunctionSignature)> {
+        match expr {
+            ast::Expression::Identifier(_, name) => self
+                .get_function_info(name)
+                .map(|(signature, _)| (*name, signature)),
+            _ => None,
+        }
+    }
+
+    pub fn lookup_function(&self, name: &str) -> Option<(&FunctionSignature, &FunctionMetadata)> {
+        self.get_function_info(name)
+    }
+
     pub(crate) fn generalize_types(&mut self, dtype_variables: &[TypeVariable]) {
         for (_, kind) in self.identifiers.iter_mut() {
             match kind {