ch4: add function optimization passes

2 files changed, 79 insertions, 7 deletions

diff --git a/src/codegen.rs b/src/codegen.rs
index 866aab7..0a7893e 100644
--- a/src/codegen.rs
+++ b/src/codegen.rs
@@ -1,6 +1,6 @@
 use std::collections::HashMap;
 
-use crate::llvm::{Builder, FnValue, Module, Value};
+use crate::llvm::{Builder, FnValue, FunctionPassManager, Module, Value};
 use crate::parser::{ExprAST, FunctionAST, PrototypeAST};
 use crate::Either;
 
@@ -10,6 +10,7 @@ type CodegenResult<T> = Result<T, String>;
 pub struct Codegen<'llvm, 'a> {
     module: &'llvm Module,
     builder: &'a Builder<'llvm>,
+    fpm: &'a FunctionPassManager<'llvm>,
 }
 
 impl<'llvm, 'a> Codegen<'llvm, 'a> {
@@ -21,6 +22,7 @@ impl<'llvm, 'a> Codegen<'llvm, 'a> {
         let cg = Codegen {
             module,
             builder: &Builder::with_ctx(module),
+            fpm: &FunctionPassManager::with_ctx(module),
         };
         let mut variables = HashMap::new();
 
@@ -127,6 +129,7 @@ impl<'llvm, 'a> Codegen<'llvm, 'a> {
         if let Ok(ret) = self.codegen_expr(body, named_values) {
             self.builder.ret(ret);
             assert!(the_function.verify());
+            self.fpm.run(the_function);
             Ok(the_function)
         } else {
             todo!("Failed to codegen function body, erase from module!");
diff --git a/src/llvm.rs b/src/llvm.rs
index ed6b930..5598cd4 100644
--- a/src/llvm.rs
+++ b/src/llvm.rs
@@ -10,14 +10,19 @@ use llvm_sys::core::{
     LLVMAddFunction, LLVMAppendBasicBlockInContext, LLVMBuildFAdd, LLVMBuildFCmp, LLVMBuildFMul,
     LLVMBuildFSub, LLVMBuildRet, LLVMBuildUIToFP, LLVMConstReal, LLVMContextCreate,
     LLVMContextDispose, LLVMCountBasicBlocks, LLVMCountParams, LLVMCreateBuilderInContext,
-    LLVMDisposeBuilder, LLVMDisposeModule, LLVMDoubleTypeInContext, LLVMDumpModule, LLVMDumpType,
-    LLVMDumpValue, LLVMGetNamedFunction, LLVMGetParam, LLVMGetReturnType, LLVMGetTypeKind,
-    LLVMGetValueKind, LLVMGetValueName2, LLVMModuleCreateWithNameInContext,
-    LLVMPositionBuilderAtEnd, LLVMSetValueName2, LLVMTypeOf,
+    LLVMCreateFunctionPassManagerForModule, LLVMDisposeBuilder, LLVMDisposeModule,
+    LLVMDisposePassManager, LLVMDoubleTypeInContext, LLVMDumpModule, LLVMDumpType, LLVMDumpValue,
+    LLVMGetNamedFunction, LLVMGetParam, LLVMGetReturnType, LLVMGetTypeKind, LLVMGetValueKind,
+    LLVMGetValueName2, LLVMInitializeFunctionPassManager, LLVMModuleCreateWithNameInContext,
+    LLVMPositionBuilderAtEnd, LLVMRunFunctionPassManager, LLVMSetValueName2, LLVMTypeOf,
 };
 use llvm_sys::prelude::{
-    LLVMBasicBlockRef, LLVMBool, LLVMBuilderRef, LLVMContextRef, LLVMModuleRef, LLVMTypeRef,
-    LLVMValueRef,
+    LLVMBasicBlockRef, LLVMBool, LLVMBuilderRef, LLVMContextRef, LLVMModuleRef, LLVMPassManagerRef,
+    LLVMTypeRef, LLVMValueRef,
+};
+use llvm_sys::transforms::{
+    instcombine::LLVMAddInstructionCombiningPass,
+    scalar::{LLVMAddCFGSimplificationPass, LLVMAddNewGVNPass, LLVMAddReassociatePass},
 };
 use llvm_sys::{LLVMRealPredicate, LLVMTypeKind, LLVMValueKind};
 
@@ -532,3 +537,67 @@ impl<'llvm> FnValue<'llvm> {
         }
     }
 }
+
+// =======================
+//   FunctionPassManager
+// =======================
+
+/// Wrapper for a LLVM Function PassManager (legacy).
+pub struct FunctionPassManager<'llvm> {
+    fpm: LLVMPassManagerRef,
+    _ctx: PhantomData<&'llvm ()>,
+}
+
+impl<'llvm> FunctionPassManager<'llvm> {
+    /// Create a new Function PassManager with the following optimization passes
+    /// - InstructionCombiningPass
+    /// - ReassociatePass
+    /// - NewGVNPass
+    /// - CFGSimplificationPass
+    ///
+    /// The list of selected optimization passes is taken from the tutorial chapter [LLVM
+    /// Optimization Passes](https://llvm.org/docs/tutorial/MyFirstLanguageFrontend/LangImpl04.html#id3).
+    pub fn with_ctx(module: &'llvm Module) -> FunctionPassManager<'llvm> {
+        let fpm = unsafe {
+            // Borrows module reference.
+            LLVMCreateFunctionPassManagerForModule(module.module)
+        };
+        assert!(!fpm.is_null());
+
+        unsafe {
+            // Do simple "peephole" optimizations and bit-twiddling optzns.
+            LLVMAddInstructionCombiningPass(fpm);
+            // Reassociate expressions.
+            LLVMAddReassociatePass(fpm);
+            // Eliminate Common SubExpressions.
+            LLVMAddNewGVNPass(fpm);
+            // Simplify the control flow graph (deleting unreachable blocks, etc).
+            LLVMAddCFGSimplificationPass(fpm);
+
+            let fail = LLVMInitializeFunctionPassManager(fpm);
+            assert_eq!(fail, 0);
+        }
+
+        FunctionPassManager {
+            fpm,
+            _ctx: PhantomData,
+        }
+    }
+
+    /// Run the optimization passes registered with the Function PassManager on the function
+    /// referenced by `fn_value`.
+    pub fn run(&'llvm self, fn_value: FnValue<'llvm>) {
+        unsafe {
+            // Returns 1 if any of the passes modified the function, false otherwise.
+            LLVMRunFunctionPassManager(self.fpm, fn_value.value_ref());
+        }
+    }
+}
+
+impl Drop for FunctionPassManager<'_> {
+    fn drop(&mut self) {
+        unsafe {
+            LLVMDisposePassManager(self.fpm);
+        }
+    }
+}