cranelift-module: Add support for passing a StackMapSink when defining functions

Fixes #2738

This follows the convention set by the existing method of passing a
TrapSink by adding another argument for a StackMapSink.

cranelift/jit/examples/jit-minimal.rs

@@ -1,5 +1,5 @@
 use cranelift::prelude::*;
-use cranelift_codegen::binemit::NullTrapSink;
+use cranelift_codegen::binemit::{NullStackMapSink, NullTrapSink};
 use cranelift_codegen::settings::{self, Configurable};
 use cranelift_jit::{JITBuilder, JITModule};
 use cranelift_module::{default_libcall_names, Linkage, Module};
@@ -49,8 +49,9 @@ fn main() {
         bcx.finalize();
     }
     let mut trap_sink = NullTrapSink {};
+    let mut stack_map_sink = NullStackMapSink {};
     module
-        .define_function(func_a, &mut ctx, &mut trap_sink)
+        .define_function(func_a, &mut ctx, &mut trap_sink, &mut stack_map_sink)
         .unwrap();
     module.clear_context(&mut ctx);
 
@@ -74,7 +75,7 @@ fn main() {
         bcx.finalize();
     }
     module
-        .define_function(func_b, &mut ctx, &mut trap_sink)
+        .define_function(func_b, &mut ctx, &mut trap_sink, &mut stack_map_sink)
         .unwrap();
     module.clear_context(&mut ctx);
 

cranelift/jit/src/backend.rs

@@ -5,7 +5,7 @@ use cranelift_codegen::isa::TargetIsa;
 use cranelift_codegen::settings::Configurable;
 use cranelift_codegen::{self, ir, settings};
 use cranelift_codegen::{
-    binemit::{self, Addend, CodeInfo, CodeOffset, Reloc, RelocSink, TrapSink},
+    binemit::{Addend, CodeInfo, CodeOffset, Reloc, RelocSink, StackMapSink, TrapSink},
     CodegenError,
 };
 use cranelift_entity::SecondaryMap;
