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