Added LLVM callbacks

lib/compiler-llvm/src/config.rs

@@ -5,7 +5,9 @@ use inkwell::targets::{
 };
 use inkwell::OptimizationLevel;
 use itertools::Itertools;
+use std::sync::Arc;
 use target_lexicon::Architecture;
+use wasm_common::{FunctionType, LocalFunctionIndex};
 use wasmer_compiler::{Compiler, CompilerConfig, CpuFeature, Features, Target, Triple};
 
 /// The InkWell ModuleInfo type
@@ -14,11 +16,26 @@ pub type InkwellModule<'ctx> = inkwell::module::Module<'ctx>;
 /// The InkWell MemoryBuffer type
 pub type InkwellMemoryBuffer = inkwell::memory_buffer::MemoryBuffer;
 
-/// Callbacks to
-pub trait LLVMCallbacks: std::any::Any + 'static {
-    fn preopt_ir_callback(&mut self, module: &InkwellModule);
-    fn postopt_ir_callback(&mut self, module: &InkwellModule);
-    fn obj_memory_buffer_callback(&mut self, memory_buffer: &InkwellMemoryBuffer);
+/// The compiled function kind, used for debugging in the `LLVMCallbacks`.
+#[derive(Debug, Clone)]
+pub enum CompiledFunctionKind {
+    // A locally-defined function in the Wasm file
+    Local(LocalFunctionIndex),
+    // A function call trampoline for a given signature
+    FunctionCallTrampoline(FunctionType),
+    // A dynamic function trampoline for a given signature
+    DynamicFunctionTrampoline(FunctionType),
+}
+
+/// Callbacks to the different LLVM compilation phases
+pub trait LLVMCallbacks: Send + Sync {
+    fn preopt_ir(&self, function: &CompiledFunctionKind, module: &InkwellModule);
+    fn postopt_ir(&self, function: &CompiledFunctionKind, module: &InkwellModule);
+    fn obj_memory_buffer(
+        &self,
+        function: &CompiledFunctionKind,
+        memory_buffer: &InkwellMemoryBuffer,
+    );
 }
 
 #[derive(Clone)]
@@ -40,6 +57,10 @@ pub struct LLVMConfig {
     /// Whether to emit PIC.
     pub is_pic: bool,
 
+    /// Callbacks that will triggered in the different compilation
+    /// phases in LLVM.
+    pub callbacks: Option<Arc<dyn LLVMCallbacks>>,
+
     features: Features,
     target: Target,
 }
