winch: Prepare for an update to the wasm-tools crates (#5238)

This commit prepares the `winch` crate for updating `wasm-tools`,
notably changing a bit about how the visitation of operators works. This
moves the function body and wasm validator out of the `CodeGen`
structure and into parameters threaded into the emission of the actual
function.

Additionally the `VisitOperator` implementation was updated to remove
the explicit calls to the validator, favoring instead a macro-generated
solution to guarantee that all validation happens before any translation
proceeds. This means that the `VisitOperator for CodeGen` impl is now
infallible and the various methods have been inlined into the trait
methods as well as removing the `Result<_>`.

Finally this commit updates translation to call `validator.finish(..)`
which is required to perform the final validation steps of the function
body.

winch/codegen/src/codegen.rs

@@ -6,7 +6,7 @@ use crate::{
     stack::Stack,
 };
 use anyhow::Result;
-use wasmparser::{FuncValidator, FunctionBody, ValType, ValidatorResources};
+use wasmparser::{BinaryReader, FuncValidator, ValType, ValidatorResources, VisitOperator};
 
 /// The code generation context.
 pub(crate) struct CodeGenContext<'a, M>
@@ -32,9 +32,6 @@ pub(crate) struct CodeGen<'a, M>
 where
     M: MacroAssembler,
 {
-    /// A reference to the function body.
-    function: FunctionBody<'a>,
-
     /// The word size in bytes, extracted from the current ABI.
     word_size: u32,
 
@@ -46,36 +43,29 @@ where
 
     /// The register allocator.
     pub regalloc: RegAlloc,
-
-    /// Function body validator.
-    pub validator: FuncValidator<ValidatorResources>,
 }
 
 impl<'a, M> CodeGen<'a, M>
 where
     M: MacroAssembler,
 {
-    pub fn new<A: ABI>(
-        context: CodeGenContext<'a, M>,
-        sig: ABISig,
-        function: FunctionBody<'a>,
-        validator: FuncValidator<ValidatorResources>,
-        regalloc: RegAlloc,
-    ) -> Self {
+    pub fn new<A: ABI>(context: CodeGenContext<'a, M>, sig: ABISig, regalloc: RegAlloc) -> Self {
         Self {
-            function,
             word_size: <A as ABI>::word_bytes(),
             sig,
             context,
             regalloc,
-            validator,
         }
     }
 
     /// Emit the function body to machine code.
-    pub fn emit(&mut self) -> Result<Vec<String>> {
+    pub fn emit(
+        &mut self,
+        body: &mut BinaryReader<'a>,
+        validator: FuncValidator<ValidatorResources>,
+    ) -> Result<Vec<String>> {
         self.emit_start()
-            .and(self.emit_body())
+            .and(self.emit_body(body, validator))
             .and(self.emit_end())?;
         let buf = self.context.masm.finalize();
         let code = Vec::from(buf);
@@ -91,7 +81,11 @@ where
         Ok(())
     }
 
-    fn emit_body(&mut self) -> Result<()> {
+    fn emit_body(
+        &mut self,
+        body: &mut BinaryReader<'a>,
+        mut validator: FuncValidator<ValidatorResources>,
+    ) -> Result<()> {
         self.spill_register_arguments();
         let defined_locals_range = &self.context.frame.defined_locals_range;
         self.context.masm.zero_mem_range(
@@ -100,12 +94,35 @@ where
             &mut self.regalloc,
         );
 
-        let mut reader = self.function.get_operators_reader()?;
-        while !reader.eof() {
-            reader.visit_with_offset(self)??;
+        let mut visitor = ValidateThenVisit(&mut validator, self);
+        while !body.eof() {
+            body.visit_operator(&mut visitor)??;
         }
+        validator.finish(body.original_position())?;
+        return Ok(());
+
+        struct ValidateThenVisit<'a, T, U>(&'a mut T, &'a mut U);
+
+        macro_rules! validate_then_visit {
+            ($( @$proposal:ident $op:ident $({ $($arg:ident: $argty:ty),* })? => $visit:ident)*) => {
+                $(
+                    fn $visit(&mut self, offset: usize $($(,$arg: $argty)*)?) -> Self::Output {
+                        self.0.$visit(offset, $($($arg.clone()),*)?)?;
+                        Ok(self.1.$visit(offset, $($($arg),*)?))
+                    }
+                )*
+            };
+        }
+
+        impl<'a, T, U> VisitOperator<'a> for ValidateThenVisit<'_, T, U>
+        where
+            T: VisitOperator<'a, Output = wasmparser::Result<()>>,
+            U: VisitOperator<'a>,
+        {
+            type Output = Result<U::Output>;
 
-        reader.ensure_end().map_err(|e| e.into())
+            wasmparser::for_each_operator!(validate_then_visit);
+        }
     }
 
     // Emit the usual function end instruction sequence.

winch/codegen/src/frame/mod.rs

@@ -2,7 +2,7 @@ use crate::abi::{align_to, local::LocalSlot, ty_size, ABIArg, ABISig, ABI};
 use anyhow::Result;
 use smallvec::SmallVec;
 use std::ops::Range;
-use wasmparser::{FuncValidator, FunctionBody, ValType, ValidatorResources};
+use wasmparser::{BinaryReader, FuncValidator, ValType, ValidatorResources};
 
 // TODO:
 // SpiderMonkey's implementation uses 16;
@@ -39,13 +39,13 @@ impl Frame {
     /// Allocate a new Frame.
     pub fn new<A: ABI>(
         sig: &ABISig,
-        function: &mut FunctionBody,
+        body: &mut BinaryReader<'_>,
         validator: &mut FuncValidator<ValidatorResources>,
         abi: &A,
     ) -> Result<Self> {
         let (mut locals, defined_locals_start) = Self::compute_arg_slots(sig, abi)?;
         let (defined_slots, defined_locals_end) =
-            Self::compute_defined_slots(function, validator, defined_locals_start)?;
+            Self::compute_defined_slots(body, validator, defined_locals_start)?;
         locals.extend(defined_slots);
         let locals_size = align_to(defined_locals_end, abi.stack_align().into());
 
@@ -117,12 +117,11 @@ impl Frame {
     }
 
     fn compute_defined_slots(
-        body_data: &mut FunctionBody,
+        reader: &mut BinaryReader<'_>,
         validator: &mut FuncValidator<ValidatorResources>,
         next_stack: u32,
     ) -> Result<(Locals, u32)> {
         let mut next_stack = next_stack;
-        let mut reader = body_data.get_binary_reader();
         let local_count = reader.read_var_u32()?;
         let mut slots: Locals = Default::default();
 

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

@@ -34,7 +34,7 @@ impl TargetIsa for Aarch64 {
     fn compile_function(
         &self,
         _sig: &FuncType,
-        mut _body: FunctionBody,
+        _body: &FunctionBody,
         mut _validator: FuncValidator<ValidatorResources>,
     ) -> Result<Vec<String>> {
         todo!()

winch/codegen/src/isa/mod.rs

@@ -85,7 +85,7 @@ pub trait TargetIsa: Send + Sync {
     fn compile_function(
         &self,
         sig: &FuncType,
-        body: FunctionBody,
+        body: &FunctionBody,
         validator: FuncValidator<ValidatorResources>,
     ) -> Result<Vec<String>>;
 

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

@@ -47,9 +47,10 @@ impl TargetIsa for X64 {
     fn compile_function(
         &self,
         sig: &FuncType,
-        mut body: FunctionBody,
+        body: &FunctionBody,
         mut validator: FuncValidator<ValidatorResources>,
     ) -> Result<Vec<String>> {
+        let mut body = body.get_binary_reader();
         let masm = MacroAssembler::new();
         let stack = Stack::new();
         let abi = abi::X64ABI::default();
@@ -58,9 +59,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(masm, stack, &frame);
-        let mut codegen =
-            CodeGen::new::<abi::X64ABI>(codegen_context, abi_sig, body, validator, regalloc);
+        let mut codegen = CodeGen::new::<abi::X64ABI>(codegen_context, abi_sig, regalloc);
 
-        codegen.emit()
+        codegen.emit(&mut body, validator)
     }
 }

winch/codegen/src/visitor.rs

@@ -7,31 +7,13 @@
 use crate::codegen::CodeGen;
 use crate::masm::{MacroAssembler, OperandSize, RegImm};
 use crate::stack::Val;
-use anyhow::Result;
 use wasmparser::ValType;
 use wasmparser::VisitOperator;
 
 impl<'a, M> CodeGen<'a, M>
 where
     M: MacroAssembler,
 {
-    fn emit_i32_add(&mut self) -> Result<()> {
-        let is_const = self
-            .context
-            .stack
-            .peek()
-            .expect("value at stack top")
-            .is_i32_const();
-
-        if is_const {
-            self.add_imm_i32();
-        } else {
-            self.add_i32();
-        }
-
-        Ok(())
-    }
-
     fn add_imm_i32(&mut self) {
         let val = self
             .context
@@ -65,41 +47,6 @@ where
         self.regalloc.free_gpr(src);
         self.context.stack.push(Val::reg(dst));
     }
-
-    fn emit_i32_const(&mut self, val: i32) -> Result<()> {
-        self.context.stack.push(Val::i32(val));
-        Ok(())
-    }
-
-    fn emit_local_get(&mut self, index: u32) -> Result<()> {
-        let context = &mut self.context;
-        let slot = context
-            .frame
-            .get_local(index)
-            .expect(&format!("valid local at slot = {}", index));
-        match slot.ty {
-            ValType::I32 | ValType::I64 => context.stack.push(Val::local(index)),
-            _ => panic!("Unsupported type {:?} for local", slot.ty),
-        }
-
-        Ok(())
-    }
-
-    // TODO verify the case where the target local is on the stack.
-    fn emit_local_set(&mut self, index: u32) -> Result<()> {
-        let context = &mut self.context;
-        let frame = context.frame;
-        let slot = frame
-            .get_local(index)
-            .expect(&format!("vald local at slot = {}", index));
-        let size: OperandSize = slot.ty.into();
-        let src = self.regalloc.pop_to_reg(context, size);
-        let addr = context.masm.local_address(&slot);
-        context.masm.store(RegImm::reg(src), addr, size);
-        self.regalloc.free_gpr(src);
-
-        Ok(())
-    }
 }
 
 /// A macro to define unsupported WebAssembly operators.
@@ -136,30 +83,53 @@ impl<'a, M> VisitOperator<'a> for CodeGen<'a, M>
 where
     M: MacroAssembler,
 {
-    type Output = Result<()>;
+    type Output = ();
 
-    fn visit_i32_const(&mut self, offset: usize, value: i32) -> Result<()> {
-        self.validator.visit_i32_const(offset, value)?;
-        self.emit_i32_const(value)
+    fn visit_i32_const(&mut self, _offset: usize, val: i32) {
+        self.context.stack.push(Val::i32(val));
     }
 
-    fn visit_i32_add(&mut self, offset: usize) -> Result<()> {
-        self.validator.visit_i32_add(offset)?;
-        self.emit_i32_add()
-    }
+    fn visit_i32_add(&mut self, _offset: usize) {
+        let is_const = self
+            .context
+            .stack
+            .peek()
+            .expect("value at stack top")
+            .is_i32_const();
 
-    fn visit_end(&mut self, offset: usize) -> Result<()> {
-        self.validator.visit_end(offset).map_err(|e| e.into())
+        if is_const {
+            self.add_imm_i32();
+        } else {
+            self.add_i32();
+        }
     }
 
-    fn visit_local_get(&mut self, offset: usize, local_index: u32) -> Result<()> {
-        self.validator.visit_local_get(offset, local_index)?;
-        self.emit_local_get(local_index)
+    fn visit_end(&mut self, _offset: usize) {}
+
+    fn visit_local_get(&mut self, _offset: usize, index: u32) {
+        let context = &mut self.context;
+        let slot = context
+            .frame
+            .get_local(index)
+            .expect(&format!("valid local at slot = {}", index));
+        match slot.ty {
+            ValType::I32 | ValType::I64 => context.stack.push(Val::local(index)),
+            _ => panic!("Unsupported type {:?} for local", slot.ty),
+        }
     }
 
-    fn visit_local_set(&mut self, offset: usize, local_index: u32) -> Result<()> {
-        self.validator.visit_local_set(offset, local_index)?;
-        self.emit_local_set(local_index)
+    // TODO verify the case where the target local is on the stack.
+    fn visit_local_set(&mut self, _offset: usize, index: u32) {
+        let context = &mut self.context;
+        let frame = context.frame;
+        let slot = frame
+            .get_local(index)
+            .expect(&format!("vald local at slot = {}", index));
+        let size: OperandSize = slot.ty.into();
+        let src = self.regalloc.pop_to_reg(context, size);
+        let addr = context.masm.local_address(&slot);
+        context.masm.store(RegImm::reg(src), addr, size);
+        self.regalloc.free_gpr(src);
     }
 
     wasmparser::for_each_operator!(def_unsupported);

winch/src/main.rs

@@ -63,7 +63,7 @@ fn compile(
     let FunctionBodyData { body, validator } = f.1;
     let validator = validator.into_validator(Default::default());
     let buffer = isa
-        .compile_function(&sig, body, validator)
+        .compile_function(&sig, &body, validator)
         .expect("Couldn't compile function");
     for i in buffer {
         println!("{}", i);