Integrate Lightbeam with latest Wasmtime master (#1232)

* Implement trap info in Lightbeam

* Start using wasm-reader instead of wasmparser for parsing operators

* Update to use wasm-reader, some reductions in allocation, support source location tracking for traps, start to support multi-value

The only thing that still needs to be supported for multi-value is stack returns, but we need to make it compatible with Cranelift.

* Error when running out of registers (although we'd hope it should be impossible) instead of panicking

* WIP: Update Lightbeam to work with latest Wasmtime

* WIP: Update Lightbeam to use current wasmtime

* WIP: Migrate to new system for builtin functions

* WIP: Update Lightbeam to work with latest Wasmtime

* Remove multi_mut

* Format

* Fix some bugs around arguments, add debuginfo offset tracking

* Complete integration with new Wasmtime

* Remove commented code

* Fix formatting

* Fix warnings, remove unused dependencies

* Fix `iter` if there are too many elements, fix compilation for latest wasmtime

* Fix float arguments on stack

* Remove wasm-reader and trap info work

* Allocate stack space _before_ passing arguments, fail if we can't zero a xmm reg

* Fix stack argument offset calculation

* Fix stack arguments in Lightbeam

* Re-add WASI because it somehow got removed during rebase

* Workaround for apparent `type_alias_impl_trait`-related bug in rustdoc

* Fix breakages caused by rebase, remove module offset info as it is unrelated to wasmtime integration PR and was broken by rebase

* Add TODO comment explaining `lightbeam::ModuleContext` trait

crates/environ/src/cranelift.rs

@@ -175,13 +175,17 @@ impl RelocSink {
     }
 }
 
-struct TrapSink {
+/// Implementation of a trap sink that simply stores all trap info in-memory
+#[derive(Default)]
+pub struct TrapSink {
+    /// The in-memory vector of trap info
     pub traps: Vec<TrapInformation>,
 }
 
 impl TrapSink {
-    fn new() -> Self {
-        Self { traps: Vec::new() }
+    /// Create a new `TrapSink`
+    pub fn new() -> Self {
+        Self::default()
     }
 }
 

crates/environ/src/func_environ.rs

@@ -85,6 +85,11 @@ impl BuiltinFunctionIndex {
         13
     }
 
+    /// Create a new `BuiltinFunctionIndex` from its index
+    pub const fn from_u32(i: u32) -> Self {
+        Self(i)
+    }
+
     /// Return the index as an u32 number.
     pub const fn index(&self) -> u32 {
         self.0
@@ -509,6 +514,9 @@ impl<'module_environment> FuncEnvironment<'module_environment> {
     }
 }
 
+// TODO: This is necessary as if Lightbeam used `FuncEnvironment` directly it would cause
+//       a circular dependency graph. We should extract common types out into a separate
+//       crate that Lightbeam can use but until then we need this trait.
 #[cfg(feature = "lightbeam")]
 impl lightbeam::ModuleContext for FuncEnvironment<'_> {
     type Signature = ir::Signature;
@@ -556,6 +564,11 @@ impl lightbeam::ModuleContext for FuncEnvironment<'_> {
             .map(DefinedMemoryIndex::as_u32)
     }
 
+    fn vmctx_builtin_function(&self, func_index: u32) -> u32 {
+        self.offsets
+            .vmctx_builtin_function(BuiltinFunctionIndex::from_u32(func_index))
+    }
+
     fn vmctx_vmfunction_import_body(&self, func_index: u32) -> u32 {
         self.offsets
             .vmctx_vmfunction_import_body(FuncIndex::from_u32(func_index))

crates/environ/src/lightbeam.rs

@@ -1,15 +1,16 @@
 //! Support for compiling with Lightbeam.
 
 use crate::cache::ModuleCacheDataTupleType;
-use crate::compilation::{Compilation, CompileError, Traps};
+use crate::compilation::{Compilation, CompileError};
 use crate::func_environ::FuncEnvironment;
+use crate::CacheConfig;
 use crate::ModuleTranslation;
 // TODO: Put this in `compilation`
 use crate::address_map::{ModuleAddressMap, ValueLabelsRanges};
-use crate::cranelift::RelocSink;
-use crate::CacheConfig;
+use crate::cranelift::{RelocSink, TrapSink};
 use cranelift_codegen::isa;
 use cranelift_entity::{PrimaryMap, SecondaryMap};
+use lightbeam::{CodeGenSession, NullOffsetSink, Sinks};
 
 /// A compiler that compiles a WebAssembly module with Lightbeam, directly translating the Wasm file.
 pub struct Lightbeam;
@@ -31,22 +32,34 @@ impl crate::compilation::Compiler for Lightbeam {
             &translation.module.local,
             &translation.tunables,
         );
-        let mut relocations = PrimaryMap::new();
-        let mut codegen_session: lightbeam::CodeGenSession<_> =
-            lightbeam::CodeGenSession::new(translation.function_body_inputs.len() as u32, &env);
+        let mut relocations = PrimaryMap::with_capacity(translation.function_body_inputs.len());
+        let mut traps = PrimaryMap::with_capacity(translation.function_body_inputs.len());
+
+        let mut codegen_session: CodeGenSession<_> = CodeGenSession::new(
+            translation.function_body_inputs.len() as u32,
+            &env,
+            lightbeam::microwasm::I32,
+        );
 
         for (i, function_body) in &translation.function_body_inputs {
             let func_index = translation.module.local.func_index(i);
-            let mut reloc_sink = RelocSink::new(func_index);
 
+            let mut reloc_sink = RelocSink::new(func_index);
+            let mut trap_sink = TrapSink::new();
             lightbeam::translate_function(
                 &mut codegen_session,
-                &mut reloc_sink,
+                Sinks {
+                    relocs: &mut reloc_sink,
+                    traps: &mut trap_sink,
+                    offsets: &mut NullOffsetSink,
+                },
                 i.as_u32(),
-                &wasmparser::FunctionBody::new(0, function_body.data),
+                wasmparser::FunctionBody::new(0, function_body.data),
             )
             .map_err(|e| CompileError::Codegen(format!("Failed to translate function: {}", e)))?;
+
             relocations.push(reloc_sink.func_relocs);
+            traps.push(trap_sink.traps);
         }
 
         let code_section = codegen_session
@@ -67,7 +80,7 @@ impl crate::compilation::Compiler for Lightbeam {
             ModuleAddressMap::new(),
             ValueLabelsRanges::new(),
             PrimaryMap::new(),
-            Traps::new(),
+            traps,
         ))
     }
 }

