Allow jump tables in wasmtime.

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"),
             }
         }