winch: Add support for the <i32|i64>.div_* instructions (#5807)

* Refactor the structure and responsibilities of `CodeGenContext`

This commit refactors how the `CodeGenContext` is used throughout the code
generation process, making it easier to pass it around when more flexibility is
desired in the MacroAssembler to perform the lowering of certain instructions.

As of this change, the responsibility of the `CodeGenContext` is to provide an
interface for operations that require an orchestration between the register
allocator, the value stack and function's frame. The MacroAssembler is removed
from the CodeGenContext as is passed as a dependency where needed, effectly
using it as an independent code generation interface only.

By giving more responsibilities to the `CodeGenContext` we can clearly separate
the concerns of the register allocator, which previously did more than it
should (e.g. popping values and spilling).

This change ultimately allows passing in the `CodeGenContext` to the
`MacroAssembler` when a given instruction cannot be generically described
through a common interface. Allowing each implementation to decide the best way
to lower a particular instruction.

* winch: Add support for the WebAssembly `<i32|i64>.div_*` instructions

Given that some architectures have very specific requirements on how to handle
division, this change uses `CodeGenContext` as a dependency to the `div`
MacroAssembler instruction to ensure that each implementation can decide on how to lower the
division. This approach also allows -- in architectures where division can be
expressed as an ordinary binary operation -- to rely on the
`CodeGenContext::i32_binop` or `CodeGenContext::i64_binop` helpers.

Author: saulecabrera

winch/codegen/src/codegen/context.rs

@@ -0,0 +1,204 @@
+use crate::{
+    frame::Frame,
+    masm::{MacroAssembler, OperandSize, RegImm},
+    reg::Reg,
+    regalloc::RegAlloc,
+    stack::{Stack, Val},
+};
+
+/// The code generation context.
+/// The code generation context is made up of three
+/// essential data structures:
+///
+/// * The register allocator, in charge of keeping the inventory of register
+///   availability.
+/// * The value stack, which keeps track of the state of the values
+///   after each operation.
+/// * The current function's frame.
+///
+/// These data structures normally require cooperating with each other
+/// to perform most of the operations needed during the code
+/// 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> {
+    /// The register allocator.
+    pub regalloc: RegAlloc,
+    /// The value stack.
+    pub stack: Stack,
+    /// The current function's frame.
+    pub frame: &'a Frame,
+}
+
+impl<'a> CodeGenContext<'a> {
+    /// Create a new code generation context.
+    pub fn new(regalloc: RegAlloc, stack: Stack, frame: &'a Frame) -> Self {
+        Self {
+            regalloc,
+            stack,
+            frame,
+        }
+    }
+
+    /// Request a specific general purpose register to the register allocator,
+    /// spilling if not available.
+    pub fn gpr<M: MacroAssembler>(&mut self, named: Reg, masm: &mut M) -> Reg {
+        self.regalloc.gpr(named, &mut |regalloc| {
+            Self::spill(&mut self.stack, regalloc, &self.frame, masm)
+        })
+    }
+
+    /// Request the next avaiable general purpose register to the register allocator,
+    /// spilling if no registers are available.
+    pub fn any_gpr<M: MacroAssembler>(&mut self, masm: &mut M) -> Reg {
+        self.regalloc
+            .any_gpr(&mut |regalloc| Self::spill(&mut self.stack, regalloc, &self.frame, masm))
+    }
+
+    /// Free the given general purpose register.
+    pub fn free_gpr(&mut self, reg: Reg) {
+        self.regalloc.free_gpr(reg);
+    }
+
+    /// Loads the stack top value into a register, if it isn't already one;
+    /// spilling if there are no registers available.
+    pub fn pop_to_reg<M: MacroAssembler>(&mut self, masm: &mut M, size: OperandSize) -> Reg {
+        if let Some(reg) = self.stack.pop_reg() {
+            return reg;
+        }
+
+        let dst = self.any_gpr(masm);
+        let val = self.stack.pop().expect("a value at stack top");
+        Self::move_val_to_reg(val, dst, masm, self.frame, size);
+        dst
+    }
+
+    /// Checks if the stack top contains the given register. The register
+    /// gets allocated otherwise, potentially causing a spill.
+    /// Once the requested register is allocated, the value at the top of the stack
+    /// gets loaded into the register.
+    pub fn pop_to_named_reg<M: MacroAssembler>(
+        &mut self,
+        masm: &mut M,
+        named: Reg,
+        size: OperandSize,
+    ) -> Reg {
+        if let Some(reg) = self.stack.pop_named_reg(named) {
+            return reg;
+        }
+
+        let dst = self.gpr(named, masm);
+        let val = self.stack.pop().expect("a value at stack top");
+        Self::move_val_to_reg(val, dst, masm, self.frame, size);
+        dst
+    }
+
+    fn move_val_to_reg<M: MacroAssembler>(
+        src: Val,
+        dst: Reg,
+        masm: &mut M,
+        frame: &Frame,
+        size: OperandSize,
+    ) {
+        match src {
+            Val::Reg(src) => masm.mov(RegImm::reg(src), RegImm::reg(dst), size),
+            Val::I32(imm) => masm.mov(RegImm::imm(imm.into()), RegImm::reg(dst), size),
+            Val::I64(imm) => masm.mov(RegImm::imm(imm), RegImm::reg(dst), size),
+            Val::Local(index) => {
+                let slot = frame
+                    .get_local(index)
+                    .expect(&format!("valid locat at index = {}", index));
+                let addr = masm.local_address(&slot);
+                masm.load(addr, dst, slot.ty.into());
+            }
+            v => panic!("Unsupported value {:?}", v),
+        };
+    }
+
+    /// Prepares arguments for emitting an i32 binary operation.
+    pub fn i32_binop<F, M>(&mut self, masm: &mut M, emit: &mut F)
+    where
+        F: FnMut(&mut M, RegImm, RegImm, OperandSize),
+        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 reg = self.pop_to_reg(masm, OperandSize::S32);
+            emit(
+                masm,
+                RegImm::reg(reg),
+                RegImm::imm(val as i64),
+                OperandSize::S32,
+            );
+            self.stack.push(Val::reg(reg));
+        } else {
+            let src = self.pop_to_reg(masm, OperandSize::S32);
+            let dst = self.pop_to_reg(masm, OperandSize::S32);
+            emit(masm, dst.into(), src.into(), OperandSize::S32);
+            self.regalloc.free_gpr(src);
+            self.stack.push(Val::reg(dst));
+        }
+    }
+
+    /// Prepares arguments for emitting an i64 binary operation.
+    pub fn i64_binop<F, M>(&mut self, masm: &mut M, emit: &mut F)
+    where
+        F: FnMut(&mut M, RegImm, RegImm, OperandSize),
+        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 reg = self.pop_to_reg(masm, OperandSize::S64);
+            emit(masm, RegImm::reg(reg), RegImm::imm(val), OperandSize::S64);
+            self.stack.push(Val::reg(reg));
+        } else {
+            let src = self.pop_to_reg(masm, OperandSize::S64);
+            let dst = self.pop_to_reg(masm, OperandSize::S64);
+            emit(masm, dst.into(), src.into(), OperandSize::S64);
+            self.regalloc.free_gpr(src);
+            self.stack.push(Val::reg(dst));
+        }
+    }
+
+    /// Spill locals and registers to memory.
+    // TODO optimize the spill range;
+    //
+    // At any point in the program, the stack
+    // might already contain Memory entries;
+    // we could effectively ignore that range;
+    // only focusing on the range that contains
+    // spillable values.
+    fn spill<M: MacroAssembler>(
+        stack: &mut Stack,
+        regalloc: &mut RegAlloc,
+        frame: &Frame,
+        masm: &mut M,
+    ) {
+        stack.inner_mut().iter_mut().for_each(|v| match v {
+            Val::Reg(r) => {
+                let offset = masm.push(*r);
+                regalloc.free_gpr(*r);
+                *v = Val::Memory(offset);
+            }
+            Val::Local(index) => {
+                let slot = frame.get_local(*index).expect("valid local at slot");
+                let addr = masm.local_address(&slot);
+                masm.load(addr, regalloc.scratch, slot.ty.into());
+                let offset = masm.push(regalloc.scratch);
+                *v = Val::Memory(offset);
+            }
+            v => {
+                println!("trying to spill something unknown {:?}", v);
+            }
+        });
+    }
+}

