winch: Drop `FuncEnv` trait (#6443)

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

In https://github.com/bytecodealliance/wasmtime/pull/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.lock

@@ -4354,7 +4354,6 @@ dependencies = [
  "wasmtime-cranelift-shared",
  "wasmtime-environ",
  "winch-codegen",
- "winch-environ",
 ]
 
 [[package]]
@@ -4525,15 +4524,6 @@ dependencies = [
  "wasmtime-environ",
 ]
 
-[[package]]
-name = "winch-environ"
-version = "0.8.0"
-dependencies = [
- "wasmparser 0.105.0",
- "wasmtime-environ",
- "winch-codegen",
-]
-
 [[package]]
 name = "winch-filetests"
 version = "0.0.0"
@@ -4548,7 +4538,6 @@ dependencies = [
  "wasmtime-environ",
  "wat",
  "winch-codegen",
- "winch-environ",
  "winch-test-macros",
 ]
 
@@ -4579,7 +4568,6 @@ dependencies = [
  "wasmtime-environ",
  "wat",
  "winch-codegen",
- "winch-environ",
  "winch-filetests",
  "winch-test-macros",
 ]

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,10 +18,11 @@ 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,
@@ -29,31 +30,28 @@ where
     /// 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,
@@ -61,7 +59,6 @@ where
             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,6 +1,6 @@
 use crate::{
     abi::ABI,
-    codegen::{CodeGen, CodeGenContext},
+    codegen::{CodeGen, CodeGenContext, FuncEnv},
 };
 
 use crate::frame::{DefinedLocals, Frame};
@@ -8,7 +8,6 @@ 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::*;

winch/codegen/src/visitor.rs

@@ -10,7 +10,7 @@ use crate::masm::{DivKind, MacroAssembler, OperandSize, RegImm, RemKind};
 use crate::stack::Val;
 use wasmparser::ValType;
 use wasmparser::VisitOperator;
-use wasmtime_environ::FuncIndex;
+use wasmtime_environ::{FuncIndex, PtrSize};
 
 /// A macro to define unsupported WebAssembly operators.
 ///
@@ -57,10 +57,11 @@ macro_rules! def_unsupported {
     (emit $unsupported:tt $($rest:tt)*) => {$($rest)*};
 }
 
-impl<'a, A, M> VisitOperator<'a> for CodeGen<'a, A, M>
+impl<'a, A, M, P> VisitOperator<'a> for CodeGen<'a, A, M, P>
 where
     M: MacroAssembler,
     A: ABI,
+    P: PtrSize,
 {
     type Output = ();
 

winch/environ/Cargo.toml

@@ -1,14 +0,0 @@
-[package]
-authors = ["The Winch Project Developers"]
-name = "winch-environ"
-description = "Implementation of Winch's function compilation environment"
-license = "Apache-2.0 WITH LLVM-exception"
-repository = "https://github.com/bytecodealliance/wasmtime"
-version = "0.8.0"
-edition.workspace = true
-
-
-[dependencies]
-winch-codegen = { workspace = true }
-wasmtime-environ = { workspace = true }
-wasmparser = { workspace = true }

winch/environ/src/lib.rs

@@ -1,39 +0,0 @@
-//! This crate implements Winch's function compilation environment,
-//! which allows Winch's code generation to resolve module and runtime
-//! specific information.  This crate mainly implements the
-//! `winch_codegen::FuncEnv` trait.
-
-use wasmparser::types::Types;
-use wasmtime_environ::{FuncIndex, Module};
-use winch_codegen::{self, Callee};
-
-/// Function environment containing module and runtime specific
-/// information.
-pub struct FuncEnv<'a> {
-    /// The translated WebAssembly module.
-    pub module: &'a Module,
-    /// Type information about a module, once it has been validated.
-    pub types: &'a Types,
-}
-
-impl<'a> winch_codegen::FuncEnv for FuncEnv<'a> {
-    fn callee_from_index(&self, index: u32) -> Callee {
-        let func = self
-            .types
-            .function_at(index)
-            .unwrap_or_else(|| panic!("function type at index: {}", index));
-
-        Callee {
-            ty: func.clone(),
-            import: self.module.is_imported_function(FuncIndex::from_u32(index)),
-            index,
-        }
-    }
-}
-
-impl<'a> FuncEnv<'a> {
-    /// Create a new function environment.
-    pub fn new(module: &'a Module, types: &'a Types) -> Self {
-        Self { module, types }
-    }
-}

winch/filetests/Cargo.toml

@@ -12,7 +12,6 @@ edition.workspace = true
 winch-test-macros = {workspace = true}
 target-lexicon = { workspace = true }
 winch-codegen = { workspace = true, features = ['all-arch'] }
-winch-environ = { workspace = true }
 wasmtime-environ = { workspace = true }
 anyhow = { workspace = true }
 wat = { workspace = true }

winch/filetests/src/lib.rs

@@ -9,13 +9,12 @@ mod test {
     use similar::TextDiff;
     use std::str::FromStr;
     use target_lexicon::Triple;
-    use wasmtime_environ::VMOffsets;
+    use wasmtime_environ::ModuleTranslation;
     use wasmtime_environ::{
         wasmparser::{Parser as WasmParser, Validator},
         DefinedFuncIndex, FunctionBodyData, ModuleEnvironment, Tunables,
     };
     use winch_codegen::{lookup, TargetIsa};
-    use winch_environ::FuncEnv;
     use winch_test_macros::generate_file_tests;
 
     #[derive(Clone, Debug, Serialize, Deserialize)]
@@ -108,14 +107,10 @@ mod test {
         let _ = types.finish();
 
         let body_inputs = std::mem::take(&mut translation.function_body_inputs);
-        let module = &translation.module;
-        let types = translation.get_types();
-        let env = FuncEnv::new(module, &types);
-        let vmoffsets = VMOffsets::new(isa.pointer_bytes(), module);
 
         let binding = body_inputs
             .into_iter()
-            .map(|func| compile(&env, &isa, &vmoffsets, func).join("\n"))
+            .map(|func| compile(&isa, &translation, func).join("\n"))
             .collect::<Vec<String>>()
             .join("\n\n");
         let actual = binding.as_str();
@@ -147,20 +142,22 @@ mod test {
     }
 
     fn compile(
-        env: &FuncEnv,
         isa: &Box<dyn TargetIsa>,
-        vmoffsets: &VMOffsets<u8>,
+        translation: &ModuleTranslation,
         f: (DefinedFuncIndex, FunctionBodyData<'_>),
     ) -> Vec<String> {
-        let index = env.module.func_index(f.0);
-        let sig = env
-            .types
+        let module = &translation.module;
+        let types = &translation.get_types();
+
+        let index = module.func_index(f.0);
+        let sig = types
             .function_at(index.as_u32())
             .expect(&format!("function type at index {:?}", index.as_u32()));
+
         let FunctionBodyData { body, validator } = f.1;
         let mut validator = validator.into_validator(Default::default());
         let buffer = isa
-            .compile_function(&sig, &body, vmoffsets, env, &mut validator)
+            .compile_function(&sig, &body, &translation, &mut validator)
             .expect("Couldn't compile function");
 
         disasm(buffer.data(), isa).unwrap()

winch/src/compile.rs

@@ -5,10 +5,9 @@ use std::{fs, path::PathBuf, str::FromStr};
 use target_lexicon::Triple;
 use wasmtime_environ::{
     wasmparser::{Parser as WasmParser, Validator},
-    DefinedFuncIndex, FunctionBodyData, ModuleEnvironment, Tunables, VMOffsets,
+    DefinedFuncIndex, FunctionBodyData, ModuleEnvironment, ModuleTranslation, Tunables,
 };
 use winch_codegen::{lookup, TargetIsa};
-use winch_environ::FuncEnv;
 use winch_filetests::disasm::disasm;
 
 #[derive(Parser, Debug)]
@@ -36,37 +35,30 @@ pub fn run(opt: &Options) -> Result<()> {
     let mut translation = ModuleEnvironment::new(&tunables, &mut validator, &mut types)
         .translate(parser, &bytes)
         .context("Failed to translate WebAssembly module")?;
-
-    let vmoffsets = VMOffsets::new(isa.pointer_bytes(), &translation.module);
     let _ = types.finish();
-
     let body_inputs = std::mem::take(&mut translation.function_body_inputs);
-    let module = &translation.module;
-    let types = translation.get_types();
-    let env = FuncEnv::new(module, &types);
 
     body_inputs
         .into_iter()
-        .try_for_each(|func| compile(&env, &isa, &vmoffsets, func))?;
+        .try_for_each(|func| compile(&isa, &translation, func))?;
 
     Ok(())
 }
 
 fn compile(
-    env: &FuncEnv,
     isa: &Box<dyn TargetIsa>,
-    vmoffsets: &VMOffsets<u8>,
+    translation: &ModuleTranslation,
     f: (DefinedFuncIndex, FunctionBodyData<'_>),
 ) -> Result<()> {
-    let index = env.module.func_index(f.0);
-    let sig = env
-        .types
+    let index = translation.module.func_index(f.0);
+    let types = &translation.get_types();
+    let sig = types
         .function_at(index.as_u32())
         .expect(&format!("function type at index {:?}", index.as_u32()));
     let FunctionBodyData { body, validator } = f.1;
     let mut validator = validator.into_validator(Default::default());
     let buffer = isa
-        .compile_function(&sig, &body, vmoffsets, env, &mut validator)
+        .compile_function(&sig, &body, &translation, &mut validator)
         .expect("Couldn't compile function");
 
     println!("Disassembly for function: {}", index.as_u32());