crates/lightbeam/Cargo.toml

@@ -11,19 +11,20 @@ keywords = ["webassembly", "wasm", "compile", "compiler", "jit"]
 edition = "2018"
 
 [dependencies]
-smallvec = "1.0.0"
+capstone = "0.6.0"
+cranelift-codegen = { path = "../../cranelift/codegen", version = "0.62.0" }
+derive_more = "0.99"
 dynasm = "0.5.2"
 dynasmrt = "0.5.2"
-wasmparser = "0.51.2"
-memoffset = "0.5.3"
+iter-enum = "0.2"
 itertools = "0.8.2"
-capstone = "0.6.0"
-thiserror = "1.0.9"
-cranelift-codegen = { path = "../../cranelift/codegen", version = "0.62.0" }
-multi_mut = "0.1"
-either = "1.5"
-typemap = "0.3"
+memoffset = "0.5.3"
 more-asserts = "0.2.1"
+smallvec = "1.0.0"
+staticvec = "0.8"
+thiserror = "1.0.9"
+typemap = "0.3"
+wasmparser = "0.51.2"
 
 [dev-dependencies]
 lazy_static = "1.2"

crates/lightbeam/src/disassemble.rs

@@ -3,9 +3,9 @@ use dynasmrt::AssemblyOffset;
 use std::error::Error;
 use std::fmt::{Display, Write};
 
