winch: Drop FuncEnv trait (#6443)

This commit is a small cleanup to drop the usage of the `FuncEnv` trait.

In #6358, we agreed on
making `winch-codegen` directly depend on `wasmtime-environ`.

Introducing a direct relatioship between `winch-codegen` and
`wasmtime-environ` means that the `FuncEnv` trait is no longer serving
its original purpose, and we can drop the usage of the trait and use the
types exposed from `winch-codegen` directly instead.

Even though this change drops the `FuncEnv` trait, it still keeps
a `FuncEnv` struct, which is used during code generation.

Cargo.toml

@@ -110,7 +110,6 @@ members = [
   "fuzz",
   "winch",
   "winch/codegen",
-  "winch/environ"
 ]
 exclude = [
   'crates/wasi-common/WASI/tools/witx-cli',
@@ -176,7 +175,6 @@ cranelift = { path = "cranelift/umbrella", version = "0.97.0" }
 wasmtime-wasi-http = { path = "crates/wasi-http", version = "=10.0.0" }
 
 winch-codegen = { path = "winch/codegen", version = "=0.8.0" }
-winch-environ = { path = "winch/environ", version = "=0.8.0" }
 winch-filetests = { path = "winch/filetests" }
 winch-test-macros = { path = "winch/test-macros" }
 

crates/winch/Cargo.toml

@@ -9,7 +9,6 @@ repository = "https://github.com/bytecodealliance/wasmtime"
 
 [dependencies]
 winch-codegen = { workspace = true }
-winch-environ = { workspace = true }
 target-lexicon = { workspace = true }
 wasmtime-environ = { workspace = true }
 anyhow = { workspace = true }

crates/winch/src/compiler.rs

@@ -6,11 +6,9 @@ use wasmparser::FuncValidatorAllocations;
 use wasmtime_cranelift_shared::{CompiledFunction, ModuleTextBuilder};
 use wasmtime_environ::{
     CompileError, DefinedFuncIndex, FilePos, FuncIndex, FunctionBodyData, FunctionLoc,
-    ModuleTranslation, ModuleTypes, PrimaryMap, TrapEncodingBuilder, Tunables, VMOffsets,
-    WasmFunctionInfo,
+    ModuleTranslation, ModuleTypes, PrimaryMap, TrapEncodingBuilder, Tunables, WasmFunctionInfo,
 };
 use winch_codegen::{TargetIsa, TrampolineKind};
-use winch_environ::FuncEnv;
 
 pub(crate) struct Compiler {
     isa: Box<dyn TargetIsa>,
@@ -68,11 +66,9 @@ impl wasmtime_environ::Compiler for Compiler {
                 .unwrap(),
         );
         let mut validator = validator.into_validator(self.take_allocations());
-        let vmoffsets = VMOffsets::new(self.isa.pointer_bytes(), &translation.module);
-        let env = FuncEnv::new(&translation.module, translation.get_types());
         let buffer = self
             .isa
-            .compile_function(&sig, &body, &vmoffsets, &env, &mut validator)
+            .compile_function(&sig, &body, &translation, &mut validator)
             .map_err(|e| CompileError::Codegen(format!("{e:?}")));
         self.save_allocations(validator.into_allocations());
         let buffer = buffer?;

scripts/publish.rs

@@ -57,7 +57,6 @@ const CRATES_TO_PUBLISH: &[&str] = &[
     "wasmtime-jit",
     "wasmtime-cache",
     "winch-codegen",
-    "winch-environ",
     "wasmtime-winch",
     "wasmtime",
     // wasi-common/wiggle

winch/Cargo.toml

@@ -13,7 +13,6 @@ path = "src/main.rs"
 
 [dependencies]
 winch-codegen = { workspace = true }
-winch-environ = { workspace = true }
 winch-filetests = { workspace = true }
 winch-test-macros = { workspace = true }
 wasmtime-environ = { workspace = true }

winch/codegen/src/codegen/env.rs

@@ -1,10 +1,41 @@
 use wasmparser::FuncType;
+use wasmtime_environ::{FuncIndex, ModuleTranslation, PtrSize, VMOffsets};
 
-/// Function environment used the by the code generation to
-/// resolve module and runtime-specific information.
-pub trait FuncEnv {
-    /// Get the callee information from a given function index.
-    fn callee_from_index(&self, index: u32) -> Callee;
+/// The function environment.
+///
+/// Contains all information about the module and runtime that is accessible to
+/// to a particular function during code generation.
+pub struct FuncEnv<'a, P> {
+    /// Offsets to the fields within the `VMContext` ptr.
+    pub vmoffsets: VMOffsets<P>,
+    /// Metadata about the translation process of a WebAssembly module.
+    pub translation: &'a ModuleTranslation<'a>,
+}
+
+impl<'a, P: PtrSize> FuncEnv<'a, P> {
+    /// Create a new function environment.
+    pub fn new(ptr: P, translation: &'a ModuleTranslation) -> Self {
+        let vmoffsets = VMOffsets::new(ptr, &translation.module);
+        Self {
+            vmoffsets,
+            translation,
+        }
+    }
+
+    /// Resolves a function [`Callee`] from an index.
+    pub fn callee_from_index(&self, idx: FuncIndex) -> Callee {
+        let types = &self.translation.get_types();
+        let ty = types
+            .function_at(idx.as_u32())
+            .unwrap_or_else(|| panic!("function type at index: {}", idx.as_u32()));
+        let import = self.translation.module.is_imported_function(idx);
+
+        Callee {
+            ty: ty.clone(),
+            import,
+            index: idx,
+        }
+    }
 }
 
 /// Metadata about a function callee.  Use by the code generation
@@ -15,5 +46,5 @@ pub struct Callee {
     /// A flag to determine if the callee is imported.
     pub import: bool,
     /// The callee index in the WebAssembly function index space.
-    pub index: u32,
+    pub index: FuncIndex,
 }

winch/codegen/src/codegen/mod.rs

@@ -9,7 +9,7 @@ use call::FnCall;
 use wasmparser::{
     BinaryReader, FuncType, FuncValidator, ValType, ValidatorResources, VisitOperator,
 };
-use wasmtime_environ::{FuncIndex, VMOffsets};
+use wasmtime_environ::{FuncIndex, PtrSize};
 
 mod context;
 pub(crate) use context::*;
@@ -18,50 +18,47 @@ pub use env::*;
 pub mod call;
 
 /// The code generation abstraction.
-pub(crate) struct CodeGen<'a, A, M>
+pub(crate) struct CodeGen<'a, A, M, P>
 where
     M: MacroAssembler,
     A: ABI,
+    P: PtrSize,
 {
     /// The ABI-specific representation of the function signature, excluding results.
     sig: ABISig,
 
     /// The code generation context.
     pub context: CodeGenContext<'a>,
 
+    /// A reference to the function compilation environment.
+    pub env: FuncEnv<'a, P>,
+
     /// The MacroAssembler.
     pub masm: &'a mut M,
 
-    /// A reference to the function compilation environment.
-    pub env: &'a dyn env::FuncEnv,
-
     /// A reference to the current ABI.
     pub abi: &'a A,
-
-    /// Offsets used with the VM context pointer.
-    vmoffsets: &'a VMOffsets<u8>,
 }
 
-impl<'a, A, M> CodeGen<'a, A, M>
+impl<'a, A, M, P> CodeGen<'a, A, M, P>
 where
     M: MacroAssembler,
     A: ABI,
+    P: PtrSize,
 {
     pub fn new(
         masm: &'a mut M,
         abi: &'a A,
         context: CodeGenContext<'a>,
-        env: &'a dyn FuncEnv,
+        env: FuncEnv<'a, P>,
         sig: ABISig,
-        vmoffsets: &'a VMOffsets<u8>,
     ) -> Self {
         Self {
             sig,
             context,
             masm,
             abi,
             env,
-            vmoffsets,
         }
     }
 
@@ -137,7 +134,7 @@ where
 
     /// Emit a direct function call.
     pub fn emit_call(&mut self, index: FuncIndex) {
-        let callee = self.env.callee_from_index(index.as_u32());
+        let callee = self.env.callee_from_index(index);
         let (sig, callee_addr): (ABISig, Option<<M as MacroAssembler>::Address>) = if callee.import
         {
             let mut params = vec![ValType::I64, ValType::I64];
@@ -146,14 +143,14 @@ where
 
             let caller_vmctx = <A as ABI>::vmctx_reg();
             let callee_vmctx = self.context.any_gpr(self.masm);
-            let callee_vmctx_offset = self.vmoffsets.vmctx_vmfunction_import_vmctx(index);
+            let callee_vmctx_offset = self.env.vmoffsets.vmctx_vmfunction_import_vmctx(index);
             let callee_vmctx_addr = self.masm.address_at_reg(caller_vmctx, callee_vmctx_offset);
             // FIXME Remove harcoded operand size, this will be needed
             // once 32-bit architectures are supported.
             self.masm
                 .load(callee_vmctx_addr, callee_vmctx, OperandSize::S64);
 
-            let callee_body_offset = self.vmoffsets.vmctx_vmfunction_import_wasm_call(index);
+            let callee_body_offset = self.env.vmoffsets.vmctx_vmfunction_import_wasm_call(index);
             let callee_addr = self.masm.address_at_reg(caller_vmctx, callee_body_offset);
 
             // Put the callee / caller vmctx at the start of the

winch/codegen/src/isa/aarch64/mod.rs

@@ -1,14 +1,14 @@
 use self::regs::{scratch, ALL_GPR};
 use crate::{
     abi::ABI,
-    codegen::{CodeGen, CodeGenContext},
+    codegen::{CodeGen, CodeGenContext, FuncEnv},
     frame::{DefinedLocals, Frame},
     isa::{Builder, CallingConvention, TargetIsa},
     masm::MacroAssembler,
     regalloc::RegAlloc,
     regset::RegSet,
     stack::Stack,
-    FuncEnv, TrampolineKind,
+    TrampolineKind,
 };
 use anyhow::Result;
 use cranelift_codegen::settings::{self, Flags};
@@ -17,7 +17,7 @@ use cranelift_codegen::{MachTextSectionBuilder, TextSectionBuilder};
 use masm::MacroAssembler as Aarch64Masm;
 use target_lexicon::Triple;
 use wasmparser::{FuncType, FuncValidator, FunctionBody, ValidatorResources};
-use wasmtime_environ::VMOffsets;
+use wasmtime_environ::ModuleTranslation;
 
 mod abi;
 mod address;
@@ -86,8 +86,7 @@ impl TargetIsa for Aarch64 {
         &self,
         sig: &FuncType,
         body: &FunctionBody,
-        vmoffsets: &VMOffsets<u8>,
-        env: &dyn FuncEnv,
+        translation: &ModuleTranslation,
         validator: &mut FuncValidator<ValidatorResources>,
     ) -> Result<MachBufferFinalized<Final>> {
         let mut body = body.get_binary_reader();
@@ -101,7 +100,8 @@ impl TargetIsa for Aarch64 {
         // TODO: Add floating point bitmask
         let regalloc = RegAlloc::new(RegSet::new(ALL_GPR, 0), scratch());
         let codegen_context = CodeGenContext::new(regalloc, stack, &frame);
-        let mut codegen = CodeGen::new(&mut masm, &abi, codegen_context, env, abi_sig, vmoffsets);
+        let env = FuncEnv::new(self.pointer_bytes(), translation);
+        let mut codegen = CodeGen::new(&mut masm, &abi, codegen_context, env, abi_sig);
 
         codegen.emit(&mut body, validator)?;
         Ok(masm.finalize())

winch/codegen/src/isa/mod.rs

@@ -1,4 +1,4 @@
-use crate::{FuncEnv, TrampolineKind};
+use crate::TrampolineKind;
 use anyhow::{anyhow, Result};
 use core::fmt::Formatter;
 use cranelift_codegen::isa::{CallConv, IsaBuilder};
@@ -10,7 +10,7 @@ use std::{
 };
 use target_lexicon::{Architecture, Triple};
 use wasmparser::{FuncType, FuncValidator, FunctionBody, ValidatorResources};
-use wasmtime_environ::VMOffsets;
+use wasmtime_environ::ModuleTranslation;
 
 #[cfg(feature = "x64")]
 pub(crate) mod x64;
@@ -149,8 +149,7 @@ pub trait TargetIsa: Send + Sync {
         &self,
         sig: &FuncType,
         body: &FunctionBody,
-        vmoffsets: &VMOffsets<u8>,
-        env: &dyn FuncEnv,
+        translation: &ModuleTranslation,
         validator: &mut FuncValidator<ValidatorResources>,
     ) -> Result<MachBufferFinalized<Final>>;
 

winch/codegen/src/isa/x64/mod.rs

@@ -1,14 +1,13 @@
 use crate::{
     abi::ABI,
-    codegen::{CodeGen, CodeGenContext},
+    codegen::{CodeGen, CodeGenContext, FuncEnv},
 };
 
 use crate::frame::{DefinedLocals, Frame};
 use crate::isa::{x64::masm::MacroAssembler as X64Masm, CallingConvention};
 use crate::masm::MacroAssembler;
 use crate::regalloc::RegAlloc;
 use crate::stack::Stack;
-use crate::FuncEnv;
 use crate::{
     isa::{Builder, TargetIsa},
     regset::RegSet,
@@ -20,7 +19,7 @@ use cranelift_codegen::{isa::x64::settings as x64_settings, Final, MachBufferFin
 use cranelift_codegen::{MachTextSectionBuilder, TextSectionBuilder};
 use target_lexicon::Triple;
 use wasmparser::{FuncType, FuncValidator, FunctionBody, ValidatorResources};
-use wasmtime_environ::VMOffsets;
+use wasmtime_environ::ModuleTranslation;
 
 use self::regs::ALL_GPR;
 
@@ -91,8 +90,7 @@ impl TargetIsa for X64 {
         &self,
         sig: &FuncType,
         body: &FunctionBody,
-        vmoffsets: &VMOffsets<u8>,
-        env: &dyn FuncEnv,
+        translation: &ModuleTranslation,
         validator: &mut FuncValidator<ValidatorResources>,
     ) -> Result<MachBufferFinalized<Final>> {
         let mut body = body.get_binary_reader();
@@ -106,7 +104,8 @@ impl TargetIsa for X64 {
         // TODO Add in floating point bitmask
         let regalloc = RegAlloc::new(RegSet::new(ALL_GPR, 0), regs::scratch());
         let codegen_context = CodeGenContext::new(regalloc, stack, &frame);
-        let mut codegen = CodeGen::new(&mut masm, &abi, codegen_context, env, abi_sig, vmoffsets);
+        let env = FuncEnv::new(self.pointer_bytes(), translation);
+        let mut codegen = CodeGen::new(&mut masm, &abi, codegen_context, env, abi_sig);
 
         codegen.emit(&mut body, validator)?;
 

winch/codegen/src/lib.rs

@@ -7,8 +7,8 @@
 #![cfg_attr(not(feature = "all-arch"), allow(dead_code))]
 
 mod abi;
+pub use codegen::FuncEnv;
 mod codegen;
-pub use codegen::{Callee, FuncEnv};
 mod frame;
 pub mod isa;
 pub use isa::*;