make terminator different to a triple value

src/compile_pass.rs

@@ -6,7 +6,7 @@ use crate::{
     repr::{
         ast::typed::Function,
         identifier::{FunctionIdx, ScopedBinding},
-        ir,
+        ir::{self, Terminator, Triple},
         ty::Ty,
     },
     stage::{
@@ -116,11 +116,17 @@ impl IRCtx for CompilePass {
     }
 }
 
+#[derive(Default)]
+struct BasicBlockBuilder {
+    triples: IndexVec<ir::TripleIdx, Triple>,
+    terminator: Option<Terminator>,
+}
+
 pub struct FunctionBuilder {
     idx: FunctionIdx,
     signature: FunctionSignature,
 
-    basic_blocks: IndexVec<ir::BasicBlockIdx, ir::BasicBlock>,
+    basic_blocks: IndexVec<ir::BasicBlockIdx, BasicBlockBuilder>,
     current_basic_block: ir::BasicBlockIdx,
 
     scope: Vec<(ScopedBinding, Ty)>,
@@ -129,7 +135,7 @@ pub struct FunctionBuilder {
 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());
+        let current_basic_block = basic_blocks.push(BasicBlockBuilder::default());
 
         Self {
             idx: function.name,
@@ -153,6 +159,17 @@ impl FunctionBuilderTrait for FunctionBuilder {
         }
     }
 
+    fn set_terminator(&mut self, terminator: ir::Terminator) {
+        let bb = &mut self.basic_blocks[self.current_basic_block];
+
+        assert!(
+            bb.terminator.is_none(),
+            "cannot set terminator if it's already been set"
+        );
+
+        bb.terminator = Some(terminator);
+    }
+
     fn current_bb(&self) -> ir::BasicBlockIdx {
         self.current_basic_block
     }
@@ -166,7 +183,7 @@ impl FunctionBuilderTrait for FunctionBuilder {
     }
 
     fn push_bb(&mut self) -> ir::BasicBlockIdx {
-        let idx = self.basic_blocks.push(ir::BasicBlock::default());
+        let idx = self.basic_blocks.push(BasicBlockBuilder::default());
 
         self.current_basic_block = idx;
 
@@ -179,7 +196,14 @@ impl FunctionBuilderTrait for FunctionBuilder {
             ir::Function {
                 identifier: self.idx,
                 signature: self.signature,
-                basic_blocks: self.basic_blocks,
+                basic_blocks: self
+                    .basic_blocks
+                    .into_iter()
+                    .map(|builder| ir::BasicBlock {
+                        triples: builder.triples,
+                        terminator: builder.terminator.expect("terminator must be set"),
+                    })
+                    .collect(),
                 scope: self.scope.into_iter().map(|(symbol, _)| symbol).collect(),
             },
         )

src/repr/ir/basic_block.rs

@@ -1,8 +1,9 @@
 use index_vec::IndexVec;
 
-use super::{Triple, TripleIdx};
+use super::{terminator::Terminator, Triple, TripleIdx};
 
-#[derive(Default, Clone, Debug)]
+#[derive(Clone, Debug)]
 pub struct BasicBlock {
     pub triples: IndexVec<TripleIdx, Triple>,
+    pub terminator: Terminator,
 }

src/repr/ir/mod.rs

@@ -1,9 +1,11 @@
 mod basic_block;
 mod function;
+mod terminator;
 mod triple;
 mod value;
 
 pub use basic_block::*;
 pub use function::*;
+pub use terminator::*;
 pub use triple::*;
 pub use value::*;

src/repr/ir/terminator.rs

@@ -0,0 +1,14 @@
+use super::*;
+
+#[derive(Clone, Debug)]
+pub enum Terminator {
+    /// Jump to the corresponding basic block.
+    Jump(BasicBlockIdx),
+    /// Return with the provided value.
+    Return(Value),
+    Switch {
+        value: Value,
+        default: BasicBlockIdx,
+        branches: Vec<(Value, BasicBlockIdx)>,
+    },
+}

src/repr/ir/triple/mod.rs

@@ -25,19 +25,10 @@ pub enum Triple {
     },
     /// Copy the provided value.
     Copy(Value),
-    /// Jump to the corresponding basic block.
-    Jump(BasicBlockIdx),
     /// Call the corresponding function.
     Call(FunctionIdx, Vec<Value>),
-    /// Return with the provided value.
-    Return(Value),
     /// Assign some symbol to some value.
     Assign(ScopedBinding, Value),
-    Switch {
-        value: Value,
-        default: BasicBlockIdx,
-        branches: Vec<(Value, BasicBlockIdx)>,
-    },
     Phi(Vec<(Value, BasicBlockIdx)>),
 }
 

src/stage/codegen/llvm/mod.rs

@@ -10,7 +10,10 @@ use inkwell::{
 
 use crate::repr::{
     identifier::{FunctionIdx, ScopedBinding},
-    ir::{BasicBlockIdx, BinaryOp, ConstantValue, Function, Triple, TripleRef, UnaryOp, Value},
+    ir::{
+        BasicBlockIdx, BinaryOp, ConstantValue, Function, Terminator, Triple, TripleRef, UnaryOp,
+        Value,
+    },
     ty::Ty,
 };
 
@@ -125,32 +128,27 @@ impl<'ctx, 'ink, Ctx: LLVMCtx> FunctionGenerator<'ctx, 'ink, Ctx> {
                 Triple::BinaryOp { lhs, rhs, op } => Some(self.gen_op_binary(lhs, rhs, op)),
                 Triple::UnaryOp { rhs, op } => Some(self.gen_op_unary(rhs, op)),
                 Triple::Copy(value) => Some(self.gen_copy(value)),
-                Triple::Jump(bb) => {
-                    self.gen_jump(bb);
-                    None
-                }
                 Triple::Call(function, params) => Some(self.gen_call(function, params)),
-                Triple::Return(value) => {
-                    self.gen_return(value);
-                    None
-                }
                 Triple::Assign(symbol, value) => {
                     self.gen_assign(symbol, value);
                     None
                 }
-                Triple::Switch {
-                    value,
-                    default,
-                    branches,
-                } => {
-                    self.gen_switch(value, default, branches);
-                    None
-                }
                 Triple::Phi(values) => Some(self.gen_phi(values)),
             };
             self.results.insert(TripleRef::new(*block_idx, idx), result);
         }
 
+        // Lower the block terminator
+        match &block.terminator {
+            Terminator::Jump(bb) => self.gen_jump(bb),
+            Terminator::Return(value) => self.gen_return(value),
+            Terminator::Switch {
+                value,
+                default,
+                branches,
+            } => self.gen_switch(value, default, branches),
+        }
+
         if let Some(prev) = prev_builder {
             self.builder.position_at_end(prev);
         }

src/stage/lower_ir/ctx.rs

@@ -1,7 +1,7 @@
 use crate::repr::{
     ast::typed as ast,
     identifier::{FunctionIdx, ScopedBinding},
-    ir::{BasicBlockIdx, Function, Triple, TripleRef},
+    ir::{BasicBlockIdx, Function, Terminator, Triple, TripleRef},
     ty::Ty,
 };
 
@@ -33,6 +33,9 @@ pub trait FunctionBuilder {
     /// Add a triple to the current basic block.
     fn add_triple(&mut self, triple: Triple) -> TripleRef;
 
+    /// Set the terminator of the current basic block. Will panic if the terminator has already been set.
+    fn set_terminator(&mut self, terminator: Terminator);
+
     /// Get the current basic block.
     fn current_bb(&self) -> BasicBlockIdx;
 
@@ -54,6 +57,7 @@ mockall::mock! {
         fn new(function: &ast::Function) -> Self;
         fn register_scoped(&mut self, ident: ScopedBinding, ty: Ty);
         fn add_triple(&mut self, triple: Triple) -> TripleRef;
+        fn set_terminator(&mut self, terminator: Terminator);
         fn current_bb(&self) -> BasicBlockIdx;
         fn goto_bb(&mut self, bb: BasicBlockIdx) ;
         fn push_bb(&mut self) -> BasicBlockIdx;

src/stage/lower_ir/lowering/mod.rs

@@ -55,7 +55,7 @@ fn lower_block(
         match statement {
             ast::Statement::Return(ast::ReturnStatement { value, .. }) => {
                 let value = lower_expression(ctx, builder, value).unwrap();
-                builder.add_triple(Triple::Return(value));
+                builder.set_terminator(Terminator::Return(value));
             }
             ast::Statement::Let(ast::LetStatement {
                 binding: name,
@@ -136,7 +136,7 @@ fn lower_expression(
 
             if let (Some(merge_bb), true) = (merge_bb, !matches!(success.ty_info.ty, Ty::Never)) {
                 // Ensure the branch returns to the merge basic block
-                builder.add_triple(Triple::Jump(merge_bb));
+                builder.set_terminator(Terminator::Jump(merge_bb));
             }
 
             let mut merge_values = vec![(success_value, success_bb)];
@@ -150,7 +150,7 @@ fn lower_expression(
 
                     if let Some(merge_bb) = merge_bb {
                         // Ensure the branch returns to the merge basic block
-                        builder.add_triple(Triple::Jump(merge_bb));
+                        builder.set_terminator(Terminator::Jump(merge_bb));
                     }
 
                     merge_values.push((otherwise_value, otherwise_bb));
@@ -166,7 +166,7 @@ fn lower_expression(
             // Revert back to original location
             builder.goto_bb(original_bb);
 
-            builder.add_triple(Triple::Switch {
+            builder.set_terminator(Terminator::Switch {
                 value: condition,
                 default: success_bb,
                 branches,