diff --git a/src/codegen.rs b/src/codegen.rs
index de13576981..6e41833792 100644
--- a/src/codegen.rs
+++ b/src/codegen.rs
@@ -10,7 +10,7 @@ use self::{
     },
     llvm_index::LlvmTypedIndex,
 };
-use crate::compile_error::CompileError;
+use crate::{compile_error::CompileError, resolver::AnnotationMap};
 
 use super::ast::*;
 use super::index::*;
@@ -93,7 +93,8 @@ impl<'ink> CodeGen<'ink> {
     /// generates all TYPEs, GLOBAL-sections and POUs of the given CompilationUnit
     pub fn generate(
         &self,
-        unit: CompilationUnit,
+        unit: &CompilationUnit,
+        _annotations: &AnnotationMap,
         global_index: &Index,
     ) -> Result<String, CompileError> {
         //Associate the index type with LLVM types
@@ -103,10 +104,10 @@ impl<'ink> CodeGen<'ink> {
         let llvm = Llvm::new(self.context, self.context.create_builder());
         let pou_generator = PouGenerator::new(llvm, global_index, &llvm_index);
         //Generate the POU stubs in the first go to make sure they can be referenced.
-        for implementation in unit.implementations {
+        for implementation in &unit.implementations {
             //Don't generate external functions
             if implementation.linkage != LinkageType::External {
-                pou_generator.generate_implementation(&implementation)?;
+                pou_generator.generate_implementation(implementation)?;
             }
         }
         Ok(self.module.print_to_string().to_string())
diff --git a/src/codegen/tests.rs b/src/codegen/tests.rs
index 263c02a1d6..f69a59d0fd 100644
--- a/src/codegen/tests.rs
+++ b/src/codegen/tests.rs
@@ -12,8 +12,9 @@ macro_rules! codegen_wihout_unwrap {
         let context = inkwell::context::Context::create();
         crate::ast::pre_process(&mut ast);
         let index = crate::index::visitor::visit(&ast);
+        let annotations = crate::resolver::TypeAnnotator::visit_unit(&index, &ast);
         let code_generator = crate::codegen::CodeGen::new(&context, "main");
-        code_generator.generate(ast, &index)
+        code_generator.generate(&ast, &annotations, &index)
     }};
 }
 
@@ -26,8 +27,9 @@ macro_rules! codegen {
         let context = inkwell::context::Context::create();
         crate::ast::pre_process(&mut ast);
         let index = crate::index::visitor::visit(&ast);
+        let annotations = crate::resolver::TypeAnnotator::visit_unit(&index, &ast);
         let code_generator = crate::codegen::CodeGen::new(&context, "main");
-        code_generator.generate(ast, &index).unwrap()
+        code_generator.generate(&ast, &annotations, &index).unwrap()
     }};
 }
 
diff --git a/src/lib.rs b/src/lib.rs
index f975d7b667..afba998b43 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -39,6 +39,7 @@ use std::{fs::File, io::Read};
 use validation::Validator;
 
 use crate::ast::CompilationUnit;
+use crate::resolver::AnnotationMap;
 mod ast;
 pub mod cli;
 mod codegen;
@@ -402,25 +403,26 @@ pub fn compile_module<'c, T: SourceContainer>(
 
     // ### PHASE 2 ###
     // annotation & validation everything
-    for (file_id, syntax_errors, container) in all_units.iter() {
-        let annotations = TypeAnnotator::visit_unit(&full_index, container);
+    type AnnotatedAst<'a> = (&'a CompilationUnit, AnnotationMap);
+    let mut annotated_units: Vec<AnnotatedAst> = Vec::new();
+    for (file_id, syntax_errors, unit) in all_units.iter() {
+        let annotations = TypeAnnotator::visit_unit(&full_index, unit);
 
         let mut validator = Validator::new();
-        validator.visit_unit(&annotations, container);
+        validator.visit_unit(&annotations, unit);
         //log errors
         report_diagnostics(*file_id, syntax_errors.iter(), &files)?;
         report_diagnostics(*file_id, validator.diagnostics().iter(), &files)?;
+
+        annotated_units.push((unit, annotations));
     }
 
     // ### PHASE 3 ###
     // - codegen
-    let mut code_gen_unit = CompilationUnit::default();
-    for (_, _, u) in all_units {
-        code_gen_unit.import(u);
-    }
-
     let code_generator = codegen::CodeGen::new(context, "main");
-    code_generator.generate(code_gen_unit, &full_index)?;
+    for (unit, annotations) in annotated_units {
+        code_generator.generate(unit, &annotations, &full_index)?;
+    }
     Ok(code_generator)
 }
 
diff --git a/src/resolver.rs b/src/resolver.rs
index 662c15d023..ec25e51456 100644
--- a/src/resolver.rs
+++ b/src/resolver.rs
@@ -177,6 +177,12 @@ impl AnnotationMap {
     }
 }
 
+impl Default for AnnotationMap {
+    fn default() -> Self {
+        Self::new()
+    }
+}
+
 impl<'i> TypeAnnotator<'i> {
     /// constructs a new TypeAnnotater that works with the given index for type-lookups
     fn new(index: &'i Index) -> TypeAnnotator<'i> {