winch: Use aarch64 backend for code emission. (#5652)

This patch introduces basic aarch64 code generation by using
`cranelift-codegen`'s backend.

This commit *does not*:

* Change the semantics of the code generation
* Adds support for other Wasm instructions

The most notable change in this patch is how addressing modes are handled at the
MacroAssembler layer: instead of having a canonical address representation, this
patch introduces the addressing mode as an associated type in the
MacroAssembler trait. This approach has the advantage that gives each ISA enough
flexiblity to describe the addressing modes and their constraints in isolation
without having to worry on how a particular addressing mode is going to affect
other ISAs. In the case of Aarch64 this becomes useful to describe indexed
addressing modes (particularly from the stack pointer).

This patch uses the concept of a shadow stack pointer (x28) as a workaround to
Aarch64's stack pointer 16-byte alignment. This constraint is enforced by:

* Introducing specialized addressing modes when using the real stack pointer; this
enables auditing when the real stack pointer is used. As of this change, the
real stack pointer is only used in the function's prologue and epilogue.

* Asserting that the real stack pointer is not used as a base for addressing
modes.

* Ensuring that at any point during the code generation process where the stack
pointer changes (e.g. when stack space is allocated / deallocated) the value of
the real stack pointer is copied into the shadow stack pointer.

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

@@ -1,57 +1,180 @@
+use super::{
+    address::Address,
+    asm::{Assembler, Operand},
+    regs,
+};
 use crate::{
-    abi::{addressing_mode::Address, local::LocalSlot},
+    abi::local::LocalSlot,
     isa::reg::Reg,
     masm::{MacroAssembler as Masm, OperandSize, RegImm},
 };
-use cranelift_codegen::{Final, MachBufferFinalized};
+use cranelift_codegen::{settings, Final, MachBufferFinalized};
 
-#[derive(Default)]
-pub(crate) struct MacroAssembler;
+/// Aarch64 MacroAssembler.
+pub(crate) struct MacroAssembler {
+    /// Low level assembler.
+    asm: Assembler,
+    /// Stack pointer offset.
+    sp_offset: u32,
+}
+
+// Conversions between generic masm arguments and aarch64 operands.
+
+impl From<RegImm> for Operand {
+    fn from(rimm: RegImm) -> Self {
+        match rimm {
+            RegImm::Reg(r) => r.into(),
+            RegImm::Imm(imm) => Operand::Imm(imm),
+        }
+    }
+}
+
+impl From<Reg> for Operand {
+    fn from(reg: Reg) -> Self {
+        Operand::Reg(reg)
+    }
+}
+
+impl From<Address> for Operand {
+    fn from(addr: Address) -> Self {
+        Operand::Mem(addr)
+    }
+}
+
+impl MacroAssembler {
+    /// Create an Aarch64 MacroAssembler.
+    pub fn new(shared_flags: settings::Flags) -> Self {
+        Self {
+            asm: Assembler::new(shared_flags),
+            sp_offset: 0u32,
+        }
+    }
+}
 
 impl Masm for MacroAssembler {
+    type Address = Address;
+
     fn prologue(&mut self) {
-        todo!()
+        let lr = regs::lr();
+        let fp = regs::fp();
+        let sp = regs::sp();
+        let addr = Address::pre_indexed_from_sp(-16);
+
+        self.asm.stp(fp, lr, addr);
+        self.asm.mov_rr(sp, fp, OperandSize::S64);
+        self.move_sp_to_shadow_sp();
     }
 
-    fn epilogue(&mut self, _locals_size: u32) {
-        todo!()
+    fn epilogue(&mut self, locals_size: u32) {
+        assert!(self.sp_offset == locals_size);
+
+        let sp = regs::sp();
+        if locals_size > 0 {
+            self.asm
+                .add_ir(locals_size as u64, sp, sp, OperandSize::S64);
+            self.move_sp_to_shadow_sp();
+        }
+
+        let lr = regs::lr();
+        let fp = regs::fp();
+        let addr = Address::post_indexed_from_sp(16);
+
+        self.asm.ldp(fp, lr, addr);
+        self.asm.ret();
     }
 
-    fn reserve_stack(&mut self, _bytes: u32) {
-        todo!()
+    fn reserve_stack(&mut self, bytes: u32) {
+        if bytes == 0 {
+            return;
+        }
+
+        let sp = regs::sp();
+        self.asm.sub_ir(bytes as u64, sp, sp, OperandSize::S64);
+        self.move_sp_to_shadow_sp();
+
+        self.increment_sp(bytes);
     }
 
-    fn local_address(&mut self, _local: &LocalSlot) -> Address {
-        todo!()
+    fn local_address(&mut self, local: &LocalSlot) -> Address {
+        let (reg, offset) = local
+            .addressed_from_sp()
+            .then(|| {
+                let offset = self.sp_offset.checked_sub(local.offset).expect(&format!(
+                    "Invalid local offset = {}; sp offset = {}",
+                    local.offset, self.sp_offset
+                ));
+                (regs::shadow_sp(), offset)
+            })
+            .unwrap_or((regs::fp(), local.offset));
+
+        Address::offset(reg, offset as i64)
     }
 
-    fn store(&mut self, _src: RegImm, _dst: Address, _size: OperandSize) {
-        todo!()
+    fn store(&mut self, src: RegImm, dst: Address, size: OperandSize) {
+        let src = match src {
+            RegImm::Imm(imm) => {
+                let scratch = regs::scratch();
+                self.asm.load_constant(imm as u64, scratch);
+                scratch
+            }
+            RegImm::Reg(reg) => reg,
+        };
+
+        self.asm.str(src, dst, size);
     }
 
-    fn load(&mut self, _src: Address, _dst: Reg, _size: OperandSize) {}
+    fn load(&mut self, src: Address, dst: Reg, size: OperandSize) {
+        self.asm.ldr(src, dst, size);
+    }
 
     fn sp_offset(&mut self) -> u32 {
-        0u32
+        self.sp_offset
     }
 
     fn finalize(self) -> MachBufferFinalized<Final> {
-        todo!()
+        self.asm.finalize()
     }
 
-    fn mov(&mut self, _src: RegImm, _dst: RegImm, _size: OperandSize) {
-        todo!()
+    fn mov(&mut self, src: RegImm, dst: RegImm, size: OperandSize) {
+        self.asm.mov(src.into(), dst.into(), size);
     }
 
-    fn add(&mut self, _dst: RegImm, __lhs: RegImm, __rhs: RegImm, _size: OperandSize) {
-        todo!()
+    fn add(&mut self, dst: RegImm, lhs: RegImm, rhs: RegImm, size: OperandSize) {
+        self.asm.add(rhs.into(), lhs.into(), dst.into(), size);
     }
 
-    fn zero(&mut self, _reg: Reg) {
-        todo!()
+    fn zero(&mut self, reg: Reg) {
+        self.asm.load_constant(0, reg);
     }
 
-    fn push(&mut self, _reg: Reg) -> u32 {
-        todo!()
+    fn push(&mut self, reg: Reg) -> u32 {
+        // The push is counted as pushing the 64-bit width in
+        // 64-bit architectures.
+        let size = 8u32;
+        self.reserve_stack(size);
+        let address = Address::from_shadow_sp(size as i64);
+        self.asm.str(reg, address, OperandSize::S64);
+
+        self.sp_offset
+    }
+}
+
+impl MacroAssembler {
+    fn increment_sp(&mut self, bytes: u32) {
+        self.sp_offset += bytes;
+    }
+
+    // Copies the value of the stack pointer to the shadow stack
+    // pointer: mov x28, sp
+
+    // This function is usually called whenever the real stack pointer
+    // changes, for example after allocating or deallocating stack
+    // space, or after performing a push or pop.
+    // For more details around the stack pointer and shadow stack
+    // pointer see the docs at regs::shadow_sp().
+    fn move_sp_to_shadow_sp(&mut self) {
+        let sp = regs::sp();
+        let shadow_sp = regs::shadow_sp();
+        self.asm.mov_rr(sp, shadow_sp, OperandSize::S64);
     }
 }