winch/codegen/src/codegen.rs renamed to winch/codegen/src/codegen/mod.rs

@@ -1,89 +1,12 @@
 use crate::{
     abi::{ABISig, ABI},
-    frame::Frame,
-    masm::{MacroAssembler, OperandSize, RegImm},
-    regalloc::RegAlloc,
-    stack::{Stack, Val},
+    masm::{MacroAssembler, OperandSize},
 };
 use anyhow::Result;
 use wasmparser::{BinaryReader, FuncValidator, ValType, ValidatorResources, VisitOperator};
 
-/// The code generation context.
-pub(crate) struct CodeGenContext<'a, M>
-where
-    M: MacroAssembler,
-{
-    pub masm: &'a mut M,
-    pub stack: Stack,
-    pub frame: &'a Frame,
-}
-
-impl<'a, M> CodeGenContext<'a, M>
-where
-    M: MacroAssembler,
-{
-    /// Create a new code generation context.
-    pub fn new(masm: &'a mut M, stack: Stack, frame: &'a Frame) -> Self {
-        Self { masm, stack, frame }
-    }
-
-    /// Prepares arguments for emitting an i32 binary operation.
-    pub fn i32_binop<F>(&mut self, regalloc: &mut RegAlloc, emit: &mut F)
-    where
-        F: FnMut(&mut M, RegImm, RegImm, OperandSize),
-    {
-        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 reg = regalloc.pop_to_reg(self, OperandSize::S32);
-            emit(
-                &mut self.masm,
-                RegImm::reg(reg),
-                RegImm::imm(val as i64),
-                OperandSize::S32,
-            );
-            self.stack.push(Val::reg(reg));
-        } else {
-            let src = regalloc.pop_to_reg(self, OperandSize::S32);
-            let dst = regalloc.pop_to_reg(self, OperandSize::S32);
-            emit(&mut self.masm, dst.into(), src.into(), OperandSize::S32);
-            regalloc.free_gpr(src);
-            self.stack.push(Val::reg(dst));
-        }
-    }
-
-    /// Prepares arguments for emitting an i64 binary operation.
-    pub fn i64_binop<F>(&mut self, regalloc: &mut RegAlloc, emit: &mut F)
-    where
-        F: FnMut(&mut M, RegImm, RegImm, OperandSize),
-    {
-        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 reg = regalloc.pop_to_reg(self, OperandSize::S64);
-            emit(
-                &mut self.masm,
-                RegImm::reg(reg),
-                RegImm::imm(val),
-                OperandSize::S64,
-            );
-            self.stack.push(Val::reg(reg));
-        } else {
-            let src = regalloc.pop_to_reg(self, OperandSize::S64);
-            let dst = regalloc.pop_to_reg(self, OperandSize::S64);
-            emit(&mut self.masm, dst.into(), src.into(), OperandSize::S64);
-            regalloc.free_gpr(src);
-            self.stack.push(Val::reg(dst));
-        }
-    }
-}
+mod context;
+pub(crate) use context::*;
 
 /// The code generation abstraction.
 pub(crate) struct CodeGen<'a, M>
