winch: Improve frame handling (#9708)

This commit addresses issues identified while working on issue #8091. It
improves the frame handling in Winch to prevent subtle bugs and enhance
the robustness of the code generation process.

Previously, there was no clear mechanism to verify when the frame was
fully set up and safe to access the local slots allocated for register
arguments, including the special slots used for the `VMContext`. As
a result, it was possible to inadvertently read from uninitialized
memory if calls were made before the frame was properly set up and
sealed.

This commit introduces two main changes with the objective to help
reduce the risk of introducing bugs related to the above:

* A `CodeGenPhase` trait, used via the type state pattern to clearly
  gate the operations allowed during each phase of the code generation
  process.

* Improve the semantics of locals, by clearly separating the notion of
  Wasm locals and special locals used by the compiler. This
  specialization allows a more accurate representation of the semantics
  of Wasm locals and their index space.

Author: saulecabrera

winch/codegen/src/abi/mod.rs

@@ -212,21 +212,6 @@ impl ABIOperand {
             _ => unreachable!(),
         }
     }
-
-    /// Get the register associated to this [`ABIOperand`].
-    pub fn get_reg(&self) -> Option<Reg> {
-        match *self {
-            ABIOperand::Reg { reg, .. } => Some(reg),
-            _ => None,
-        }
-    }
-
-    /// Get the type associated to this [`ABIOperand`].
-    pub fn ty(&self) -> WasmValType {
-        match *self {
-            ABIOperand::Reg { ty, .. } | ABIOperand::Stack { ty, .. } => ty,
-        }
-    }
 }
 
 /// Information about the [`ABIOperand`] information used in [`ABISig`].

winch/codegen/src/codegen/bounds.rs

