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

Cargo.lock

@@ -106,9 +106,9 @@ dependencies = [
 
 [[package]]
 name = "backtrace-sys"
-version = "0.1.35"
+version = "0.1.36"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "7de8aba10a69c8e8d7622c5710229485ec32e9d55fdad160ea559c086fdcd118"
+checksum = "78848718ee1255a2485d1309ad9cdecfc2e7d0362dd11c6829364c6b35ae1bc7"
 dependencies = [
  "cc",
  "libc",
@@ -242,9 +242,9 @@ dependencies = [
 
 [[package]]
 name = "cc"
-version = "1.0.50"
+version = "1.0.52"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "95e28fa049fda1c330bcf9d723be7663a899c4679724b34c81e9f5a326aab8cd"
+checksum = "c3d87b23d6a92cd03af510a5ade527033f6aa6fa92161e2d5863a907d4c5e31d"
 dependencies = [
  "jobserver",
 ]
@@ -648,6 +648,28 @@ dependencies = [
  "syn",
 ]
 
+[[package]]
+name = "derive_more"
+version = "0.99.5"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "e2323f3f47db9a0e77ce7a300605d8d2098597fc451ed1a97bb1f6411bb550a7"
+dependencies = [
+ "proc-macro2",
+ "quote",
+ "syn",
+]
+
+[[package]]
+name = "derive_utils"
+version = "0.9.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "114aa287358087a616096186f3de19525d8083f83d437dc6b583f895316b02e6"
+dependencies = [
+ "proc-macro2",
+ "quote",
+ "syn",
+]
+
 [[package]]
 name = "diff"
 version = "0.1.12"
@@ -958,9 +980,9 @@ dependencies = [
 
 [[package]]
 name = "hermit-abi"
-version = "0.1.10"
+version = "0.1.11"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "725cf19794cf90aa94e65050cb4191ff5d8fa87a498383774c47b332e3af952e"
+checksum = "8a0d737e0f947a1864e93d33fdef4af8445a00d1ed8dc0c8ddb73139ea6abf15"
 dependencies = [
  "libc",
 ]
@@ -995,6 +1017,17 @@ dependencies = [
  "regex",
 ]
 
+[[package]]
+name = "iter-enum"
+version = "0.2.3"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "58107b833f974ca7bc9f98e032881967613dc214b6ab7ececb45b1ab9f0e7008"
+dependencies = [
+ "derive_utils",
+ "quote",
+ "syn",
+]
+
 [[package]]
 name = "itertools"
 version = "0.8.2"
@@ -1042,9 +1075,9 @@ checksum = "3576a87f2ba00f6f106fdfcd16db1d698d648a26ad8e0573cad8537c3c362d2a"
 
 [[package]]
 name = "libc"
-version = "0.2.68"
+version = "0.2.69"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "dea0c0405123bba743ee3f91f49b1c7cfb684eef0da0a50110f758ccf24cdff0"
+checksum = "99e85c08494b21a9054e7fe1374a732aeadaff3980b6990b94bfd3a70f690005"
 
 [[package]]
 name = "libfuzzer-sys"
@@ -1063,16 +1096,17 @@ dependencies = [
  "anyhow",
  "capstone",
  "cranelift-codegen",
+ "derive_more",
  "dynasm",
  "dynasmrt",
- "either",
+ "iter-enum",
  "itertools",
  "lazy_static",
  "memoffset",
  "more-asserts",
- "multi_mut",
  "quickcheck",
  "smallvec",
+ "staticvec",
  "thiserror",
  "typemap",
  "wasmparser",
@@ -1143,12 +1177,6 @@ version = "0.2.1"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "0debeb9fcf88823ea64d64e4a815ab1643f33127d995978e099942ce38f25238"
 
-[[package]]
-name = "multi_mut"
-version = "0.1.3"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "816df386e5557ac1843a96f1ba8a7cbf4ab175d05ccc15c87a3cda27b4fbdece"
-
 [[package]]
 name = "num-traits"
 version = "0.2.11"
@@ -1160,9 +1188,9 @@ dependencies = [
 
 [[package]]
 name = "num_cpus"
-version = "1.12.0"
+version = "1.13.0"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "46203554f085ff89c235cd12f7075f3233af9b11ed7c9e16dfe2560d03313ce6"
+checksum = "05499f3756671c15885fee9034446956fff3f243d6077b91e5767df161f766b3"
 dependencies = [
  "hermit-abi",
  "libc",
@@ -1640,9 +1668,9 @@ dependencies = [
 
 [[package]]
 name = "ryu"
-version = "1.0.3"
+version = "1.0.4"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "535622e6be132bccd223f4bb2b8ac8d53cda3c7a6394944d3b2b33fb974f9d76"
+checksum = "ed3d612bc64430efeb3f7ee6ef26d590dce0c43249217bddc62112540c7941e1"
 
 [[package]]
 name = "same-file"
@@ -1749,6 +1777,12 @@ version = "1.1.1"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "dba1a27d3efae4351c8051072d619e3ade2820635c3958d826bfea39d59b54c8"
 
+[[package]]
+name = "staticvec"
+version = "0.8.9"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "ad858b0675d898276da104afc602fd7f42790123d3150646dac32338a92132af"
+
 [[package]]
 name = "string-interner"
 version = "0.7.1"
@@ -1766,9 +1800,9 @@ checksum = "8ea5119cdb4c55b55d432abb513a0429384878c15dde60cc77b1c99de1a95a6a"
 
 [[package]]
 name = "structopt"
-version = "0.3.13"
+version = "0.3.14"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "ff6da2e8d107dfd7b74df5ef4d205c6aebee0706c647f6bc6a2d5789905c00fb"
+checksum = "863246aaf5ddd0d6928dfeb1a9ca65f505599e4e1b399935ef7e75107516b4ef"
 dependencies = [
  "clap",
  "lazy_static",
@@ -1777,9 +1811,9 @@ dependencies = [
 
 [[package]]
 name = "structopt-derive"
-version = "0.4.6"
+version = "0.4.7"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "a489c87c08fbaf12e386665109dd13470dcc9c4583ea3e10dd2b4523e5ebd9ac"
+checksum = "d239ca4b13aee7a2142e6795cbd69e457665ff8037aed33b3effdc430d2f927a"
 dependencies = [
  "heck",
  "proc-macro-error",
@@ -1938,9 +1972,9 @@ dependencies = [
 
 [[package]]
 name = "typenum"
-version = "1.11.2"
+version = "1.12.0"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "6d2783fe2d6b8c1101136184eb41be8b1ad379e4657050b8aaff0c79ee7575f9"
+checksum = "373c8a200f9e67a0c95e62a4f52fbf80c23b4381c05a17845531982fa99e6b33"
 
 [[package]]
 name = "unicode-segmentation"
@@ -2400,9 +2434,9 @@ checksum = "ac3b87c63620426dd9b991e5ce0329eff545bccbbb34f3be09ff6fb6ab51b7b6"
 
 [[package]]
 name = "winapi-util"
-version = "0.1.4"
+version = "0.1.5"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "fa515c5163a99cc82bab70fd3bfdd36d827be85de63737b40fcef2ce084a436e"
+checksum = "70ec6ce85bb158151cae5e5c87f95a8e97d2c0c4b001223f33a334e3ce5de178"
 dependencies = [
  "winapi",
 ]

build.rs

@@ -176,8 +176,6 @@ fn ignore(testsuite: &str, testname: &str, strategy: &str) -> bool {
             ("multi_value", _) => return true,
             ("reference_types", _) => return true,
             ("bulk_memory_operations", _) => return true,
-            // Lightbeam doesn't support float arguments on the stack.
-            ("spec_testsuite", "call") => return true,
             _ => (),
         },
         "Cranelift" => match (testsuite, testname) {

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 {
+    /// Create a new `TrapSink`
-        Self { traps: Vec::new() }
+    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::cranelift::{RelocSink, TrapSink};
-use crate::CacheConfig;
 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 relocations = PrimaryMap::with_capacity(translation.function_body_inputs.len());
-        let mut codegen_session: lightbeam::CodeGenSession<_> =
+        let mut traps = PrimaryMap::with_capacity(translation.function_body_inputs.len());
-            lightbeam::CodeGenSession::new(translation.function_body_inputs.len() as u32, &env);
+
+        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"
+iter-enum = "0.2"
-memoffset = "0.5.3"
 itertools = "0.8.2"
-capstone = "0.6.0"
+memoffset = "0.5.3"
-thiserror = "1.0.9"
-cranelift-codegen = { path = "../../cranelift/codegen", version = "0.62.0" }
-multi_mut = "0.1"
-either = "1.5"
-typemap = "0.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,
+    _code: CodeSectionReader,
-    translation_ctx: &SimpleContext,
+    _translation_ctx: &SimpleContext,
 ) -> Result<TranslatedCodeSection, Error> {
-    let func_count = code.get_count();
+    // TODO: Remove the Lightbeam harness entirely, this is just to make this compile.
-    let mut session = CodeGenSession::new(func_count, translation_ctx);
+    //       We do all our testing through Wasmtime now, there's no reason to duplicate
-
+    //       writing a WebAssembly environment in Lightbeam too.
-    for (idx, body) in code.into_iter().enumerate() {
+    unimplemented!("Incomplete migration to wasm-reader");
-        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()?)
 }
 
 /// Parses the Data section of the wasm module.