implement TypeCheckCtx

src/compile_pass.rs

@@ -1,11 +1,22 @@
+use std::collections::HashMap;
+
+use index_vec::IndexVec;
+
 use crate::{
-    stage::parse::ParseCtx,
+    repr::identifier::FunctionIdx,
+    stage::{
+        parse::ParseCtx,
+        type_check::{FunctionSignature, TypeCheckCtx},
+    },
     util::symbol_map::{interner_symbol_map::*, SymbolMap},
 };
 
 #[derive(Default)]
 pub struct CompilePass {
     symbols: InternerSymbolMap,
+
+    function_signatures: IndexVec<FunctionIdx, FunctionSignature>,
+    function_symbols: HashMap<Symbol, FunctionIdx>,
 }
 
 impl SymbolMap for CompilePass {
@@ -28,3 +39,19 @@ impl SymbolMap for CompilePass {
 }
 
 impl ParseCtx for CompilePass {}
+
+impl TypeCheckCtx for CompilePass {
+    fn register_function(&mut self, symbol: Symbol, signature: FunctionSignature) -> FunctionIdx {
+        let idx = self.function_signatures.push(signature);
+        self.function_symbols.insert(symbol, idx);
+        idx
+    }
+
+    fn get_function(&self, idx: FunctionIdx) -> FunctionSignature {
+        self.function_signatures[idx].clone()
+    }
+
+    fn lookup_function_symbol(&self, symbol: Symbol) -> Option<FunctionIdx> {
+        self.function_symbols.get(&symbol).cloned()
+    }
+}

src/main.rs

@@ -26,22 +26,19 @@ fn main() -> int {
         }
     };
 