@@ -601,6 +601,7 @@ impl Module for JITModule {
         id: FuncId,
         ctx: &mut cranelift_codegen::Context,
         trap_sink: &mut dyn TrapSink,
+        stack_map_sink: &mut dyn StackMapSink,
     ) -> ModuleResult<ModuleCompiledFunction> {
         info!("defining function {}: {}", id, ctx.func.display(self.isa()));
         let CodeInfo {
@@ -625,16 +626,7 @@ impl Module for JITModule {
             .expect("TODO: handle OOM etc.");
 
         let mut reloc_sink = JITRelocSink::default();
-        let mut stack_map_sink = binemit::NullStackMapSink {};
+        unsafe { ctx.emit_to_memory(&*self.isa, ptr, &mut reloc_sink, trap_sink, stack_map_sink) };
-        unsafe {
-            ctx.emit_to_memory(
-                &*self.isa,
-                ptr,
-                &mut reloc_sink,
-                trap_sink,
-                &mut stack_map_sink,
-            )
-        };
 
         self.record_function_for_perf(ptr, size, &decl.name);
         self.compiled_functions[id] = Some(CompiledBlob {

cranelift/jit/tests/basic.rs

@@ -1,4 +1,4 @@
-use cranelift_codegen::binemit::NullTrapSink;
+use cranelift_codegen::binemit::{NullStackMapSink, NullTrapSink};
 use cranelift_codegen::ir::*;
 use cranelift_codegen::isa::CallConv;
 use cranelift_codegen::settings::{self, Configurable};
@@ -57,8 +57,9 @@ fn define_simple_function(module: &mut JITModule) -> FuncId {
     }
 
     let mut trap_sink = NullTrapSink {};
+    let mut stack_map_sink = NullStackMapSink {};
     module
-        .define_function(func_id, &mut ctx, &mut trap_sink)
+        .define_function(func_id, &mut ctx, &mut trap_sink, &mut stack_map_sink)
         .unwrap();
 
     func_id
@@ -205,8 +206,9 @@ fn libcall_function() {
     }
 
     let mut trap_sink = NullTrapSink {};
+    let mut stack_map_sink = NullStackMapSink {};
     module
-        .define_function(func_id, &mut ctx, &mut trap_sink)
+        .define_function(func_id, &mut ctx, &mut trap_sink, &mut stack_map_sink)
         .unwrap();
 
     module.finalize_definitions();

cranelift/module/src/module.rs

@@ -469,6 +469,7 @@ pub trait Module {
         func: FuncId,
         ctx: &mut Context,
         trap_sink: &mut dyn binemit::TrapSink,
+        stack_map_sink: &mut dyn binemit::StackMapSink,
     ) -> ModuleResult<ModuleCompiledFunction>;
 
     /// Define a function, taking the function body from the given `bytes`.
@@ -562,8 +563,9 @@ impl<M: Module> Module for &mut M {
         func: FuncId,
         ctx: &mut Context,
         trap_sink: &mut dyn binemit::TrapSink,
+        stack_map_sink: &mut dyn binemit::StackMapSink,
     ) -> ModuleResult<ModuleCompiledFunction> {
-        (**self).define_function(func, ctx, trap_sink)
+        (**self).define_function(func, ctx, trap_sink, stack_map_sink)
     }
 
     fn define_function_bytes(

cranelift/object/src/backend.rs

@@ -5,7 +5,7 @@ use cranelift_codegen::entity::SecondaryMap;
 use cranelift_codegen::isa::TargetIsa;
 use cranelift_codegen::{self, ir};
 use cranelift_codegen::{
-    binemit::{Addend, CodeInfo, CodeOffset, NullStackMapSink, Reloc, RelocSink, TrapSink},
+    binemit::{Addend, CodeInfo, CodeOffset, Reloc, RelocSink, StackMapSink, TrapSink},
     CodegenError,
 };
 use cranelift_module::{
@@ -248,6 +248,7 @@ impl Module for ObjectModule {
         func_id: FuncId,
         ctx: &mut cranelift_codegen::Context,
         trap_sink: &mut dyn TrapSink,
+        stack_map_sink: &mut dyn StackMapSink,
     ) -> ModuleResult<ModuleCompiledFunction> {
         info!(
             "defining function {}: {}",
@@ -260,7 +261,6 @@ impl Module for ObjectModule {
         } = ctx.compile(self.isa())?;
         let mut code: Vec<u8> = vec![0; code_size as usize];
         let mut reloc_sink = ObjectRelocSink::default();
-        let mut stack_map_sink = NullStackMapSink {};
 
         unsafe {
             ctx.emit_to_memory(
@@ -268,7 +268,7 @@ impl Module for ObjectModule {
                 code.as_mut_ptr(),
                 &mut reloc_sink,
                 trap_sink,
-                &mut stack_map_sink,
+                stack_map_sink,
             )
         };
 

cranelift/object/tests/basic.rs

@@ -1,6 +1,9 @@
 use cranelift_codegen::ir::*;
 use cranelift_codegen::isa::CallConv;
-use cranelift_codegen::{binemit::NullTrapSink, settings};
+use cranelift_codegen::{
+    binemit::{NullStackMapSink, NullTrapSink},
+    settings,
+};
 use cranelift_codegen::{ir::types::I16, Context};
 use cranelift_entity::EntityRef;
 use cranelift_frontend::*;
@@ -51,8 +54,9 @@ fn define_simple_function(module: &mut ObjectModule) -> FuncId {
     }
 
     let mut trap_sink = NullTrapSink {};
+    let mut stack_map_sink = NullStackMapSink {};
     module
-        .define_function(func_id, &mut ctx, &mut trap_sink)
+        .define_function(func_id, &mut ctx, &mut trap_sink, &mut stack_map_sink)
         .unwrap();
 
     func_id
@@ -191,8 +195,9 @@ fn libcall_function() {
     }
 
     let mut trap_sink = NullTrapSink {};
+    let mut stack_map_sink = NullStackMapSink {};
     module
-        .define_function(func_id, &mut ctx, &mut trap_sink)
+        .define_function(func_id, &mut ctx, &mut trap_sink, &mut stack_map_sink)
         .unwrap();
 
     module.finish();