@@ -78,6 +99,7 @@ impl LLVMConfig {
             is_pic: false,
             features,
             target,
+            callbacks: None,
         }
     }
     fn reloc_mode(&self) -> RelocMode {

lib/compiler-llvm/src/lib.rs

@@ -21,4 +21,6 @@ mod trampoline;
 mod translator;
 
 pub use crate::compiler::LLVMCompiler;
-pub use crate::config::{InkwellMemoryBuffer, InkwellModule, LLVMCallbacks, LLVMConfig};
+pub use crate::config::{
+    CompiledFunctionKind, InkwellMemoryBuffer, InkwellModule, LLVMCallbacks, LLVMConfig,
+};

lib/compiler-llvm/src/trampoline/wasm.rs

@@ -1,4 +1,4 @@
-use crate::config::LLVMConfig;
+use crate::config::{CompiledFunctionKind, LLVMConfig};
 use crate::object_file::load_object_file;
 use crate::translator::intrinsics::{
     func_type_to_llvm, type_to_llvm, type_to_llvm_ptr, Intrinsics,
@@ -36,6 +36,8 @@ impl FuncTrampoline {
         ty: &FunctionType,
         config: &LLVMConfig,
     ) -> Result<FunctionBody, CompileError> {
+        // The function type, used for the callbacks
+        let function = CompiledFunctionKind::FunctionCallTrampoline(ty.clone());
         let module = self.ctx.create_module("");
         let target_triple = config.target_triple();
         let target_machine = config.target_machine();
@@ -61,8 +63,9 @@ impl FuncTrampoline {
             .set_section(FUNCTION_SECTION);
         generate_trampoline(trampoline_func, ty, &self.ctx, &intrinsics)?;
 
-        // TODO: remove debugging
-        //module.print_to_stderr();
+        if let Some(ref callbacks) = config.callbacks {
+            callbacks.preopt_ir(&function, &module);
+        }
 
         let pass_manager = PassManager::create(());
 
@@ -74,22 +77,17 @@ impl FuncTrampoline {
 
         pass_manager.run_on(&module);
 
-        // TODO: remove debugging
-        //module.print_to_stderr();
+        if let Some(ref callbacks) = config.callbacks {
+            callbacks.postopt_ir(&function, &module);
+        }
 
         let memory_buffer = target_machine
             .write_to_memory_buffer(&module, FileType::Object)
             .unwrap();
 
-        /*
-        // TODO: remove debugging
-        let mem_buf_slice = memory_buffer.as_slice();
-        let mut file = fs::File::create("/home/nicholas/trampoline.o").unwrap();
-        let mut pos = 0;
-        while pos < mem_buf_slice.len() {
-            pos += file.write(&mem_buf_slice[pos..]).unwrap();
+        if let Some(ref callbacks) = config.callbacks {
+            callbacks.obj_memory_buffer(&function, &memory_buffer);
         }
-        */
 
         let mem_buf_slice = memory_buffer.as_slice();
         let (function, sections) =
@@ -127,6 +125,8 @@ impl FuncTrampoline {
         ty: &FunctionType,
         config: &LLVMConfig,
     ) -> Result<FunctionBody, CompileError> {
+        // The function type, used for the callbacks
+        let function = CompiledFunctionKind::DynamicFunctionTrampoline(ty.clone());
         let module = self.ctx.create_module("");
         let target_triple = config.target_triple();
         let target_machine = config.target_machine();
@@ -149,8 +149,9 @@ impl FuncTrampoline {
             .set_section(FUNCTION_SECTION);
         generate_dynamic_trampoline(trampoline_func, ty, &self.ctx, &intrinsics)?;
 
-        // TODO: remove debugging
-        //module.print_to_stderr();
+        if let Some(ref callbacks) = config.callbacks {
+            callbacks.preopt_ir(&function, &module);
+        }
 
         let pass_manager = PassManager::create(());
 
@@ -162,22 +163,17 @@ impl FuncTrampoline {
 
         pass_manager.run_on(&module);
 
-        // TODO: remove debugging
-        //module.print_to_stderr();
+        if let Some(ref callbacks) = config.callbacks {
+            callbacks.postopt_ir(&function, &module);
+        }
 
         let memory_buffer = target_machine
             .write_to_memory_buffer(&module, FileType::Object)
             .unwrap();
 
-        /*
-        // TODO: remove debugging
-        let mem_buf_slice = memory_buffer.as_slice();
-        let mut file = fs::File::create("/home/nicholas/trampoline.o").unwrap();
-        let mut pos = 0;
-        while pos < mem_buf_slice.len() {
-            pos += file.write(&mem_buf_slice[pos..]).unwrap();
+        if let Some(ref callbacks) = config.callbacks {
+            callbacks.obj_memory_buffer(&function, &memory_buffer);
         }
-        */
 
         let mem_buf_slice = memory_buffer.as_slice();
         let (function, sections) =

lib/compiler-llvm/src/translator/code.rs

@@ -23,7 +23,7 @@ use inkwell::{
 };
 use smallvec::SmallVec;
 
-use crate::config::LLVMConfig;
+use crate::config::{CompiledFunctionKind, LLVMConfig};
 use crate::object_file::load_object_file;
 use wasm_common::entity::{PrimaryMap, SecondaryMap};
 use wasm_common::{
@@ -89,6 +89,8 @@ impl FuncTranslator {
         _table_plans: &PrimaryMap<TableIndex, TablePlan>,
         func_names: &SecondaryMap<FunctionIndex, String>,
     ) -> Result<(CompiledFunction, CustomSections), CompileError> {
+        // The function type, used for the callbacks
+        let function = CompiledFunctionKind::Local(local_func_index.clone());
         let func_index = wasm_module.func_index(*local_func_index);
         let func_name = &func_names[func_index];
         let module_name = match wasm_module.name.as_ref() {
@@ -217,10 +219,10 @@ impl FuncTranslator {
             }
         }
 
-        // TODO: debugging
-        //module.print_to_stderr();
+        if let Some(ref callbacks) = config.callbacks {
+            callbacks.preopt_ir(&function, &module);
+        }
 
-        // TODO: llvm-callbacks pre-opt-ir
         let pass_manager = PassManager::create(());
 
         if config.enable_verifier {
@@ -260,21 +262,17 @@ impl FuncTranslator {
 
         pass_manager.run_on(&module);
 
-        // TODO: llvm-callbacks llvm post-opt-ir
+        if let Some(ref callbacks) = config.callbacks {
+            callbacks.postopt_ir(&function, &module);
+        }
 
         let memory_buffer = target_machine
             .write_to_memory_buffer(&module, FileType::Object)
             .unwrap();
 
-        // TODO: remove debugging.
-        /*
-        let mem_buf_slice = memory_buffer.as_slice();
-        let mut file = std::fs::File::create(format!("/home/nicholas/code{}.o", func_name)).unwrap();
-        let mut pos = 0;
-        while pos < mem_buf_slice.len() {
-            pos += file.write(&mem_buf_slice[pos..]).unwrap();
+        if let Some(ref callbacks) = config.callbacks {
+            callbacks.obj_memory_buffer(&function, &memory_buffer);
         }
-        */
 
         let mem_buf_slice = memory_buffer.as_slice();
         load_object_file(