-    let program = match program.ty_solve() {
+    let program = match program.ty_solve(&mut ctx) {
         Ok(program) => program,
         Err(e) => {
             eprintln!("{e}");
             return;
         }
     };
 
-    let ir_ctx = ir::lower(program);
-    let _main = ir_ctx
-        .function_for_name("main")
-        .expect("main function to exist");
+    let main = program.main.name;
 
+    let ir_ctx = ir::lower(program);
     let ctx = inkwell::context::Context::create();
 
-    let main = ir_ctx.symbol_map.get("main").unwrap();
     let function_ids = ir_ctx.functions.keys().cloned().collect::<Vec<_>>();
     let mut llvm_pass = Pass::new(&ctx, ir_ctx);
     function_ids.into_iter().for_each(|function| {

src/repr/ast/typed.rs

@@ -1,6 +1,6 @@
 use crate::{
-    generate_ast, repr::ty::Ty, stage::lower_ir::FunctionIdx,
-    util::symbol_map::interner_symbol_map::Symbol,
+    generate_ast,
+    repr::{identifier::FunctionIdx, ty::Ty},
 };
 
 #[derive(Clone, Debug)]
@@ -11,6 +11,5 @@ pub struct TyInfo {
 
 generate_ast! {
     TyInfo: TyInfo,
-    // TODO: Something else for this
-    FnIdentifier: Symbol
+    FnIdentifier: FunctionIdx
 }

src/repr/identifier/function.rs

@@ -0,0 +1,3 @@
+use index_vec::define_index_type;
+
+define_index_type! {pub struct FunctionIdx = usize;}

src/repr/identifier/mod.rs

@@ -0,0 +1,3 @@
+mod function;
+
+pub use function::FunctionIdx;

src/repr/ir/function.rs

@@ -3,8 +3,8 @@ use std::collections::HashSet;
 use index_vec::IndexVec;
 
 use crate::{
-    repr::ast::typed as ast,
-    stage::{lower_ir::FunctionIdx, type_check::FunctionSignature},
+    repr::{ast::typed as ast, identifier::FunctionIdx},
+    stage::type_check::FunctionSignature,
     util::symbol_map::interner_symbol_map::Symbol,
 };
 
@@ -16,8 +16,7 @@ index_vec::define_index_type! {
 
 #[derive(Debug, Clone)]
 pub struct Function {
-    // TODO: This should be something else
-    pub symbol: Symbol,
+    pub symbol: FunctionIdx,
     pub signature: FunctionSignature,
     pub basic_blocks: IndexVec<BasicBlockIdx, BasicBlock>,
     pub scope: HashSet<Symbol>,

src/repr/ir/triple/mod.rs

@@ -1,6 +1,6 @@
 use index_vec::define_index_type;
 
-use crate::{stage::lower_ir::FunctionIdx, util::symbol_map::interner_symbol_map::Symbol};
+use crate::{repr::identifier::FunctionIdx, util::symbol_map::interner_symbol_map::Symbol};
 
 use super::{BasicBlockIdx, Value};
 
@@ -25,7 +25,7 @@ pub enum Triple {
     /// Jump to the corresponding basic block.
     Jump(BasicBlockIdx),
     /// Call the corresponding function.
-    Call(Symbol),
+    Call(FunctionIdx),
     /// Return with the provided value.
     Return(Value),
     /// Assign some symbol to some value.

src/repr/mod.rs

@@ -1,4 +1,5 @@
 pub mod ast;
+pub mod identifier;
 pub mod ir;
 pub mod token;
 pub mod ty;

src/stage/codegen/llvm/mod.rs

@@ -12,8 +12,11 @@ use inkwell::{
 };
 
 use crate::{
-    repr::ir::{BinaryOp, Triple, UnaryOp, Value},
-    stage::lower_ir::{FunctionIdx, IRCtx},
+    repr::{
+        identifier::FunctionIdx,
+        ir::{BinaryOp, Triple, UnaryOp, Value},
+    },
+    stage::lower_ir::IRCtx,
     util::symbol_map::interner_symbol_map::Symbol,
 };
 
@@ -25,8 +28,8 @@ pub struct Pass<'ctx> {
     module: Module<'ctx>,
 
     symbols: HashMap<Symbol, PointerValue<'ctx>>,
-    basic_blocks: HashMap<(Symbol, BasicBlockIdx), inkwell::basic_block::BasicBlock<'ctx>>,
-    pub function_values: HashMap<Symbol, FunctionValue<'ctx>>,
+    basic_blocks: HashMap<(FunctionIdx, BasicBlockIdx), inkwell::basic_block::BasicBlock<'ctx>>,
+    pub function_values: HashMap<FunctionIdx, FunctionValue<'ctx>>,
 }
 
 impl<'ctx> Pass<'ctx> {
@@ -60,23 +63,23 @@ impl<'ctx> Pass<'ctx> {
     }
 
     /// Compile the provided function, returning the LLVM handle to it.
-    pub fn compile(&mut self, function_id: Symbol) -> FunctionValue<'ctx> {
+    pub fn compile(&mut self, function_idx: FunctionIdx) -> FunctionValue<'ctx> {
         let function = self
             .ir_ctx
             .functions
-            .get(&function_id)
+            .get(&function_idx)
             .expect("function to exist");
 
         let builder = self.llvm_ctx.create_builder();
 
         let fn_value = self
             .function_values
-            .get(&function_id)
+            .get(&function_idx)
             .expect("function to exist")
             .to_owned();
 
         // BUG: This won't work with multiple functions
-        let entry_bb = (function_id, BasicBlockIdx::new(0));
+        let entry_bb = (function_idx, BasicBlockIdx::new(0));
         let entry = *self.basic_blocks.entry(entry_bb).or_insert_with(|| {
             self.llvm_ctx
                 .append_basic_block(fn_value, &format!("bb_{:?}_{:?}", entry_bb.0, entry_bb.1))
@@ -108,7 +111,7 @@ impl<'ctx> Pass<'ctx> {
     fn compile_basic_block(
         &mut self,
         function: &FunctionValue<'ctx>,
-        basic_block_id: (Symbol, BasicBlockIdx),
+        basic_block_id: (FunctionIdx, BasicBlockIdx),
     ) {
         let bb = *self.basic_blocks.entry(basic_block_id).or_insert_with(|| {
             self.llvm_ctx

src/stage/lower_ir/ctx.rs

@@ -1,18 +1,14 @@
 use std::collections::HashMap;
 
-use index_vec::{define_index_type, IndexVec};
-
 use crate::{
-    repr::ir::Function,
-    util::symbol_map::interner_symbol_map::{InternerSymbolMap, Symbol},
+    repr::{identifier::FunctionIdx, ir::Function},
+    util::symbol_map::interner_symbol_map::InternerSymbolMap,
 };
 
-define_index_type! {pub struct FunctionIdx = usize;}
-
 #[derive(Default, Clone, Debug)]
 pub struct IRCtx {
     /// Map of function symbol to the basic block entry point
-    pub functions: HashMap<Symbol, Function>,
+    pub functions: HashMap<FunctionIdx, Function>,
     pub symbol_map: InternerSymbolMap,
 }
 
@@ -23,9 +19,4 @@ impl IRCtx {
             ..Default::default()
         }
     }
-
-    pub fn function_for_name(&self, s: &str) -> Option<Function> {
-        let symbol = self.symbol_map.get(s)?;
-        self.functions.get(&symbol).cloned()
-    }
 }

src/stage/lower_ir/lowering/mod.rs

@@ -1,9 +1,6 @@
-use crate::{
-    repr::{ast::typed as ast, ir::*},
-    util::symbol_map::interner_symbol_map::Symbol,
-};
+use crate::repr::{ast::typed as ast, identifier::FunctionIdx, ir::*};
 
-use super::{FunctionIdx, IRCtx};
+use super::IRCtx;
 
 pub fn lower(program: ast::Program) -> IRCtx {
     let mut ir = IRCtx::new(program.symbols);
@@ -39,7 +36,7 @@ impl FunctionLoweringCtx<'_> {
     }
 }
 
-fn lower_function(ir_ctx: &mut IRCtx, function: ast::Function) -> Symbol {
+fn lower_function(ir_ctx: &mut IRCtx, function: ast::Function) -> FunctionIdx {
     let mut repr_function = Function::new(&function);
 
     // Insert entry basic block

src/stage/type_check/ctx.rs

@@ -0,0 +1,142 @@
+use index_vec::{define_index_type, IndexVec};
+
+use crate::{
+    repr::{identifier::FunctionIdx, ty::Ty},
+    util::symbol_map::interner_symbol_map::Symbol,
+};
+
+use super::FunctionSignature;
+
+pub trait TypeCheckCtx {
+    /// Register a function's signature and associated symbol, to produce a unique identifier for the function.
+    fn register_function(&mut self, symbol: Symbol, signature: FunctionSignature) -> FunctionIdx;
+
+    /// Get the signature associated with a function identifier.
+    fn get_function(&self, idx: FunctionIdx) -> FunctionSignature;
+
+    /// Attempt to look up a symbol, returning the associated function's identifier if it exists.
+    fn lookup_function_symbol(&self, symbol: Symbol) -> Option<FunctionIdx>;
+}
+
+#[cfg(test)]
+mockall::mock! {
+    pub TypeCheckCtx {}
+
+    impl TypeCheckCtx for TypeCheckCtx {
+        fn register_function(&mut self, symbol: Symbol, signature: FunctionSignature) -> FunctionIdx;
+        fn get_function(&self, idx: FunctionIdx) -> FunctionSignature;
+        fn lookup_function_symbol(&self, symbol: Symbol) -> Option<FunctionIdx>;
+    }
+}
+
+define_index_type! {pub struct ScopeIdx = usize;}
+define_index_type! {pub struct BindingIdx = usize;}
+
+/// A binding within a specific scope.
+pub struct ScopedBinding(ScopeIdx, BindingIdx);
+
+pub struct ScopePart {
+    /// Indicates that this scope (and potentially a descendant) is active.
+    active: bool,
+
+    /// All of the bindings present within this scope.
+    bindings: IndexVec<BindingIdx, (Symbol, Ty)>,
+}
+
+impl ScopePart {
+    /// Create a new scope part, and automatically mark it as active.
+    pub fn new() -> Self {
+        Self {
+            active: true,
+            bindings: IndexVec::new(),
+        }
+    }
+
+    pub fn exit(&mut self) {
+        self.active = false;
+    }
+
+    pub fn add(&mut self, symbol: Symbol, ty: Ty) -> BindingIdx {
+        self.bindings.push((symbol, ty))
+    }
+}
+
+pub struct Scope {
+    scopes: IndexVec<ScopeIdx, ScopePart>,
+}
+
+impl Scope {
+    /// Create a new instance, and automatically enter the base scope.
+    pub fn new() -> Self {
+        let mut scope = Self {
+            scopes: IndexVec::new(),
+        };
+
+        // Automatically enter the first scope
+        scope.enter();
+
+        scope
+    }
+
+    /// Enter a new scope.
+    pub fn enter(&mut self) -> ScopeIdx {
+        self.scopes.push(ScopePart::new())
+    }
+
+    /// Exit the most recent scope.
+    pub fn leave(&mut self) {
+        let Some(scope) = self.scopes.last_mut() else {
+            return;
+        };
+
+        scope.exit();
+    }
+
+    /// Register a symbol and associated type, and produce a unique binding for it.
+    pub fn register(&mut self, symbol: Symbol, ty: Ty) -> ScopedBinding {
+        // Fetch the currently active scope
+        let active_scope_idx = self.active_scope();
+        let active_scope = &mut self.scopes[active_scope_idx];
+
+        // Register the binding type and symbol
+        let binding_idx = active_scope.add(symbol, ty);
+
+        ScopedBinding(active_scope_idx, binding_idx)
+    }
+
+    /// Attempt to retrieve a binding and type for the provided symbol. Will only search for items that are in scope.
+    pub fn resolve(&mut self, symbol: Symbol) -> Option<(ScopedBinding, Ty)> {
+        // Run through all possible scopes
+        self.scopes
+            .iter_enumerated()
+            // Only include active scopes
+            .filter(|(_, scope)| scope.active)
+            // Extract all of the bindings from the scopes
+            .flat_map(|(scope_idx, scope)| {
+                scope
+                    .bindings
+                    .iter_enumerated()
+                    // Only consider bindings that match the symbol
+                    .filter(|(_, (test_symbol, _))| *test_symbol == symbol)
+                    // Generate the scoped binding representation to track which scope the binding originated from
+                    .map(move |(binding_idx, (_, ty))| (ScopedBinding(scope_idx, binding_idx), *ty))
+            })
+            .next_back()
+    }
+
+    /// Find the currently activated scope identifier.
+    fn active_scope(&self) -> ScopeIdx {
+        self.scopes
+            .iter_enumerated()
+            .rev()
+            .find(|(_, scope)| scope.active)
+            .map(|(idx, _)| idx)
+            .expect("should always be at least one scope active")
+    }
+}
+
+impl Default for Scope {
+    fn default() -> Self {
+        Self::new()
+    }
+}