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(
-        &mut self,
-        _offset: binemit::CodeOffset,
-        _reloc: binemit::Reloc,
-        _jt: ir::JumpTable,
-    ) {
-        panic!("jump tables not yet implemented");
+    fn reloc_jt(&mut self, offset: binemit::CodeOffset, reloc: binemit::Reloc, jt: ir::JumpTable) {
+        self.func_relocs.push(Relocation {
+            reloc,
+            reloc_target: RelocationTarget::JumpTable(self.func_index, jt),
+            offset,
+            addend: 0,
+        });
     }
 }
 
 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)| {
-                functions.push(function);
+            .for_each(|(function, func_jt_offsets, relocs, address_transform)| {
+                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(),
         ))