create and implement IRCtx

src/compile_pass.rs

@@ -3,8 +3,9 @@ use std::collections::HashMap;
 use index_vec::IndexVec;
 
 use crate::{
-    repr::identifier::FunctionIdx,
+    repr::{ast::typed::Function, identifier::FunctionIdx, ir, ty::Ty},
     stage::{
+        lower_ir::{FunctionBuilder as FunctionBuilderTrait, IRCtx},
         parse::ParseCtx,
         type_check::{FunctionSignature, TypeCheckCtx},
     },
@@ -17,6 +18,8 @@ pub struct CompilePass {
 
     function_signatures: IndexVec<FunctionIdx, FunctionSignature>,
     function_symbols: HashMap<Symbol, FunctionIdx>,
+
+    ir_functions: HashMap<FunctionIdx, ir::Function>,
 }
 
 impl SymbolMap for CompilePass {
@@ -55,3 +58,88 @@ impl TypeCheckCtx for CompilePass {
         self.function_symbols.get(&symbol).cloned()
     }
 }
+
+impl IRCtx for CompilePass {
+    type FunctionBuilder = FunctionBuilder;
+
+    fn register_function(&mut self, idx: FunctionIdx, function: ir::Function) {
+        self.ir_functions.insert(idx, function);
+    }
+
+    fn all_functions(&self) -> Vec<(FunctionIdx, ir::Function)> {
+        self.ir_functions
+            .iter()
+            .map(|(idx, function)| (*idx, function.clone()))
+            .collect()
+    }
+}
+
+pub struct FunctionBuilder {
+    idx: FunctionIdx,
+    signature: FunctionSignature,
+
+    basic_blocks: IndexVec<ir::BasicBlockIdx, ir::BasicBlock>,
+    current_basic_block: ir::BasicBlockIdx,
+
+    scope: Vec<(Symbol, Ty)>,
+}
+
+impl FunctionBuilderTrait for FunctionBuilder {
+    fn new(function: &Function) -> Self {
+        let mut basic_blocks = IndexVec::new();
+        let current_basic_block = basic_blocks.push(ir::BasicBlock::default());
+
+        Self {
+            idx: function.name,
+            signature: FunctionSignature::from(function),
+            basic_blocks,
+            current_basic_block,
+            scope: Vec::new(),
+        }
+    }
+
+    fn register_scoped(&mut self, symbol: Symbol, ty: Ty) {
+        self.scope.push((symbol, ty));
+    }
+
+    fn add_triple(&mut self, triple: ir::Triple) -> ir::TripleRef {
+        ir::TripleRef {
+            basic_block: self.current_basic_block,
+            triple: self.basic_blocks[self.current_basic_block]
+                .triples
+                .push(triple),
+        }
+    }
+
+    fn current_bb(&self) -> ir::BasicBlockIdx {
+        self.current_basic_block
+    }
+
+    fn goto_bb(&mut self, bb: ir::BasicBlockIdx) {
+        assert!(
+            bb < self.basic_blocks.len_idx(),
+            "can only goto basic block if it exists"
+        );
+        self.current_basic_block = bb;
+    }
+
+    fn push_bb(&mut self) -> ir::BasicBlockIdx {
+        let idx = self.basic_blocks.push(ir::BasicBlock::default());
+
+        self.current_basic_block = idx;
+
+        idx
+    }
+
+    fn build<Ctx: IRCtx<FunctionBuilder = Self>>(self, ctx: &mut Ctx) {
+        ctx.register_function(
+            self.idx,
+            ir::Function {
+                symbol: self.idx,
+                signature: self.signature,
+                basic_blocks: self.basic_blocks,
+                scope: self.scope.into_iter().map(|(symbol, _)| symbol).collect(),
+            },
+        )
+    }
+}

src/main.rs

@@ -1,6 +1,11 @@
 use lumina::{
     compile_pass::CompilePass,
-    stage::{codegen::llvm::Pass, lex::Lexer, lower_ir as ir, parse::parse},
+    stage::{
+        codegen::llvm::Pass,
+        lex::Lexer,
+        lower_ir::{self as ir, IRCtx},
+        parse::parse,
+    },
     util::source::Source,
 };
 
@@ -36,11 +41,15 @@ fn main() -> int {
 
     let main = program.main.name;
 
-    let ir_ctx = ir::lower(program);
-    let ctx = inkwell::context::Context::create();
+    ir::lower(&mut ctx, program);
+    let llvm_ctx = inkwell::context::Context::create();
 
-    let function_ids = ir_ctx.functions.keys().cloned().collect::<Vec<_>>();
-    let mut llvm_pass = Pass::new(&ctx, ir_ctx);
+    let function_ids = ctx
+        .all_functions()
+        .iter()
+        .map(|(idx, _)| *idx)
+        .collect::<Vec<_>>();
+    let mut llvm_pass = Pass::new(&llvm_ctx, ctx);
     function_ids.into_iter().for_each(|function| {
         llvm_pass.compile(function);
     });

src/stage/codegen/llvm/mod.rs

@@ -14,45 +14,48 @@ use inkwell::{
 use crate::{
     repr::{
         identifier::FunctionIdx,
-        ir::{BinaryOp, Triple, UnaryOp, Value},
+        ir::{BasicBlockIdx, BinaryOp, ConstantValue, Triple, TripleIdx, UnaryOp, Value},
     },
     stage::lower_ir::IRCtx,
     util::symbol_map::interner_symbol_map::Symbol,
 };
 
-use crate::repr::ir::{BasicBlockIdx, ConstantValue, TripleIdx};
-
-pub struct Pass<'ctx> {
+pub struct Pass<'ctx, I> {
     llvm_ctx: &'ctx LLVMContext,
-    ir_ctx: IRCtx,
+    ir_ctx: I,
     module: Module<'ctx>,
 
     symbols: HashMap<Symbol, PointerValue<'ctx>>,
     basic_blocks: HashMap<(FunctionIdx, BasicBlockIdx), inkwell::basic_block::BasicBlock<'ctx>>,
     pub function_values: HashMap<FunctionIdx, FunctionValue<'ctx>>,
 }
 
-impl<'ctx> Pass<'ctx> {
-    pub fn new(llvm_ctx: &'ctx LLVMContext, ir_ctx: IRCtx) -> Self {
+impl<'ctx, I> Pass<'ctx, I>
+where
+    I: IRCtx,
+{
+    pub fn new(llvm_ctx: &'ctx LLVMContext, ir_ctx: I) -> Self {
         let module = llvm_ctx.create_module("module");
 
         Self {
-            function_values: HashMap::from_iter(ir_ctx.functions.iter().map(|(idx, function)| {
-                (*idx, {
-                    // Forward-declare all the functions
-                    // TODO: Pick appropriate return type depending on signature
-                    let fn_type = llvm_ctx.i64_type().fn_type(&[], false);
-                    let fn_value = module.add_function(
-                        // TODO: Determine function name from identifier
-                        // ir_ctx.symbol_map.resolve(function.symbol).unwrap(),
-                        "my function",
-                        fn_type,
-                        None,
-                    );
-
-                    fn_value
-                })
-            })),
+            function_values: HashMap::from_iter(ir_ctx.all_functions().iter().map(
+                |(idx, _function)| {
+                    (*idx, {
+                        // Forward-declare all the functions
+                        // TODO: Pick appropriate return type depending on signature
+                        let fn_type = llvm_ctx.i64_type().fn_type(&[], false);
+                        let fn_value = module.add_function(
+                            // TODO: Determine function name from identifier
+                            // ir_ctx.symbol_map.resolve(function.symbol).unwrap(),
+                            "my function",
+                            fn_type,
+                            None,
+                        );
+
+                        fn_value
+                    })
+                },
+            )),
             llvm_ctx,
             ir_ctx,
             module,
@@ -66,9 +69,12 @@ impl<'ctx> Pass<'ctx> {
     pub fn compile(&mut self, function_idx: FunctionIdx) -> FunctionValue<'ctx> {
         let function = self
             .ir_ctx
-            .functions
-            .get(&function_idx)
-            .expect("function to exist");
+            .all_functions()
+            .iter()
+            .find(|(idx, _)| idx == &function_idx)
+            .expect("function to exist")
+            .1
+            .clone();
 
         let builder = self.llvm_ctx.create_builder();
 
@@ -92,13 +98,7 @@ impl<'ctx> Pass<'ctx> {
             (
                 symbol.to_owned(),
                 builder
-                    .build_alloca(
-                        self.llvm_ctx.i64_type(),
-                        self.ir_ctx
-                            .symbol_map
-                            .resolve(*symbol)
-                            .expect("symbol to exist in map"),
-                    )
+                    .build_alloca(self.llvm_ctx.i64_type(), "todo: work out symbol name")
                     .unwrap(),
             )
         }));
@@ -123,9 +123,12 @@ impl<'ctx> Pass<'ctx> {
 
         let basic_block = self
             .ir_ctx
-            .functions
-            .get(&basic_block_id.0)
+            .all_functions()
+            .iter()
+            .find(|(idx, _)| idx == &basic_block_id.0)
             .expect("function to exist")
+            .1
+            .clone()
             .basic_blocks
             .get(basic_block_id.1)
             .expect("requested basic block must exist")

src/stage/lower_ir/ctx.rs

@@ -1,22 +1,78 @@
-use std::collections::HashMap;
-
 use crate::{
-    repr::{identifier::FunctionIdx, ir::Function},
-    util::symbol_map::interner_symbol_map::InternerSymbolMap,
+    repr::{
+        ast::typed as ast,
+        identifier::FunctionIdx,
+        ir::{BasicBlockIdx, Function, Triple, TripleRef},
+        ty::Ty,
+    },
+    util::symbol_map::interner_symbol_map::Symbol,
 };
 
-#[derive(Default, Clone, Debug)]
-pub struct IRCtx {
-    /// Map of function symbol to the basic block entry point
-    pub functions: HashMap<FunctionIdx, Function>,
-    pub symbol_map: InternerSymbolMap,
+pub trait IRCtx {
+    type FunctionBuilder: FunctionBuilder;
+
+    /// Register an IR function implementation against it's identifier.
+    fn register_function(&mut self, idx: FunctionIdx, function: Function);
+
+    /// Prepare and return a new builder for the provided AST function representation.
+    fn new_builder(&self, function: &ast::Function) -> Self::FunctionBuilder {
+        Self::FunctionBuilder::new(function)
+    }
+
+    // TODO: Probably get rid of this
+    fn all_functions(&self) -> Vec<(FunctionIdx, Function)>;
+}
+
+#[cfg(test)]
+mockall::mock! {
+    pub IRCtx {}
+
+    impl IRCtx for IRCtx {
+        type FunctionBuilder = MockFunctionBuilder;
+
+        fn register_function(&mut self, idx: FunctionIdx, function: Function);
+        fn new_builder(&self, function: &ast::Function) -> MockFunctionBuilder;
+        fn all_functions(&self) -> Vec<(FunctionIdx, Function)>;
+    }
 }
 
-impl IRCtx {
-    pub fn new(symbol_map: InternerSymbolMap) -> Self {
-        Self {
-            symbol_map,
-            ..Default::default()
-        }
+/// A stateful representation of a function that is being constructed. A function consists of basic
+///  blocks, and the builder is always 'located' at a basic block.
+pub trait FunctionBuilder {
+    /// Initialise a new builder with the provided function, positioned at the entry point.
+    fn new(function: &ast::Function) -> Self;
+
+    /// Register the provided symbol with the given type into the function scope.
+    fn register_scoped(&mut self, symbol: Symbol, ty: Ty);
+
+    /// Add a triple to the current basic block.
+    fn add_triple(&mut self, triple: Triple) -> TripleRef;
+
+    /// Get the current basic block.
+    fn current_bb(&self) -> BasicBlockIdx;
+
+    /// Go to a specific basic block.
+    fn goto_bb(&mut self, bb: BasicBlockIdx);
+
+    /// Create a new basic block, switch to it, and return its identifier.
+    fn push_bb(&mut self) -> BasicBlockIdx;
+
+    /// Consume the builder, and register the built function against the context.
+    fn build<Ctx: IRCtx<FunctionBuilder = Self>>(self, ctx: &mut Ctx);
+}
+
+#[cfg(test)]
+mockall::mock! {
+    pub FunctionBuilder {}
+
+    impl FunctionBuilder for FunctionBuilder {
+        fn new(function: &ast::Function) -> Self;
+        fn register_scoped(&mut self, symbol: Symbol, ty: Ty);
+        fn add_triple(&mut self, triple: Triple) -> TripleRef;
+        fn current_bb(&self) -> BasicBlockIdx;
+        fn goto_bb(&mut self, bb: BasicBlockIdx) ;
+        fn push_bb(&mut self) -> BasicBlockIdx;
+        #[mockall::concretize]
+        fn build<Ctx: IRCtx<FunctionBuilder = MockFunctionBuilder>>(self, ctx: &mut Ctx);
     }
 }