Address review comments

* Fix typos
* Better documentation for `new_unchecked`
* Introduce `max` for `BitSet`
* Make allocatable property `u64`

winch/codegen/src/abi/mod.rs

@@ -72,7 +72,7 @@ pub(crate) trait ABI {
     /// function type.
     fn sig(wasm_sig: &WasmFuncType, call_conv: &CallingConvention) -> ABISig;
 
-    /// Construct an ABI siganture from WasmType params and returns.
+    /// Construct an ABI signature from WasmType params and returns.
     fn sig_from(params: &[WasmType], returns: &[WasmType], call_conv: &CallingConvention)
         -> ABISig;
 

winch/codegen/src/codegen/call.rs

@@ -91,16 +91,25 @@ impl<'a> FnCall<'a> {
             call_stack_space: None,
         };
         // When all the information is known upfront, we can optimize and
-        // calcualte the stack space needed by the call right away.
+        // calculate the stack space needed by the call right away.
         call.save_live_registers(context, masm);
         call
     }
 
-    /// Creates a new [`FnCall`] from an [`ABISIg`] without checking if the
-    /// stack has the correct amount of values for the call.
-    /// This happens in situations in which not all the information is known
-    /// upfront in order to fulfill the call, like for example with dealing
-    /// with libcalls.
+    /// Creates a new [`FnCall`] from an [`ABISIg`] without checking if the stack has the correct
+    /// amount of values for the call. This happens in situations in which not all the information
+    /// is known upfront in order to fulfill the call, like for example with dealing with libcalls.
+    /// Libcalls don't necessarily match 1-to-1 to WebAssembly instructions, so it's possible that
+    /// before emittiing the libcall we need to adjust the value stack with the right values to be
+    /// used as parameters. We can't preemptively adjust the stack since in some cases we might
+    /// need to ensure that the stack is balanced right until we emit the function call because
+    /// there's a dependency on certain values on the stack. A good example of this is when lazily
+    /// initializing funcrefs: in order to correctly get the value of a function reference we need
+    /// to determine if a libcall is needed, in order to do so we preemptively prepare the compiler
+    /// to emit one since we can't know ahead-of-time of time if one will be required or not. That
+    /// is the main reason why this function is unchecked, it's the caller's responsibility to
+    /// ensure -- depending on the libcall -- that the value is correctly balanaced before the
+    /// call.
     pub fn new_unchecked(sig: &'a ABISig) -> Self {
         Self {
             abi_sig: sig,

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

@@ -1,4 +1,4 @@
-use self::regs::{ALL_GPR, NON_ALLOCATABLE_GPR};
+use self::regs::{ALL_GPR, MAX_FPR, MAX_GPR, NON_ALLOCATABLE_GPR};
 use crate::{
     abi::ABI,
     codegen::{CodeGen, CodeGenContext, FuncEnv},
@@ -97,9 +97,13 @@ impl TargetIsa for Aarch64 {
 
         let defined_locals = DefinedLocals::new(translation, &mut body, validator)?;
         let frame = Frame::new::<abi::Aarch64ABI>(&abi_sig, &defined_locals)?;
-        let gpr = RegBitSet::int(ALL_GPR, NON_ALLOCATABLE_GPR);
+        let gpr = RegBitSet::int(
+            ALL_GPR.into(),
+            NON_ALLOCATABLE_GPR.into(),
+            usize::try_from(MAX_GPR).unwrap(),
+        );
         // TODO: Add floating point bitmask
-        let fpr = RegBitSet::float(0, 0);
+        let fpr = RegBitSet::float(0, 0, usize::try_from(MAX_FPR).unwrap());
         let regalloc = RegAlloc::from(gpr, fpr);
         let codegen_context = CodeGenContext::new(regalloc, stack, &frame);
         let env = FuncEnv::new(

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

@@ -4,15 +4,20 @@ use crate::{isa::reg::Reg, masm::OperandSize};
 use regalloc2::{PReg, RegClass};
 use smallvec::{smallvec, SmallVec};
 
+/// FPR index bound.
+pub(crate) const MAX_FPR: u32 = 32;
+/// FPR index bound.
+pub(crate) const MAX_GPR: u32 = 32;
+
 /// Construct a X-register from an index.
 pub(crate) const fn xreg(num: u8) -> Reg {
-    assert!(num < 32);
+    assert!((num as u32) < MAX_GPR);
     Reg::new(PReg::new(num as usize, RegClass::Int))
 }
 
 /// Construct a V-register from an index.
 pub(crate) const fn vreg(num: u8) -> Reg {
-    assert!(num < 32);
+    assert!((num as u32) < MAX_FPR);
     Reg::new(PReg::new(num as usize, RegClass::Float))
 }
 

winch/codegen/src/isa/reg.rs

@@ -34,8 +34,8 @@ impl Reg {
     }
 
     /// Get the encoding of the underlying register.
-    pub const fn hw_enc(self) -> u8 {
-        self.0.hw_enc() as u8
+    pub const fn hw_enc(self) -> usize {
+        self.0.hw_enc()
     }
 
     /// Get the physical register representation.

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

@@ -21,7 +21,7 @@ use target_lexicon::Triple;
 use wasmparser::{FuncValidator, FunctionBody, ValidatorResources};
 use wasmtime_environ::{ModuleTranslation, ModuleTypes, WasmFuncType};
 
-use self::regs::{ALL_FPR, ALL_GPR, NON_ALLOCATABLE_FPR, NON_ALLOCATABLE_GPR};
+use self::regs::{ALL_FPR, ALL_GPR, MAX_FPR, MAX_GPR, NON_ALLOCATABLE_FPR, NON_ALLOCATABLE_GPR};
 
 mod abi;
 mod address;
@@ -106,8 +106,16 @@ impl TargetIsa for X64 {
 
         let defined_locals = DefinedLocals::new(translation, &mut body, validator)?;
         let frame = Frame::new::<abi::X64ABI>(&abi_sig, &defined_locals)?;
-        let gpr = RegBitSet::int(ALL_GPR, NON_ALLOCATABLE_GPR);
-        let fpr = RegBitSet::float(ALL_FPR, NON_ALLOCATABLE_FPR);
+        let gpr = RegBitSet::int(
+            ALL_GPR.into(),
+            NON_ALLOCATABLE_GPR.into(),
+            usize::try_from(MAX_GPR).unwrap(),
+        );
+        let fpr = RegBitSet::float(
+            ALL_FPR.into(),
+            NON_ALLOCATABLE_FPR.into(),
+            usize::try_from(MAX_FPR).unwrap(),
+        );
 
         let regalloc = RegAlloc::from(gpr, fpr);
         let codegen_context = CodeGenContext::new(regalloc, stack, &frame);

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

@@ -174,6 +174,10 @@ pub(crate) fn scratch_xmm() -> Reg {
 const GPR: u32 = 16;
 /// FPR count.
 const FPR: u32 = 16;
+/// GPR index bound.
+pub(crate) const MAX_GPR: u32 = GPR;
+/// GPR index bound.
+pub(crate) const MAX_FPR: u32 = FPR;
 const ALLOCATABLE_GPR: u32 = (1 << GPR) - 1;
 const ALLOCATABLE_FPR: u32 = (1 << FPR) - 1;
 /// Bitmask of non-alloctable GPRs.

winch/codegen/src/regset.rs

@@ -37,31 +37,38 @@ pub struct RegBitSet {
     /// The register class.
     class: RegClass,
     /// The set of allocatable
-    allocatable: u32,
+    allocatable: u64,
     /// The set of non-alloctable registers.
-    non_allocatable: u32,
+    non_allocatable: u64,
+    /// The max number of registers.
+    /// Invariant:
+    /// When allocating or freeing a register the encoding (index) of the
+    /// register must be less than the max property.
+    max: usize,
 }
 
 impl RegBitSet {
     /// Creates an integer register class bitset.
-    pub fn int(allocatable: u32, non_allocatable: u32) -> Self {
+    pub fn int(allocatable: u64, non_allocatable: u64, max: usize) -> Self {
         // Assert that one set is the complement of the other.
         debug_assert!(allocatable & non_allocatable == 0);
         Self {
             class: RegClass::Int,
             allocatable,
             non_allocatable,
+            max,
         }
     }
 
     /// Creates a float register class bitset.
-    pub fn float(allocatable: u32, non_allocatable: u32) -> Self {
+    pub fn float(allocatable: u64, non_allocatable: u64, max: usize) -> Self {
         // Assert that one set is the complement of the other.
         debug_assert!(allocatable & non_allocatable == 0);
         Self {
             class: RegClass::Float,
             allocatable,
             non_allocatable,
+            max,
         }
     }
 }
@@ -81,7 +88,7 @@ impl RegSet {
         self.available(class).then(|| {
             let bitset = &self[class];
             let index = bitset.allocatable.trailing_zeros();
-            self.allocate(class, index);
+            self.allocate(class, index.into());
             Reg::from(class, index as usize)
         })
     }
@@ -90,15 +97,18 @@ impl RegSet {
     pub fn reg(&mut self, reg: Reg) -> Option<Reg> {
         let index = reg.hw_enc();
         self.named_reg_available(reg).then(|| {
-            self.allocate(reg.class(), index.into());
+            self.allocate(reg.class(), index.try_into().unwrap());
             reg
         })
     }
 
     /// Marks the specified register as available, utilizing the
     /// register class to determine the bitset that requires updating.
     pub fn free(&mut self, reg: Reg) {
-        let index = reg.hw_enc() as u32;
+        let bitset = &self[reg.class()];
+        let index = reg.hw_enc();
+        assert!(index < bitset.max);
+        let index = u64::try_from(index).unwrap();
         if !self.is_non_allocatable(reg.class(), index) {
             self[reg.class()].allocatable |= 1 << index;
         }
@@ -107,23 +117,25 @@ impl RegSet {
     /// Returns true if the specified register is allocatable.
     pub fn named_reg_available(&self, reg: Reg) -> bool {
         let bitset = &self[reg.class()];
+        assert!(reg.hw_enc() < bitset.max);
         let index = 1 << reg.hw_enc();
+
         (!bitset.allocatable & index) == 0
-            || self.is_non_allocatable(reg.class(), reg.hw_enc().into())
+            || self.is_non_allocatable(reg.class(), reg.hw_enc().try_into().unwrap())
     }
 
     fn available(&self, class: RegClass) -> bool {
         let bitset = &self[class];
         bitset.allocatable != 0
     }
 
-    fn allocate(&mut self, class: RegClass, index: u32) {
+    fn allocate(&mut self, class: RegClass, index: u64) {
         if !self.is_non_allocatable(class, index) {
             self[class].allocatable &= !(1 << index);
         }
     }
 
-    fn is_non_allocatable(&self, class: RegClass, index: u32) -> bool {
+    fn is_non_allocatable(&self, class: RegClass, index: u64) -> bool {
         let bitset = &self[class];
         let non_allocatable = bitset.non_allocatable;
         non_allocatable != 0 && !non_allocatable & (1 << index) == 0

winch/codegen/src/visitor.rs

@@ -639,7 +639,7 @@ where
             &mut self.context,
             &builtin,
             |cx, masm, call, callee| {
-                // Calcualte the table element address.
+                // Calculate the table element address.
                 let index = cx.pop_to_reg(masm, None);
                 let elem_addr =
                     masm.table_elem_address(index.into(), index.ty.into(), &table_data, cx);