@@ -97,22 +20,22 @@ where
     sig: ABISig,
 
     /// The code generation context.
-    pub context: CodeGenContext<'a, M>,
+    pub context: CodeGenContext<'a>,
 
-    /// The register allocator.
-    pub regalloc: RegAlloc,
+    /// The MacroAssembler.
+    pub masm: &'a mut M,
 }
 
 impl<'a, M> CodeGen<'a, M>
 where
     M: MacroAssembler,
 {
-    pub fn new<A: ABI>(context: CodeGenContext<'a, M>, sig: ABISig, regalloc: RegAlloc) -> Self {
+    pub fn new<A: ABI>(masm: &'a mut M, context: CodeGenContext<'a>, sig: ABISig) -> Self {
         Self {
             word_size: <A as ABI>::word_bytes(),
             sig,
             context,
-            regalloc,
+            masm,
         }
     }
 
@@ -131,10 +54,8 @@ where
 
     // TODO stack checks
     fn emit_start(&mut self) -> Result<()> {
-        self.context.masm.prologue();
-        self.context
-            .masm
-            .reserve_stack(self.context.frame.locals_size);
+        self.masm.prologue();
+        self.masm.reserve_stack(self.context.frame.locals_size);
         Ok(())
     }
 
@@ -145,10 +66,10 @@ where
     ) -> Result<()> {
         self.spill_register_arguments();
         let defined_locals_range = &self.context.frame.defined_locals_range;
-        self.context.masm.zero_mem_range(
+        self.masm.zero_mem_range(
             defined_locals_range.as_range(),
             self.word_size,
-            &mut self.regalloc,
+            &mut self.context.regalloc,
         );
 
         while !body.eof() {
@@ -185,7 +106,7 @@ where
     // Emit the usual function end instruction sequence.
     pub fn emit_end(&mut self) -> Result<()> {
         self.handle_abi_result();
-        self.context.masm.epilogue(self.context.frame.locals_size);
+        self.masm.epilogue(self.context.frame.locals_size);
         Ok(())
     }
 
@@ -206,14 +127,14 @@ where
                     .frame
                     .get_local(index as u32)
                     .expect("valid local slot at location");
-                let addr = self.context.masm.local_address(local);
+                let addr = self.masm.local_address(local);
                 let src = arg
                     .get_reg()
                     .expect("arg should be associated to a register");
 
                 match &ty {
-                    ValType::I32 => self.context.masm.store(src.into(), addr, OperandSize::S32),
-                    ValType::I64 => self.context.masm.store(src.into(), addr, OperandSize::S64),
+                    ValType::I32 => self.masm.store(src.into(), addr, OperandSize::S32),
+                    ValType::I64 => self.masm.store(src.into(), addr, OperandSize::S64),
                     _ => panic!("Unsupported type {:?}", ty),
                 }
             });
@@ -225,8 +146,8 @@ where
         }
         let named_reg = self.sig.result.result_reg();
         let reg = self
-            .regalloc
-            .pop_to_named_reg(&mut self.context, named_reg, OperandSize::S64);
-        self.regalloc.free_gpr(reg);
+            .context
+            .pop_to_named_reg(self.masm, named_reg, OperandSize::S64);
+        self.context.regalloc.free_gpr(reg);
     }
 }