winch: Simplify constant handling, part 2/N (#10989)

* winch: Introduce register allocation for scratch registers

This commit introduces a register allocator for scratch registers. The
objective of this change is to make it generally safer to work with
scratch registers and prevent accidental clobbering of said registers.

This approach also has the advantage that allows for a more natural
abstraction over ISA-dependent scratch register definitions, e.g., we
can easily encode that fact that in aarch64 x16 and x17 are considered
scratch registers, while in x64 Winch's ABI defines a single global
scratch register.

* winch: Hook up the scratch allocator in each of the backends

This commit makes use of the scratch register allocator in both the
x64 and aarch64 backends for immediate value loading.

Given that the MacroAssembler is the boundary between ISA-agnostic
code and ISA-dependent code, it seems to be the natural location for
this allocator to live.

The allocator gives exclusive access to a scratch register of a
particular class, through the `Masm::with_scratch` method. Note that
the semantics of this allocator don't involve spilling or any other
form of register availability resolution. If a register is requested
and it's not available, this method will panic.

* Apply `cargo fmt`

* Update disassembly tests

Even though the entire change doens't contain major funcitonal
changes, a side effect of improving the constant handling in aarch64
is that we perform better instruction selection for instructions that
deal with immediattes, improving the generated code in some cases.

* Review edits

Author: saulecabrera

tests/disas/winch/aarch64/br_table/large.wat

@@ -757,8 +757,8 @@
 ;;       stur    x1, [x28, #8]
 ;;       stur    w2, [x28, #4]
 ;;       ldur    w0, [x28, #4]
-;;       mov     x16, #0x6027
-;;       cmp     x0, x16, uxtx
+;;       mov     x1, #0x6027
+;;       cmp     w0, w1, uxtx
 ;;       b.hs    #0x18110
 ;;   54: csel    x1, xzr, x0, hs
 ;;       csdb

tests/disas/winch/aarch64/br_table/nested_br_table_loop_block.wat

@@ -29,43 +29,45 @@
 ;;       movk    x17, #0x18
 ;;       add     x16, x16, x17
 ;;       cmp     sp, x16
-;;       b.lo    #0xc4
+;;       b.lo    #0xcc
 ;;   2c: mov     x9, x0
 ;;       sub     x28, x28, #0x18
 ;;       mov     sp, x28
 ;;       stur    x0, [x28, #0x10]
 ;;       stur    x1, [x28, #8]
 ;;       stur    w2, [x28, #4]
 ;;       ldur    w0, [x28, #4]
-;;       cmp     x0, #2
-;;       b.hs    #0x74
-;;   50: csel    x1, xzr, x0, hs
+;;       mov     x1, #2
+;;       cmp     w0, w1, uxtx
+;;       b.hs    #0x78
+;;   54: csel    x1, xzr, x0, hs
 ;;       csdb
-;;       adr     x16, #0x68
+;;       adr     x16, #0x6c
 ;;       ldrsw   x1, [x16, w1, uxtw #2]
 ;;       add     x16, x16, x1
 ;;       br      x16
-;;   68: .byte   0xdc, 0xff, 0xff, 0xff
+;;   6c: .byte   0xd8, 0xff, 0xff, 0xff
 ;;       .byte   0x0c, 0x00, 0x00, 0x00
 ;;       b       #0x44
-;;   74: mov     x0, #0
+;;   78: mov     x0, #0
 ;;       stur    w0, [x28, #4]
 ;;       ldur    w0, [x28, #4]
-;;       cmp     x0, #2
-;;       b.hs    #0x7c
-;;   88: csel    x1, xzr, x0, hs
+;;       mov     x1, #2
+;;       cmp     w0, w1, uxtx
+;;       b.hs    #0x80
+;;   90: csel    x1, xzr, x0, hs
 ;;       csdb
-;;       adr     x16, #0xa0
+;;       adr     x16, #0xa8
 ;;       ldrsw   x1, [x16, w1, uxtw #2]
 ;;       add     x16, x16, x1
 ;;       br      x16
-;;   a0: .byte   0x08, 0x00, 0x00, 0x00
-;;       .byte   0xdc, 0xff, 0xff, 0xff
+;;   a8: .byte   0x08, 0x00, 0x00, 0x00
+;;       .byte   0xd8, 0xff, 0xff, 0xff
 ;;       mov     x0, #3
 ;;       add     x28, x28, #0x18
 ;;       mov     sp, x28
 ;;       mov     sp, x28
 ;;       ldr     x28, [sp], #0x10
 ;;       ldp     x29, x30, [sp], #0x10
 ;;       ret
-;;   c4: .byte   0x1f, 0xc1, 0x00, 0x00
+;;   cc: .byte   0x1f, 0xc1, 0x00, 0x00

tests/disas/winch/aarch64/i32_rotl/16_const.wat

@@ -26,8 +26,8 @@
 ;;       stur    x0, [x28, #8]
 ;;       stur    x1, [x28]
 ;;       mov     x0, #1
-;;       neg     w0, w0
 ;;       mov     x16, #0x200
+;;       neg     w16, w16
 ;;       ror     w0, w0, w16
 ;;       add     x28, x28, #0x10
 ;;       mov     sp, x28

tests/disas/winch/aarch64/i64_rotl/16_const.wat

@@ -26,8 +26,8 @@
 ;;       stur    x0, [x28, #8]
 ;;       stur    x1, [x28]
 ;;       mov     x0, #1
-;;       neg     x0, x0
 ;;       mov     x16, #0x200
+;;       neg     x16, x16
 ;;       ror     x0, x0, x16
 ;;       add     x28, x28, #0x10
 ;;       mov     sp, x28

winch/codegen/src/codegen/context.rs

@@ -9,8 +9,8 @@ use crate::{
     frame::Frame,
     isa::reg::RegClass,
     masm::{
-        ExtractLaneKind, MacroAssembler, MemMoveDirection, OperandSize, RegImm, ReplaceLaneKind,
-        SPOffset, ShiftKind, StackSlot,
+        ExtractLaneKind, Imm, MacroAssembler, MemMoveDirection, OperandSize, RegImm,
+        ReplaceLaneKind, SPOffset, ShiftKind, StackSlot,
     },
     reg::{Reg, WritableReg, writable},
     regalloc::RegAlloc,
@@ -64,7 +64,7 @@ impl<'a> CodeGenContext<'a, Emission> {
             let typed_reg = self.pop_to_reg(masm, None)?;
             masm.shift_ir(
                 writable!(typed_reg.reg),
-                val as u64,
+                Imm::i32(val),
                 typed_reg.reg,
                 kind,
                 OperandSize::S32,
@@ -93,7 +93,7 @@ impl<'a> CodeGenContext<'a, Emission> {
             let typed_reg = self.pop_to_reg(masm, None)?;
             masm.shift_ir(
                 writable!(typed_reg.reg),
-                val as u64,
+                Imm::i64(val),
                 typed_reg.reg,
                 kind,
                 OperandSize::S64,

winch/codegen/src/codegen/mod.rs

@@ -1102,7 +1102,7 @@ where
         let pow = heap_data.memory.page_size_log2;
         self.masm.shift_ir(
             writable!(dst.reg),
-            pow as u64,
+            Imm::i32(pow as i32),
             dst.into(),
             ShiftKind::ShrU,
             heap_data.index_type().try_into()?,

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

@@ -122,12 +122,14 @@ impl ABI for Aarch64ABI {
 
     fn scratch_for(ty: &WasmValType) -> Reg {
         match ty {
+            #[expect(deprecated, reason = "Pending migration to Masm::with_scratch")]
             WasmValType::I32
             | WasmValType::I64
             | WasmValType::Ref(WasmRefType {
                 heap_type: WasmHeapType::Func,
                 ..
             }) => regs::scratch(),
+            #[expect(deprecated, reason = "Pending migration to Masm::with_scratch")]
             WasmValType::F32 | WasmValType::F64 | WasmValType::V128 => regs::float_scratch(),
             _ => unimplemented!(),
         }