@@ -4,7 +4,7 @@
 use super::env::HeapData;
 use crate::{
     abi::{scratch, vmctx},
-    codegen::CodeGenContext,
+    codegen::{CodeGenContext, Emission},
     isa::reg::{writable, Reg},
     masm::{IntCmpKind, MacroAssembler, OperandSize, RegImm, TrapCode},
     stack::TypedReg,
@@ -82,7 +82,7 @@ impl Index {
 
 /// Loads the bounds of the dynamic heap.
 pub(crate) fn load_dynamic_heap_bounds<M>(
-    context: &mut CodeGenContext,
+    context: &mut CodeGenContext<Emission>,
     masm: &mut M,
     heap: &HeapData,
     ptr_size: OperandSize,
@@ -149,7 +149,7 @@ pub(crate) fn ensure_index_and_offset<M: MacroAssembler>(
 /// criteria is in bounds.
 pub(crate) fn load_heap_addr_checked<M, F>(
     masm: &mut M,
-    context: &mut CodeGenContext,
+    context: &mut CodeGenContext<Emission>,
     ptr_size: OperandSize,
     heap: &HeapData,
     enable_spectre_mitigation: bool,

winch/codegen/src/codegen/call.rs

@@ -58,7 +58,7 @@
 
 use crate::{
     abi::{scratch, vmctx, ABIOperand, ABISig, RetArea},
-    codegen::{BuiltinFunction, BuiltinType, Callee, CodeGenContext},
+    codegen::{BuiltinFunction, BuiltinType, Callee, CodeGenContext, Emission},
     masm::{
         CalleeKind, ContextArgs, MacroAssembler, MemMoveDirection, OperandSize, SPOffset,
         VMContextLoc,
@@ -85,7 +85,7 @@ impl FnCall {
     pub fn emit<M: MacroAssembler>(
         env: &mut FuncEnv<M::Ptr>,
         masm: &mut M,
-        context: &mut CodeGenContext,
+        context: &mut CodeGenContext<Emission>,
         callee: Callee,
     ) {
         let (kind, callee_context) = Self::lower(env, context.vmoffsets, &callee, context, masm);
@@ -129,7 +129,7 @@ impl FnCall {
         env: &mut FuncEnv<M::Ptr>,
         vmoffsets: &VMOffsets<u8>,
         callee: &Callee,
-        context: &mut CodeGenContext,
+        context: &mut CodeGenContext<Emission>,
         masm: &mut M,
     ) -> (CalleeKind, ContextArgs) {
         let ptr = vmoffsets.ptr.size();
@@ -177,7 +177,7 @@ impl FnCall {
     fn lower_import<M: MacroAssembler, P: PtrSize>(
         index: FuncIndex,
         sig: &ABISig,
-        context: &mut CodeGenContext,
+        context: &mut CodeGenContext<Emission>,
         masm: &mut M,
         vmoffsets: &VMOffsets<P>,
     ) -> (CalleeKind, ContextArgs) {
@@ -204,7 +204,7 @@ impl FnCall {
     fn lower_funcref<M: MacroAssembler>(
         sig: &ABISig,
         ptr: impl PtrSize,
-        context: &mut CodeGenContext,
+        context: &mut CodeGenContext<Emission>,
         masm: &mut M,
     ) -> (CalleeKind, ContextArgs) {
         // Pop the funcref pointer to a register and allocate a register to hold the
@@ -275,7 +275,7 @@ impl FnCall {
         sig: &ABISig,
         callee_context: &ContextArgs,
         ret_area: Option<&RetArea>,
-        context: &mut CodeGenContext,
+        context: &mut CodeGenContext<Emission>,
         masm: &mut M,
     ) {
         let arg_count = sig.params.len_without_retptr();
@@ -337,7 +337,7 @@ impl FnCall {
         reserved_space: u32,
         ret_area: Option<RetArea>,
         masm: &mut M,
-        context: &mut CodeGenContext,
+        context: &mut CodeGenContext<Emission>,
     ) {
         // Free any registers holding any function references.
         match callee_kind {

winch/codegen/src/codegen/context.rs

@@ -3,6 +3,7 @@ use wasmtime_environ::{VMOffsets, WasmHeapType, WasmValType};
 use super::ControlStackFrame;
 use crate::{
     abi::{scratch, vmctx, ABIOperand, ABIResults, RetArea},
+    codegen::{CodeGenPhase, Emission, Prologue},
     frame::Frame,
     isa::reg::RegClass,
     masm::{MacroAssembler, OperandSize, RegImm, SPOffset, ShiftKind, StackSlot},
@@ -26,25 +27,78 @@ use crate::{
 /// generation process. The code generation context should
 /// be generally used as the single entry point to access
 /// the compound functionality provided by its elements.
-pub(crate) struct CodeGenContext<'a> {
+pub(crate) struct CodeGenContext<'a, P: CodeGenPhase> {
     /// The register allocator.
     pub regalloc: RegAlloc,
     /// The value stack.
     pub stack: Stack,
     /// The current function's frame.
-    pub frame: Frame,
+    pub frame: Frame<P>,
     /// Reachability state.
     pub reachable: bool,
     /// A reference to the VMOffsets.
     pub vmoffsets: &'a VMOffsets<u8>,
 }
 
-impl<'a> CodeGenContext<'a> {
+impl<'a> CodeGenContext<'a, Emission> {
+    /// Prepares arguments for emitting an i32 shift operation.
+    pub fn i32_shift<M>(&mut self, masm: &mut M, kind: ShiftKind)
+    where
+        M: MacroAssembler,
+    {
+        let top = self.stack.peek().expect("value at stack top");
+
+        if top.is_i32_const() {
+            let val = self
+                .stack
+                .pop_i32_const()
+                .expect("i32 const value at stack top");
+            let typed_reg = self.pop_to_reg(masm, None);
+            masm.shift_ir(
+                writable!(typed_reg.reg),
+                val as u64,
+                typed_reg.reg,
+                kind,
+                OperandSize::S32,
+            );
+            self.stack.push(typed_reg.into());
+        } else {
+            masm.shift(self, kind, OperandSize::S32);
+        }
+    }
+
+    /// Prepares arguments for emitting an i64 binary operation.
+    pub fn i64_shift<M>(&mut self, masm: &mut M, kind: ShiftKind)
+    where
+        M: MacroAssembler,
+    {
+        let top = self.stack.peek().expect("value at stack top");
+        if top.is_i64_const() {
+            let val = self
+                .stack
+                .pop_i64_const()
+                .expect("i64 const value at stack top");
+            let typed_reg = self.pop_to_reg(masm, None);
+            masm.shift_ir(
+                writable!(typed_reg.reg),
+                val as u64,
+                typed_reg.reg,
+                kind,
+                OperandSize::S64,
+            );
+            self.stack.push(typed_reg.into());
+        } else {
+            masm.shift(self, kind, OperandSize::S64);
+        };
+    }
+}
+
+impl<'a> CodeGenContext<'a, Prologue> {
     /// Create a new code generation context.
     pub fn new(
         regalloc: RegAlloc,
         stack: Stack,
-        frame: Frame,
+        frame: Frame<Prologue>,
         vmoffsets: &'a VMOffsets<u8>,
     ) -> Self {
         Self {
@@ -56,6 +110,19 @@ impl<'a> CodeGenContext<'a> {
         }
     }
 
+    /// Prepares the frame for the [`Emission`] code generation phase.
+    pub fn for_emission(self) -> CodeGenContext<'a, Emission> {
+        CodeGenContext {
+            regalloc: self.regalloc,
+            stack: self.stack,
+            reachable: self.reachable,
+            vmoffsets: self.vmoffsets,
+            frame: self.frame.for_emission(),
+        }
+    }
+}
+
+impl<'a> CodeGenContext<'a, Emission> {
     /// Request a specific register to the register allocator,
     /// spilling if not available.
     pub fn reg<M: MacroAssembler>(&mut self, named: Reg, masm: &mut M) -> Reg {
@@ -325,57 +392,6 @@ impl<'a> CodeGenContext<'a> {
         };
     }
 
-    /// Prepares arguments for emitting an i32 shift operation.
-    pub fn i32_shift<M>(&mut self, masm: &mut M, kind: ShiftKind)
-    where
-        M: MacroAssembler,
-    {
-        let top = self.stack.peek().expect("value at stack top");
-
-        if top.is_i32_const() {
-            let val = self
-                .stack
-                .pop_i32_const()
-                .expect("i32 const value at stack top");
-            let typed_reg = self.pop_to_reg(masm, None);
-            masm.shift_ir(
-                writable!(typed_reg.reg),
-                val as u64,
-                typed_reg.reg,
-                kind,
-                OperandSize::S32,
-            );
-            self.stack.push(typed_reg.into());
-        } else {
-            masm.shift(self, kind, OperandSize::S32);
-        }
-    }
-
-    /// Prepares arguments for emitting an i64 binary operation.
-    pub fn i64_shift<M>(&mut self, masm: &mut M, kind: ShiftKind)
-    where
-        M: MacroAssembler,
-    {
-        let top = self.stack.peek().expect("value at stack top");
-        if top.is_i64_const() {
-            let val = self
-                .stack
-                .pop_i64_const()
-                .expect("i64 const value at stack top");
-            let typed_reg = self.pop_to_reg(masm, None);
-            masm.shift_ir(
-                writable!(typed_reg.reg),
-                val as u64,
-                typed_reg.reg,
-                kind,
-                OperandSize::S64,
-            );
-            self.stack.push(typed_reg.into());
-        } else {
-            masm.shift(self, kind, OperandSize::S64);
-        };
-    }
-
     /// Prepares arguments for emitting a convert operation.
     pub fn convert_op<F, M>(&mut self, masm: &mut M, dst_ty: WasmValType, mut emit: F)
     where
@@ -511,7 +527,7 @@ impl<'a> CodeGenContext<'a> {
         mut calculate_ret_area: F,
     ) where
         M: MacroAssembler,
-        F: FnMut(&ABIResults, &mut CodeGenContext, &mut M) -> Option<RetArea>,
+        F: FnMut(&ABIResults, &mut CodeGenContext<Emission>, &mut M) -> Option<RetArea>,
     {
         let area = results
             .on_stack()
@@ -558,7 +574,7 @@ impl<'a> CodeGenContext<'a> {
     where
         M: MacroAssembler,
     {
-        let addr = masm.local_address(&self.frame.vmctx_slot);
+        let addr = masm.local_address(&self.frame.vmctx_slot());
         masm.load_ptr(addr, writable!(vmctx!(M)));
     }
 
@@ -570,7 +586,7 @@ impl<'a> CodeGenContext<'a> {
     fn spill_impl<M: MacroAssembler>(
         stack: &mut Stack,
         regalloc: &mut RegAlloc,
-        frame: &Frame,
+        frame: &Frame<Emission>,
         masm: &mut M,
     ) {
         stack.inner_mut().iter_mut().for_each(|v| match v {