  1
+codegen.rs - source☰
  1
   2
   3
   4
@@ -135,10 +133,37 @@
 133
 134
 135
+136
+137
+138
+139
+140
+141
+142
+143
+144
+145
+146
+147
+148
+149
+150
+151
+152
+153
+154
+155
+156
+157
+158
+159
+160
+161
+162
 
 use std::collections::HashMap;
 
-use crate::llvm::{Builder, FnValue, Module, Value};
+use crate::llvm::{FnValue, FunctionPassManager, IRBuilder, Module, Value};
 use crate::parser::{ExprAST, FunctionAST, PrototypeAST};
 use crate::Either;
 
@@ -147,18 +172,23 @@
 /// Code generator from kaleidoscope AST to LLVM IR.
 pub struct Codegen<'llvm, 'a> {
     module: &'llvm Module,
-    builder: &'a Builder<'llvm>,
+    builder: &'a IRBuilder<'llvm>,
+    fpm: &'a FunctionPassManager<'llvm>,
+    fn_protos: &'a mut HashMap<String, PrototypeAST>,
 }
 
 impl<'llvm, 'a> Codegen<'llvm, 'a> {
     /// Compile either a [`PrototypeAST`] or a [`FunctionAST`] into the LLVM `module`.
     pub fn compile(
         module: &'llvm Module,
+        fn_protos: &mut HashMap<String, PrototypeAST>,
         compilee: Either<&PrototypeAST, &FunctionAST>,
     ) -> CodegenResult<FnValue<'llvm>> {
-        let cg = Codegen {
+        let mut cg = Codegen {
             module,
-            builder: &Builder::with_ctx(module),
+            builder: &IRBuilder::with_ctx(module),
+            fpm: &FunctionPassManager::with_ctx(module),
+            fn_protos,
         };
         let mut variables = HashMap::new();
 
@@ -195,7 +225,7 @@
                     _ => Err("invalid binary operator".into()),
                 }
             }
-            ExprAST::Call(callee, args) => match self.module.get_fn(callee) {
+            ExprAST::Call(callee, args) => match self.get_function(callee) {
                 Some(callee) => {
                     if callee.args() != args.len() {
                         return Err("Incorrect # arguments passed".into());
@@ -235,14 +265,16 @@
     }
 
     fn codegen_function(
-        &self,
+        &mut self,
         FunctionAST(proto, body): &FunctionAST,
         named_values: &mut HashMap<&'llvm str, Value<'llvm>>,
     ) -> CodegenResult<FnValue<'llvm>> {
-        let the_function = match self.module.get_fn(&proto.0) {
-            Some(f) => f,
-            None => self.codegen_prototype(proto),
-        };
+        // Insert the function prototype into the `fn_protos` map to keep track for re-generating
+        // declarations in other modules.
+        self.fn_protos.insert(proto.0.clone(), proto.clone());
+
+        let the_function = self.get_function(&proto.0)
+            .expect("If proto not already generated, get_function will do for us since we updated fn_protos before-hand!");
 
         if the_function.basic_blocks() > 0 {
             return Err("Function cannot be redefined.".into());
@@ -265,11 +297,33 @@
         if let Ok(ret) = self.codegen_expr(body, named_values) {
             self.builder.ret(ret);
             assert!(the_function.verify());
+
+            // Run the optimization passes on the function.
+            self.fpm.run(the_function);
+
             Ok(the_function)
         } else {
             todo!("Failed to codegen function body, erase from module!");
         }
     }
+
+    /// Lookup function with `name` in the LLVM module and return the corresponding value reference.
+    /// If the function is not available in the module, check if the prototype is known and codegen
+    /// it.
+    /// Return [`None`] if the prototype is not known.
+    fn get_function(&self, name: &str) -> Option<FnValue<'llvm>> {
+        let callee = match self.module.get_fn(name) {
+            Some(callee) => callee,
+            None => {
+                let proto = self.fn_protos.get(name)?;
+                self.codegen_prototype(proto)
+            }
+        };
+
+        Some(callee)
+    }
 }
 
-

\ No newline at end of file
+