Refactor where results of compilation are stored (#2086)

* Refactor where results of compilation are stored

This commit refactors the internals of compilation in Wasmtime to change
where results of individual function compilation are stored. Previously
compilation resulted in many maps being returned, and compilation
results generally held all these maps together. This commit instead
switches this to have all metadata stored in a `CompiledFunction`
instead of having a separate map for each item that can be stored.

The motivation for this is primarily to help out with future
module-linking-related PRs. What exactly "module level" is depends on
how we interpret modules and how many modules are in play, so it's a bit
easier for operations in wasmtime to work at the function level where
possible. This means that we don't have to pass around multiple
different maps and a function index, but instead just one map or just
one entry representing a compiled function.

Additionally this change updates where the parallelism of compilation
happens, pushing it into `wasmtime-jit` instead of `wasmtime-environ`.
This is another goal where `wasmtime-jit` will have more knowledge about
module-level pieces with module linking in play. User-facing-wise this
should be the same in terms of parallel compilation, though.

The ultimate goal of this refactoring is to make it easier for the
results of compilation to actually be a set of wasm modules. This means
we won't be able to have a map-per-metadata where the primary key is the
function index, because there will be many modules within one "object
file".

* Don't clear out fields, just don't store them

Persist a smaller set of fields in `CompilationArtifacts` instead of
trying to clear fields out and dynamically not accessing them.

crates/environ/src/compilation.rs

@@ -1,19 +1,21 @@
 //! A `Compilation` contains the compiled function bodies for a WebAssembly
 //! module.
 
-use crate::address_map::{ModuleAddressMap, ValueLabelsRanges};
-use crate::ModuleTranslation;
+use crate::{FunctionAddressMap, FunctionBodyData, ModuleTranslation};
 use cranelift_codegen::{binemit, ir, isa, isa::unwind::UnwindInfo};
 use cranelift_entity::PrimaryMap;
 use cranelift_wasm::{DefinedFuncIndex, FuncIndex, WasmError};
 use serde::{Deserialize, Serialize};
-use std::ops::Range;
 use thiserror::Error;
 
+#[allow(missing_docs)]
+pub type CompiledFunctions = PrimaryMap<DefinedFuncIndex, CompiledFunction>;
+
 /// Compiled function: machine code body, jump table offsets, and unwind information.
-#[derive(Serialize, Deserialize, Debug, Clone, PartialEq, Eq)]
+#[derive(Serialize, Deserialize, Debug, Clone, PartialEq, Eq, Default)]
+#[allow(missing_docs)]
 pub struct CompiledFunction {
-    /// The function body.
+    /// The machine code for this function.
     pub body: Vec<u8>,
 
     /// The jump tables offsets (in the body).
@@ -21,93 +23,13 @@ pub struct CompiledFunction {
 
     /// The unwind information.
     pub unwind_info: Option<UnwindInfo>,
-}
 
-type Functions = PrimaryMap<DefinedFuncIndex, CompiledFunction>;
-
-/// The result of compiling a WebAssembly module's functions.
-#[derive(Deserialize, Serialize, Debug, PartialEq, Eq)]
-pub struct Compilation {
-    /// Compiled machine code for the function bodies.
-    functions: Functions,
-}
-
-impl Compilation {
-    /// Creates a compilation artifact from a contiguous function buffer and a set of ranges
-    pub fn new(functions: Functions) -> Self {
-        Self { functions }
-    }
-
-    /// Allocates the compilation result with the given function bodies.
-    pub fn from_buffer(
-        buffer: Vec<u8>,
-        functions: impl IntoIterator<Item = (Range<usize>, ir::JumpTableOffsets)>,
-    ) -> Self {
-        Self::new(
-            functions
-                .into_iter()
-                .map(|(body_range, jt_offsets)| CompiledFunction {
-                    body: buffer[body_range].to_vec(),
-                    jt_offsets,
-                    unwind_info: None, // not implemented for lightbeam currently
-                })
-                .collect(),
-        )
-    }
-
-    /// Gets the bytes of a single function
-    pub fn get(&self, func: DefinedFuncIndex) -> &CompiledFunction {
-        &self.functions[func]
-    }
-
-    /// Gets the number of functions defined.
-    pub fn len(&self) -> usize {
-        self.functions.len()
-    }
-
-    /// Returns whether there are no functions defined.
-    pub fn is_empty(&self) -> bool {
-        self.functions.is_empty()
-    }
-
-    /// Returns unwind info for all defined functions.
-    pub fn unwind_info(&self) -> PrimaryMap<DefinedFuncIndex, &Option<UnwindInfo>> {
-        self.functions
-            .iter()
-            .map(|(_, func)| &func.unwind_info)
-            .collect::<PrimaryMap<DefinedFuncIndex, _>>()
-    }
-
-    /// Gets functions jump table offsets.
-    pub fn get_jt_offsets(&self) -> PrimaryMap<DefinedFuncIndex, ir::JumpTableOffsets> {
-        self.functions
-            .iter()
-            .map(|(_, func)| func.jt_offsets.clone())
-            .collect::<PrimaryMap<DefinedFuncIndex, _>>()
-    }
-}
-
-impl<'a> IntoIterator for &'a Compilation {
-    type IntoIter = Iter<'a>;
-    type Item = <Self::IntoIter as Iterator>::Item;
-
-    fn into_iter(self) -> Self::IntoIter {
-        Iter {
-            iterator: self.functions.iter(),
-        }
-    }
-}
-
-pub struct Iter<'a> {
-    iterator: <&'a Functions as IntoIterator>::IntoIter,
-}
-
-impl<'a> Iterator for Iter<'a> {
-    type Item = &'a CompiledFunction;
-
-    fn next(&mut self) -> Option<Self::Item> {
-        self.iterator.next().map(|(_, b)| b)
-    }
+    pub relocations: Vec<Relocation>,
+    pub address_map: FunctionAddressMap,
+    pub value_labels_ranges: cranelift_codegen::ValueLabelsRanges,
+    pub stack_slots: ir::StackSlots,
+    pub traps: Vec<TrapInformation>,
+    pub stack_maps: Vec<StackMapInformation>,
 }
 
 /// A record of a relocation to perform.
@@ -134,9 +56,6 @@ pub enum RelocationTarget {
     JumpTable(FuncIndex, ir::JumpTable),
 }
 
-/// Relocations to apply to function bodies.
-pub type Relocations = PrimaryMap<DefinedFuncIndex, Vec<Relocation>>;
-
 /// Information about trap.
 #[derive(Serialize, Deserialize, Debug, PartialEq, Eq, Clone)]
 pub struct TrapInformation {
@@ -148,9 +67,6 @@ pub struct TrapInformation {
     pub trap_code: ir::TrapCode,
 }
 
-/// Information about traps associated with the functions where the traps are placed.
-pub type Traps = PrimaryMap<DefinedFuncIndex, Vec<TrapInformation>>;
-
 /// The offset within a function of a GC safepoint, and its associated stack
 /// map.
 #[derive(Serialize, Deserialize, Debug, PartialEq, Eq, Clone)]
@@ -163,10 +79,6 @@ pub struct StackMapInformation {
     pub stack_map: binemit::Stackmap,
 }
 
-/// Information about GC safepoints and their associated stack maps within each
-/// function.
-pub type StackMaps = PrimaryMap<DefinedFuncIndex, Vec<StackMapInformation>>;
-
 /// An error while compiling WebAssembly to machine code.
 #[derive(Error, Debug)]
 pub enum CompileError {
@@ -183,22 +95,15 @@ pub enum CompileError {
     DebugInfoNotSupported,
 }
 
-/// A type alias for the result of `Compiler::compile_module`
-pub type CompileResult = (
-    Compilation,
-    Relocations,
-    ModuleAddressMap,
-    ValueLabelsRanges,
-    PrimaryMap<DefinedFuncIndex, ir::StackSlots>,
-    Traps,
-    StackMaps,
-);
-
-/// An implementation of a compiler from parsed WebAssembly module to native code.
-pub trait Compiler {
-    /// Compile a parsed module with the given `TargetIsa`.
-    fn compile_module(
-        translation: &ModuleTranslation,
+/// An implementation of a compiler from parsed WebAssembly module to native
+/// code.
+pub trait Compiler: Send + Sync {
+    /// Compile a function with the given `TargetIsa`.
+    fn compile_function(
+        &self,
+        translation: &ModuleTranslation<'_>,
+        index: DefinedFuncIndex,
+        data: &FunctionBodyData<'_>,
         isa: &dyn isa::TargetIsa,
-    ) -> Result<CompileResult, CompileError>;
+    ) -> Result<CompiledFunction, CompileError>;
 }