Allow jump tables in wasmtime.

wasmtime-environ/src/compilation.rs

@@ -9,7 +9,17 @@ use cranelift_wasm::{DefinedFuncIndex, FuncIndex, WasmError};
 use std::ops::Range;
 use std::vec::Vec;
 
-type Functions = PrimaryMap<DefinedFuncIndex, Vec<u8>>;
+/// Compiled machine code: body and jump table offsets.
+#[derive(Debug, Clone)]
+pub struct CodeAndJTOffsets {
+    /// The function body.
+    pub body: Vec<u8>,
+
+    /// The jump tables offsets (in the body).
+    pub jt_offsets: ir::JumpTableOffsets,
+}
+
+type Functions = PrimaryMap<DefinedFuncIndex, CodeAndJTOffsets>;
 
 /// The result of compiling a WebAssembly module's functions.
 #[derive(Debug)]
@@ -25,17 +35,23 @@ impl Compilation {
     }
 
     /// Allocates the compilation result with the given function bodies.
-    pub fn from_buffer(buffer: Vec<u8>, functions: impl IntoIterator<Item = Range<usize>>) -> Self {
+    pub fn from_buffer(
+        buffer: Vec<u8>,
+        functions: impl IntoIterator<Item = (Range<usize>, ir::JumpTableOffsets)>,
+    ) -> Self {
         Self::new(
             functions
                 .into_iter()
-                .map(|range| buffer[range].to_vec())
+                .map(|(range, jt_offsets)| CodeAndJTOffsets {
+                    body: buffer[range].to_vec(),
+                    jt_offsets,
+                })
                 .collect(),
         )
     }
 
     /// Gets the bytes of a single function
-    pub fn get(&self, func: DefinedFuncIndex) -> &[u8] {
+    pub fn get(&self, func: DefinedFuncIndex) -> &CodeAndJTOffsets {
         &self.functions[func]
     }
 
@@ -43,6 +59,14 @@ impl Compilation {
     pub fn len(&self) -> usize {
         self.functions.len()
     }
+
+    /// Gets functions jump table offsets.
+    pub fn get_jt_offsets(&self) -> PrimaryMap<DefinedFuncIndex, ir::JumpTableOffsets> {
+        self.functions
+            .iter()
+            .map(|(_, code_and_jt)| code_and_jt.jt_offsets.clone())
+            .collect::<PrimaryMap<DefinedFuncIndex, _>>()
+    }
 }
 
 impl<'a> IntoIterator for &'a Compilation {
@@ -61,10 +85,10 @@ pub struct Iter<'a> {
 }
 
 impl<'a> Iterator for Iter<'a> {
-    type Item = &'a [u8];
+    type Item = &'a CodeAndJTOffsets;
 
     fn next(&mut self) -> Option<Self::Item> {
-        self.iterator.next().map(|(_, b)| &b[..])
+        self.iterator.next().map(|(_, b)| b)
     }
 }
 
@@ -96,6 +120,8 @@ pub enum RelocationTarget {
     Memory32Size,
     /// Function for query current size of an imported 32-bit linear memory.
     ImportedMemory32Size,
+    /// Jump table index.
+    JumpTable(FuncIndex, ir::JumpTable),
 }
 
 /// Relocations to apply to function bodies.

wasmtime-environ/src/cranelift.rs

@@ -1,7 +1,7 @@
 //! Support for compiling with Cranelift.
 
 use crate::compilation::{
-    AddressTransforms, Compilation, CompileError, FunctionAddressTransform,
+    AddressTransforms, CodeAndJTOffsets, Compilation, CompileError, FunctionAddressTransform,
     InstructionAddressTransform, Relocation, RelocationTarget, Relocations,
 };
 use crate::func_environ::{
@@ -22,6 +22,9 @@ use std::vec::Vec;
 
 /// Implementation of a relocation sink that just saves all the information for later
 pub struct RelocSink {
+    /// Current function index.
+    func_index: FuncIndex,
+
     /// Relocations recorded for the function.
     pub func_relocs: Vec<Relocation>,
 }
@@ -66,20 +69,21 @@ impl binemit::RelocSink for RelocSink {
             addend,
         });
     }
-    fn reloc_jt(
+    fn reloc_jt(&mut self, offset: binemit::CodeOffset, reloc: binemit::Reloc, jt: ir::JumpTable) {
-        &mut self,
+        self.func_relocs.push(Relocation {
-        _offset: binemit::CodeOffset,
+            reloc,
-        _reloc: binemit::Reloc,
+            reloc_target: RelocationTarget::JumpTable(self.func_index, jt),
-        _jt: ir::JumpTable,
+            offset,
-    ) {
+            addend: 0,
-        panic!("jump tables not yet implemented");
+        });
     }
 }
 
 impl RelocSink {
     /// Return a new `RelocSink` instance.
-    pub fn new() -> Self {
+    pub fn new(func_index: FuncIndex) -> Self {
         Self {
+            func_index,
             func_relocs: Vec::new(),
         }
     }
@@ -147,12 +151,14 @@ impl crate::compilation::Compiler for Cranelift {
                     .map_err(CompileError::Wasm)?;
 
                 let mut code_buf: Vec<u8> = Vec::new();
-                let mut reloc_sink = RelocSink::new();
+                let mut reloc_sink = RelocSink::new(func_index);
                 let mut trap_sink = binemit::NullTrapSink {};
                 context
                     .compile_and_emit(isa, &mut code_buf, &mut reloc_sink, &mut trap_sink)
                     .map_err(CompileError::Codegen)?;
 
+                let jt_offsets = context.func.jt_offsets.clone();
+
                 let address_transform = if generate_debug_info {
                     let body_len = code_buf.len();
                     let at = get_address_transform(&context, isa);
@@ -165,12 +171,20 @@ impl crate::compilation::Compiler for Cranelift {
                     None
                 };
 
-                Ok((code_buf, reloc_sink.func_relocs, address_transform))
+                Ok((
+                    code_buf,
+                    jt_offsets,
+                    reloc_sink.func_relocs,
+                    address_transform,
+                ))
             })
             .collect::<Result<Vec<_>, CompileError>>()?
             .into_iter()
-            .for_each(|(function, relocs, address_transform)| {
+            .for_each(|(function, func_jt_offsets, relocs, address_transform)| {
-                functions.push(function);
+                functions.push(CodeAndJTOffsets {
+                    body: function,
+                    jt_offsets: func_jt_offsets,
+                });
                 relocations.push(relocs);
                 if let Some(address_transform) = address_transform {
                     address_transforms.push(address_transform);

wasmtime-environ/src/lightbeam.rs

@@ -7,7 +7,7 @@ use crate::module_environ::FunctionBodyData;
 // TODO: Put this in `compilation`
 use crate::cranelift::RelocSink;
 use cranelift_codegen::isa;
-use cranelift_entity::PrimaryMap;
+use cranelift_entity::{PrimaryMap, SecondaryMap};
 use cranelift_wasm::DefinedFuncIndex;
 use lightbeam;
 
@@ -30,7 +30,8 @@ impl crate::compilation::Compiler for Lightbeam {
             lightbeam::CodeGenSession::new(function_body_inputs.len() as u32, &env);
 
         for (i, function_body) in &function_body_inputs {
-            let mut reloc_sink = RelocSink::new();
+            let func_index = module.func_index(i);
+            let mut reloc_sink = RelocSink::new(func_index);
 
             lightbeam::translate_function(
                 &mut codegen_session,
@@ -46,8 +47,15 @@ impl crate::compilation::Compiler for Lightbeam {
             .into_translated_code_section()
             .expect("Failed to generate output code. TODO: Stop this from panicking");
 
+        // TODO pass jump table offsets to Compilation::from_buffer() when they
+        // are implemented in lightbeam -- using empty set of offsets for now.
+        let code_section_ranges_and_jt = code_section
+            .funcs()
+            .into_iter()
+            .map(|r| (r, SecondaryMap::new()));
+
         Ok((
-            Compilation::from_buffer(code_section.buffer().to_vec(), code_section.funcs()),
+            Compilation::from_buffer(code_section.buffer().to_vec(), code_section_ranges_and_jt),
             relocations,
             AddressTransforms::new(),
         ))

wasmtime-jit/src/compiler.rs

@@ -77,6 +77,7 @@ impl Compiler {
     ) -> Result<
         (
             PrimaryMap<DefinedFuncIndex, *mut [VMFunctionBody]>,
+            PrimaryMap<DefinedFuncIndex, ir::JumpTableOffsets>,
             Relocations,
             Option<Vec<u8>>,
         ),
@@ -104,7 +105,7 @@ impl Compiler {
             let mut funcs = Vec::new();
             for (i, allocated) in allocated_functions.into_iter() {
                 let ptr = (*allocated) as *const u8;
-                let body_len = compilation.get(i).len();
+                let body_len = compilation.get(i).body.len();
                 funcs.push((ptr, body_len));
             }
             let bytes = emit_debugsections_image(
@@ -120,7 +121,9 @@ impl Compiler {
             None
         };
 
-        Ok((allocated_functions, relocations, dbg))
+        let jt_offsets = compilation.get_jt_offsets();
+
+        Ok((allocated_functions, jt_offsets, relocations, dbg))
     }
 
     /// Create a trampoline for invoking a function.
@@ -259,7 +262,10 @@ fn allocate_functions(
     // Allocate code for all function in one continuous memory block.
     // First, collect all function bodies into vector to pass to the
     // allocate_copy_of_byte_slices.
-    let bodies = compilation.into_iter().collect::<Vec<&[u8]>>();
+    let bodies = compilation
+        .into_iter()
+        .map(|code_and_jt| &code_and_jt.body[..])
+        .collect::<Vec<&[u8]>>();
     let fat_ptrs = code_memory.allocate_copy_of_byte_slices(&bodies)?;
     // Second, create a PrimaryMap from result vector of pointers.
     let mut result = PrimaryMap::with_capacity(compilation.len());

wasmtime-jit/src/instantiate.rs

@@ -77,7 +77,7 @@ impl<'data> RawCompiledModule<'data> {
             None
         };
 
-        let (allocated_functions, relocations, dbg_image) = compiler.compile(
+        let (allocated_functions, jt_offsets, relocations, dbg_image) = compiler.compile(
             &translation.module,
             translation.function_body_inputs,
             debug_data,
@@ -86,6 +86,7 @@ impl<'data> RawCompiledModule<'data> {
         let imports = link_module(
             &translation.module,
             &allocated_functions,
+            &jt_offsets,
             relocations,
             resolver,
         )

wasmtime-jit/src/link.rs

@@ -3,6 +3,7 @@
 use crate::resolver::Resolver;
 use core::ptr::write_unaligned;
 use cranelift_codegen::binemit::Reloc;
+use cranelift_codegen::ir::JumpTableOffsets;
 use cranelift_entity::PrimaryMap;
 use cranelift_wasm::{DefinedFuncIndex, Global, GlobalInit, Memory, Table, TableElementType};
 use std::collections::HashSet;
@@ -20,6 +21,7 @@ use wasmtime_runtime::{
 pub fn link_module(
     module: &Module,
     allocated_functions: &PrimaryMap<DefinedFuncIndex, *mut [VMFunctionBody]>,
+    jt_offsets: &PrimaryMap<DefinedFuncIndex, JumpTableOffsets>,
     relocations: Relocations,
     resolver: &mut dyn Resolver,
 ) -> Result<Imports, LinkError> {
@@ -194,7 +196,7 @@ pub fn link_module(
     }
 
     // Apply relocations, now that we have virtual addresses for everything.
-    relocate(allocated_functions, relocations, module);
+    relocate(allocated_functions, jt_offsets, relocations, module);
 
     Ok(Imports::new(
         dependencies,
@@ -299,6 +301,7 @@ fn is_memory_compatible(exported: &MemoryPlan, imported: &MemoryPlan) -> bool {
 /// Performs the relocations inside the function bytecode, provided the necessary metadata.
 fn relocate(
     allocated_functions: &PrimaryMap<DefinedFuncIndex, *mut [VMFunctionBody]>,
+    jt_offsets: &PrimaryMap<DefinedFuncIndex, JumpTableOffsets>,
     relocations: PrimaryMap<DefinedFuncIndex, Vec<Relocation>>,
     module: &Module,
 ) {
@@ -335,6 +338,19 @@ fn relocate(
                         other => panic!("unexpected libcall: {}", other),
                     }
                 }
+                RelocationTarget::JumpTable(func_index, jt) => {
+                    match module.defined_func_index(func_index) {
+                        Some(f) => {
+                            let offset = *jt_offsets
+                                .get(f)
+                                .and_then(|ofs| ofs.get(jt))
+                                .expect("func jump table");
+                            let fatptr: *const [VMFunctionBody] = allocated_functions[f];
+                            fatptr as *const VMFunctionBody as usize + offset as usize
+                        }
+                        None => panic!("func index of jump table"),
+                    }
+                }
             };
 
             let fatptr: *const [VMFunctionBody] = allocated_functions[i];
@@ -363,6 +379,9 @@ fn relocate(
                 Reloc::X86CallPCRel4 => {
                     // ignore
                 }
+                Reloc::X86PCRelRodata4 => {
+                    // ignore
+                }
                 _ => panic!("unsupported reloc kind"),
             }
         }

wasmtime-obj/src/function.rs

@@ -62,11 +62,11 @@ pub fn emit_functions(
         .expect("Missing enable_verifier setting");
 
     for (i, _function_relocs) in relocations.iter() {
-        let body = compilation.get(i);
+        let body = &compilation.get(i).body;
         let func_index = module.func_index(i);
         let string_name = format!("_wasm_function_{}", func_index.index());
 
-        obj.define(string_name, body.to_vec())
+        obj.define(string_name, body.clone())
             .map_err(|err| format!("{}", err))?;
     }