diff --git a/examples/reid/arithmetic.reid b/examples/reid/arithmetic.reid
index 9de6f4f..270cfcb 100644
--- a/examples/reid/arithmetic.reid
+++ b/examples/reid/arithmetic.reid
@@ -5,6 +5,10 @@ fn main() {
     let simpleAdd = 2 + 2;
     let simpleMult = 7 * 2; // 14
     let arithmetic = 3 + 2 * 5 + 1 * 2; // 15
+    
+    if simpleAdd < test {
+        return 3;
+    }
 
     return arithmetic + simpleMult * arithmetic;
 }
\ No newline at end of file
diff --git a/src/codegen/llvm.rs b/src/codegen/llvm.rs
index cc5a2fa..03df00e 100644
--- a/src/codegen/llvm.rs
+++ b/src/codegen/llvm.rs
@@ -24,7 +24,7 @@ pub enum IRType {
 }
 
 impl IRType {
-    fn in_context(&self, context: &mut IRContext) -> *mut LLVMType {
+    fn in_context(&self, context: &IRContext) -> *mut LLVMType {
         use IRType::*;
         unsafe {
             return match self {
@@ -38,14 +38,14 @@ impl IRType {
 pub struct IRValue(pub IRType, *mut LLVMValue);
 
 impl IRValue {
-    pub fn from_literal(literal: &ast::Literal, block: &mut IRBlock) -> Self {
+    pub fn from_literal(literal: &ast::Literal, module: &IRModule) -> Self {
         use ast::Literal;
         match literal {
             Literal::I32(v) => {
                 let ir_type = IRType::I32;
                 unsafe {
                     let ir_value = LLVMConstInt(
-                        ir_type.in_context(block.function.module.context),
+                        ir_type.in_context(module.context),
                         mem::transmute(*v as i64),
                         1,
                     );
@@ -116,13 +116,13 @@ impl<'a> Drop for IRModule<'a> {
 }
 
 pub struct IRFunction<'a, 'b> {
-    module: &'b mut IRModule<'a>,
+    pub module: &'b IRModule<'a>,
     /// The actual function
     value: *mut LLVMValue,
 }
 
 impl<'a, 'b> IRFunction<'a, 'b> {
-    pub fn new(name: &String, module: &'b mut IRModule<'a>) -> IRFunction<'a, 'b> {
+    pub fn new(name: &String, module: &'b IRModule<'a>) -> IRFunction<'a, 'b> {
         unsafe {
             // TODO, fix later!
 
@@ -140,23 +140,26 @@ impl<'a, 'b> IRFunction<'a, 'b> {
             }
         }
     }
+
+    pub fn attach(&mut self, block: IRBlock) {
+        unsafe { LLVMAppendExistingBasicBlock(self.value, block.blockref) }
+    }
 }
 
-pub struct IRBlock<'a, 'b, 'c> {
-    function: &'a mut IRFunction<'b, 'c>,
+pub struct IRBlock<'a, 'b> {
+    pub module: &'b IRModule<'a>,
     blockref: *mut LLVMBasicBlock,
 }
 
-impl<'a, 'b, 'c> IRBlock<'a, 'b, 'c> {
-    pub fn new(function: &'a mut IRFunction<'b, 'c>) -> IRBlock<'a, 'b, 'c> {
+impl<'a, 'b, 'c> IRBlock<'a, 'b> {
+    pub fn new(module: &'b IRModule<'a>) -> IRBlock<'a, 'b> {
         unsafe {
             let blockref = LLVMCreateBasicBlockInContext(
-                function.module.context.context,
+                module.context.context,
                 into_cstring("entryblock").as_ptr(),
             );
-            LLVMPositionBuilderAtEnd(function.module.context.builder, blockref);
 
-            IRBlock { function, blockref }
+            IRBlock { module, blockref }
         }
     }
 
@@ -166,11 +169,12 @@ impl<'a, 'b, 'c> IRBlock<'a, 'b, 'c> {
         IRValue(rhs_t, rhs_v): IRValue,
     ) -> Result<IRValue, Error> {
         unsafe {
+            LLVMPositionBuilderAtEnd(self.module.context.builder, self.blockref);
             if lhs_t == rhs_t {
                 Ok(IRValue(
                     lhs_t,
                     LLVMBuildAdd(
-                        self.function.module.context.builder,
+                        self.module.context.builder,
                         lhs_v,
                         rhs_v,
                         c"tmpadd".as_ptr(),
@@ -188,11 +192,12 @@ impl<'a, 'b, 'c> IRBlock<'a, 'b, 'c> {
         IRValue(rhs_t, rhs_v): IRValue,
     ) -> Result<IRValue, Error> {
         unsafe {
+            LLVMPositionBuilderAtEnd(self.module.context.builder, self.blockref);
             if lhs_t == rhs_t {
                 Ok(IRValue(
                     lhs_t,
                     LLVMBuildMul(
-                        self.function.module.context.builder,
+                        self.module.context.builder,
                         lhs_v,
                         rhs_v,
                         c"tmpadd".as_ptr(),
@@ -204,21 +209,14 @@ impl<'a, 'b, 'c> IRBlock<'a, 'b, 'c> {
         }
     }
 
-    pub fn add_return(self, value: Option<IRValue>) {
+    pub fn add_return(&mut self, value: Option<IRValue>) {
         unsafe {
+            LLVMPositionBuilderAtEnd(self.module.context.builder, self.blockref);
             if let Some(IRValue(_, value)) = value {
-                LLVMBuildRet(self.function.module.context.builder, value);
+                LLVMBuildRet(self.module.context.builder, value);
             } else {
-                LLVMBuildRetVoid(self.function.module.context.builder);
+                LLVMBuildRetVoid(self.module.context.builder);
             }
         }
     }
 }
-
-impl<'a, 'b, 'c> Drop for IRBlock<'a, 'b, 'c> {
-    fn drop(&mut self) {
-        unsafe {
-            LLVMAppendExistingBasicBlock(self.function.value, self.blockref);
-        }
-    }
-}
diff --git a/src/codegen/mod.rs b/src/codegen/mod.rs
index 88abc3b..9b1d2d6 100644
--- a/src/codegen/mod.rs
+++ b/src/codegen/mod.rs
@@ -6,8 +6,8 @@ use llvm::{Error, IRBlock, IRContext, IRFunction, IRModule, IRValue};
 
 use crate::{
     ast::{
-        BinaryOperator, Block, BlockLevelStatement, Expression, ExpressionKind, FunctionDefinition,
-        LetStatement, ReturnType,
+        Block, BlockLevelStatement, Expression, ExpressionKind, FunctionDefinition, LetStatement,
+        ReturnType,
     },
     TopLevelStatement,
 };
@@ -42,19 +42,23 @@ impl FunctionDefinition {
     fn codegen(&self, scope: &mut ScopeData, module: &mut IRModule) {
         let FunctionDefinition(signature, block, _) = self;
         let mut ir_function = IRFunction::new(&signature.name, module);
-        let ir_block = IRBlock::new(&mut ir_function);
-        block.codegen(scope.inner(ir_block));
+
+        let ir_block = IRBlock::new(&module);
+        let mut scope = scope.inner(ir_block);
+        block.codegen(&mut scope);
+
+        ir_function.attach(scope.block);
     }
 }
 
 impl Block {
-    fn codegen(&self, mut scope: Scope) {
+    fn codegen(&self, scope: &mut Scope) {
         for statement in &self.0 {
-            statement.codegen(&mut scope);
+            statement.codegen(scope);
         }
 
         if let Some((_, return_exp)) = &self.1 {
-            let value = return_exp.codegen(&mut scope);
+            let value = return_exp.codegen(scope);
             scope.block.add_return(Some(value));
         }
     }
@@ -83,7 +87,7 @@ impl Expression {
 
         use ExpressionKind::*;
         match kind {
-            Literal(lit) => IRValue::from_literal(lit, &mut scope.block),
+            Literal(lit) => IRValue::from_literal(lit, &mut scope.block.module),
             VariableName(v) => scope.data.fetch(v),
             Binop(op, lhs, rhs) => {
                 let lhs = lhs.codegen(scope);
@@ -112,11 +116,11 @@ impl ScopeData {
         }
     }
 
-    fn with_block<'a, 'b, 'c>(self, block: IRBlock<'a, 'b, 'c>) -> Scope<'a, 'b, 'c> {
+    fn with_block<'a, 'b>(self, block: IRBlock<'a, 'b>) -> Scope<'a, 'b> {
         Scope { data: self, block }
     }
 
-    fn inner<'a, 'b, 'c>(&self, block: IRBlock<'a, 'b, 'c>) -> Scope<'a, 'b, 'c> {
+    fn inner<'a, 'b>(&self, block: IRBlock<'a, 'b>) -> Scope<'a, 'b> {
         self.clone().with_block(block)
     }
 
@@ -135,13 +139,13 @@ impl ScopeData {
     }
 }
 
-struct Scope<'a, 'b, 'c> {
+struct Scope<'a, 'b> {
     data: ScopeData,
-    block: IRBlock<'a, 'b, 'c>,
+    block: IRBlock<'a, 'b>,
 }
 
-impl<'a, 'b, 'c> Scope<'a, 'b, 'c> {
-    fn inner(&self, block: IRBlock<'a, 'b, 'c>) -> Scope<'a, 'b, 'c> {
+impl<'a, 'b> Scope<'a, 'b> {
+    fn inner(&self, block: IRBlock<'a, 'b>) -> Scope<'a, 'b> {
         self.data.clone().with_block(block)
     }
 }