-pub fn disassemble(
+pub fn disassemble<D: Display>(
     mem: &[u8],
-    mut ops: &[(AssemblyOffset, impl Display)],
+    mut ops: &[(AssemblyOffset, D)],
 ) -> Result<(), Box<dyn Error>> {
     let cs = Capstone::new()
         .x86()

crates/lightbeam/src/lib.rs

@@ -8,7 +8,7 @@ mod backend;
 mod disassemble;
 mod error;
 mod function_body;
-mod microwasm;
+pub mod microwasm;
 mod module;
 mod translate_sections;
 
@@ -16,7 +16,10 @@ mod translate_sections;
 mod benches;
 
 pub use crate::backend::CodeGenSession;
-pub use crate::function_body::translate_wasm as translate_function;
+pub use crate::function_body::{
+    translate_wasm as translate_function, NullOffsetSink, OffsetSink, Sinks,
+};
 pub use crate::module::{
     translate, ExecutableModule, ExecutionError, ModuleContext, Signature, TranslatedModule,
 };
+pub use disassemble::disassemble;

crates/lightbeam/src/module.rs

@@ -300,8 +300,9 @@ impl Signature for CraneliftSignature {
         // TODO: We want to instead add the `VMContext` to the signature used by
         //       cranelift, removing the special-casing from the internals.
         assert_eq!(self.params[0].purpose, ir::ArgumentPurpose::VMContext);
+        // `self.params[1]` should be caller vmctx
         assert_eq!(self.call_conv, isa::CallConv::SystemV);
-        &self.params[1..]
+        &self.params[2..]
     }
 
     fn returns(&self) -> &[Self::Type] {
@@ -332,6 +333,7 @@ pub trait ModuleContext {
     type Signature: Signature;
     type GlobalType: SigType;
 
+    fn vmctx_builtin_function(&self, index: u32) -> u32;
     fn vmctx_vmglobal_definition(&self, index: u32) -> u32;
     fn vmctx_vmglobal_import_from(&self, index: u32) -> u32;
     fn vmctx_vmmemory_import_from(&self, memory_index: u32) -> u32;
@@ -424,6 +426,10 @@ impl ModuleContext for SimpleContext {
         Some(index)
     }
 
+    fn vmctx_builtin_function(&self, _index: u32) -> u32 {
+        unimplemented!()
+    }
+
     fn vmctx_vmfunction_import_body(&self, _func_index: u32) -> u32 {
         unimplemented!()
     }

crates/lightbeam/src/translate_sections.rs

@@ -1,6 +1,5 @@
-use crate::backend::{CodeGenSession, TranslatedCodeSection};
+use crate::backend::TranslatedCodeSection;
 use crate::error::Error;
-use crate::function_body;
 use crate::module::SimpleContext;
 use cranelift_codegen::{binemit, ir};
 use wasmparser::{
@@ -106,20 +105,13 @@ impl binemit::RelocSink for UnimplementedRelocSink {
 
 /// Parses the Code section of the wasm module.
 pub fn code(
-    code: CodeSectionReader,
-    translation_ctx: &SimpleContext,
+    _code: CodeSectionReader,
+    _translation_ctx: &SimpleContext,
 ) -> Result<TranslatedCodeSection, Error> {
-    let func_count = code.get_count();
-    let mut session = CodeGenSession::new(func_count, translation_ctx);
-
-    for (idx, body) in code.into_iter().enumerate() {
-        let body = body?;
-        let mut relocs = UnimplementedRelocSink;
-
-        function_body::translate_wasm(&mut session, &mut relocs, idx as u32, &body)?;
-    }
-
-    Ok(session.into_translated_code_section()?)
+    // TODO: Remove the Lightbeam harness entirely, this is just to make this compile.
+    //       We do all our testing through Wasmtime now, there's no reason to duplicate
+    //       writing a WebAssembly environment in Lightbeam too.
+    unimplemented!("Incomplete migration to wasm-reader");
 }
 
 /// Parses the Data section of the wasm module.