Cranelift: return programmatic error rather than panic when temporaries run out.

cranelift/codegen/src/machinst/abi.rs

@@ -1545,7 +1545,7 @@ impl<M: ABIMachineSpec> Callee<M> {
                 // We need to dereference the pointer.
                 let base = match &pointer {
                     &ABIArgSlot::Reg { reg, ty, .. } => {
-                        let tmp = vregs.alloc(ty).unwrap().only_reg().unwrap();
+                        let tmp = vregs.alloc_with_deferred_error(ty).only_reg().unwrap();
                         self.reg_args.push(ArgPair {
                             vreg: Writable::from_reg(tmp),
                             preg: reg.into(),
@@ -1607,9 +1607,10 @@ impl<M: ABIMachineSpec> Callee<M> {
                                     if n < word_bits =>
                                 {
                                     let signed = ext == ir::ArgumentExtension::Sext;
-                                    let dst = writable_value_regs(vregs.alloc(ty).unwrap())
-                                        .only_reg()
-                                        .unwrap();
+                                    let dst =
+                                        writable_value_regs(vregs.alloc_with_deferred_error(ty))
+                                            .only_reg()
+                                            .unwrap();
                                     ret.push(M::gen_extend(
                                         dst, from_reg, signed, from_bits,
                                         /* to_bits = */ word_bits,
@@ -1650,9 +1651,10 @@ impl<M: ABIMachineSpec> Callee<M> {
                                 {
                                     assert_eq!(M::word_reg_class(), from_reg.class());
                                     let signed = ext == ir::ArgumentExtension::Sext;
-                                    let dst = writable_value_regs(vregs.alloc(ty).unwrap())
-                                        .only_reg()
-                                        .unwrap();
+                                    let dst =
+                                        writable_value_regs(vregs.alloc_with_deferred_error(ty))
+                                            .only_reg()
+                                            .unwrap();
                                     ret.push(M::gen_extend(
                                         dst, from_reg, signed, from_bits,
                                         /* to_bits = */ word_bits,

cranelift/codegen/src/machinst/lower.rs

@@ -1075,6 +1075,12 @@ impl<'func, I: VCodeInst> Lower<'func, I> {
             }
 
             self.finish_bb();
+
+            // Check for any deferred vreg-temp allocation errors, and
+            // bubble one up at this time if it exists.
+            if let Some(e) = self.vregs.take_deferred_error() {
+                return Err(e);
+            }
         }
 
         // Now that we've emitted all instructions into the
@@ -1325,7 +1331,7 @@ impl<'func, I: VCodeInst> Lower<'func, I> {
 impl<'func, I: VCodeInst> Lower<'func, I> {
     /// Get a new temp.
     pub fn alloc_tmp(&mut self, ty: Type) -> ValueRegs<Writable<Reg>> {
-        writable_value_regs(self.vregs.alloc(ty).unwrap())
+        writable_value_regs(self.vregs.alloc_with_deferred_error(ty))
     }
 
     /// Emit a machine instruction.

cranelift/codegen/src/machinst/vcode.rs

@@ -1454,6 +1454,12 @@ pub struct VRegAllocator<I> {
     /// reftype-status of each vreg) for efficient iteration.
     reftyped_vregs: Vec<VReg>,
 
+    /// A deferred error, to be bubbled up to the top level of the
+    /// lowering algorithm. We take this approach because we cannot
+    /// currently propagate a `Result` upward through ISLE code (the
+    /// lowering rules) or some ABI code.
+    deferred_error: Option<CodegenError>,
+
     /// The type of instruction that this allocator makes registers for.
     _inst: core::marker::PhantomData<I>,
 }
@@ -1466,12 +1472,16 @@ impl<I: VCodeInst> VRegAllocator<I> {
             vreg_types: vec![],
             reftyped_vregs_set: FxHashSet::default(),
             reftyped_vregs: vec![],
+            deferred_error: None,
             _inst: core::marker::PhantomData::default(),
         }
     }
 
     /// Allocate a fresh ValueRegs.
     pub fn alloc(&mut self, ty: Type) -> CodegenResult<ValueRegs<Reg>> {
+        if self.deferred_error.is_some() {
+            return Err(CodegenError::CodeTooLarge);
+        }
         let v = self.next_vreg;
         let (regclasses, tys) = I::rc_for_type(ty)?;
         self.next_vreg += regclasses.len();
@@ -1493,6 +1503,38 @@ impl<I: VCodeInst> VRegAllocator<I> {
         Ok(regs)
     }
 
+    /// Allocate a fresh ValueRegs, deferring any out-of-vregs
+    /// errors. This is useful in places where we cannot bubble a
+    /// `CodegenResult` upward easily, and which are known to be
+    /// invoked from within the lowering loop that checks the deferred
+    /// error status below.
+    pub fn alloc_with_deferred_error(&mut self, ty: Type) -> ValueRegs<Reg> {
+        match self.alloc(ty) {
+            Ok(x) => x,
+            Err(e) => {
+                self.deferred_error = Some(e);
+                self.bogus_for_deferred_error(ty)
+            }
+        }
+    }
+
+    /// Take any deferred error that was accumulated by `alloc_with_deferred_error`.
+    pub fn take_deferred_error(&mut self) -> Option<CodegenError> {
+        self.deferred_error.take()
+    }
+
+    /// Produce an bogus VReg placeholder with the proper number of
+    /// registers for the given type. This is meant to be used with
+    /// deferred allocation errors (see `Lower::alloc_tmp()`).
+    fn bogus_for_deferred_error(&self, ty: Type) -> ValueRegs<Reg> {
+        let (regclasses, _tys) = I::rc_for_type(ty).expect("must have valid type");
+        match regclasses {
+            &[rc0] => ValueRegs::one(VReg::new(0, rc0).into()),
+            &[rc0, rc1] => ValueRegs::two(VReg::new(0, rc0).into(), VReg::new(1, rc1).into()),
+            _ => panic!("Value must reside in 1 or 2 registers"),
+        }
+    }
+
     /// Set the type of this virtual register.
     pub fn set_vreg_type(&mut self, vreg: VirtualReg, ty: Type) {
         if self.vreg_types.len() <= vreg.index() {

tests/all/code_too_large.rs

@@ -0,0 +1,36 @@
+#![cfg(not(miri))]
+
+use anyhow::Result;
+use wasmtime::*;
+
+#[test]
+fn code_too_large_without_panic() -> Result<()> {
+    const N: usize = 120000;
+
+    // Build a module with a function whose body will allocate too many
+    // temporaries for our current (Cranelift-based) compiler backend to
+    // handle. This test ensures that we propagate the failure upward
+    // and return it programmatically, rather than panic'ing. If we ever
+    // improve our compiler backend to actually handle such a large
+    // function body, we'll need to increase the limits here too!
+    let mut s = String::new();
+    s.push_str("(module\n");
+    s.push_str("(table 1 1 funcref)\n");
+    s.push_str("(func (export \"\") (result i32)\n");
+    s.push_str("i32.const 0\n");
+    for _ in 0..N {
+        s.push_str("table.get 0\n");
+        s.push_str("ref.is_null\n");
+    }
+    s.push_str("))\n");
+
+    let store = Store::<()>::default();
+    let result = Module::new(store.engine(), &s);
+    match result {
+        Err(e) => assert!(e
+            .to_string()
+            .starts_with("Compilation error: Code for function is too large")),
+        Ok(_) => panic!("Please adjust limits to make the module too large to compile!"),
+    }
+    Ok(())
+}

tests/all/main.rs

@@ -3,6 +3,7 @@
 mod async_functions;
 mod call_hook;
 mod cli_tests;
+mod code_too_large;
 mod component_model;
 mod coredump;
 mod custom_signal_handler;