Revert to 32-bit immediate offsets in heap_addr

This commit updates the generation of addresses in wasm code to always
use 32-bit offsets for `heap_addr`, and if the calculated offset is
bigger than 32-bits we emit a manual add with an overflow check.

Author: alexcrichton

cranelift/codegen/meta/src/shared/formats.rs

@@ -273,7 +273,7 @@ impl Formats {
             heap_addr: Builder::new("HeapAddr")
                 .imm(&entities.heap)
                 .value()
-                .imm(&imm.uimm64)
+                .imm(&imm.uimm32)
                 .build(),
 
             // Accessing a WebAssembly table.

cranelift/codegen/meta/src/shared/immediates.rs

@@ -14,8 +14,8 @@ pub(crate) struct Immediates {
     /// counts on shift instructions.
     pub uimm8: OperandKind,
 
-    /// An unsigned 64-bit immediate integer operand.
-    pub uimm64: OperandKind,
+    /// An unsigned 32-bit immediate integer operand.
+    pub uimm32: OperandKind,
 
     /// An unsigned 128-bit immediate integer operand.
     ///
@@ -97,8 +97,8 @@ impl Immediates {
             imm64: new_imm("imm", "ir::immediates::Imm64").with_doc("A 64-bit immediate integer."),
             uimm8: new_imm("imm", "ir::immediates::Uimm8")
                 .with_doc("An 8-bit immediate unsigned integer."),
-            uimm64: new_imm("imm", "ir::immediates::Uimm64")
-                .with_doc("A 64-bit immediate unsigned integer."),
+            uimm32: new_imm("imm", "ir::immediates::Uimm32")
+                .with_doc("A 32-bit immediate unsigned integer."),
             uimm128: new_imm("imm", "ir::Immediate")
                 .with_doc("A 128-bit immediate unsigned integer."),
             pool_constant: new_imm("constant_handle", "ir::Constant")

cranelift/codegen/meta/src/shared/instructions.rs

@@ -1534,7 +1534,7 @@ pub(crate) fn define(
 
     let H = &Operand::new("H", &entities.heap);
     let p = &Operand::new("p", HeapOffset);
-    let Size = &Operand::new("Size", &imm.uimm64).with_doc("Size in bytes");
+    let Size = &Operand::new("Size", &imm.uimm32).with_doc("Size in bytes");
 
     ig.push(
         Inst::new(

cranelift/codegen/src/ir/immediates.rs

@@ -291,6 +291,12 @@ impl From<Uimm32> for u32 {
     }
 }
 
+impl From<Uimm32> for u64 {
+    fn from(val: Uimm32) -> u64 {
+        val.0.into()
+    }
+}
+
 impl From<Uimm32> for i64 {
     fn from(val: Uimm32) -> i64 {
         i64::from(val.0)

cranelift/codegen/src/ir/instructions.rs

@@ -302,8 +302,8 @@ impl InstructionData {
             // 32-bit
             &InstructionData::UnaryIeee32 { imm, .. } => Some(DataValue::from(imm)),
             &InstructionData::HeapAddr { imm, .. } => {
-                let imm: u64 = imm.into();
-                Some(DataValue::from(imm as i64)) // Note the switch from unsigned to signed.
+                let imm: u32 = imm.into();
+                Some(DataValue::from(imm as i32)) // Note the switch from unsigned to signed.
             }
             &InstructionData::Load { offset, .. }
             | &InstructionData::LoadComplex { offset, .. }

cranelift/codegen/src/legalizer/heap.rs

@@ -25,7 +25,7 @@ pub fn expand_heap_addr(
             imm,
         } => {
             debug_assert_eq!(opcode, ir::Opcode::HeapAddr);
-            (heap, arg, imm.into())
+            (heap, arg, u64::from(imm))
         }
         _ => panic!("Wanted heap_addr: {}", func.dfg.display_inst(inst, None)),
     };

cranelift/reader/src/parser.rs

@@ -3064,7 +3064,7 @@ impl<'a> Parser<'a> {
                 self.match_token(Token::Comma, "expected ',' between operands")?;
                 let arg = self.match_value("expected SSA value heap address")?;
                 self.match_token(Token::Comma, "expected ',' between operands")?;
-                let imm = self.match_uimm64("expected 64-bit integer size")?;
+                let imm = self.match_uimm32("expected 32-bit integer size")?;
                 InstructionData::HeapAddr {
                     opcode,
                     heap,

cranelift/wasm/src/code_translator.rs

@@ -2224,20 +2224,58 @@ fn prepare_addr<FE: FuncEnvironment + ?Sized>(
         cmp::max(memarg.offset / offset_guard_size * offset_guard_size, 1)
     };
     debug_assert!(adjusted_offset > 0); // want to bounds check at least 1 byte
-    let base = builder
-        .ins()
-        .heap_addr(environ.pointer_type(), heap, addr, adjusted_offset);
+    let (addr, offset) = match u32::try_from(adjusted_offset) {
+        // If our adjusted offset fits within a u32, then we can place the
+        // entire offset into the offset of the `heap_addr` instruction. After
+        // the `heap_addr` instruction, though, we need to factor the the offset
+        // into the returned address. This is either an immediate if the offset
+        // further fits within `i32`, or a manual add instruction otherwise.
+        //
+        // Note that native instructions take a signed offset hence the switch
+        // to i32. Note also the lack of overflow checking in the offset
+        // addition, which should be ok since if `heap_addr` passed we're
+        // guaranteed that this won't overflow.
+        Ok(offset) => {
+            let base = builder
+                .ins()
+                .heap_addr(environ.pointer_type(), heap, addr, offset);
+            match i32::try_from(memarg.offset) {
+                Ok(val) => (base, val),
+                Err(_) => {
+                    let adj = builder.ins().iadd_imm(base, i64::from(offset));
+                    (adj, 0)
+                }
+            }
+        }
 
-    // Native load/store instructions take a signed `Offset32` immediate, so adjust the base
-    // pointer if necessary.
-    let (addr, offset) = match i32::try_from(memarg.offset) {
-        Ok(val) => (base, val),
+        // If the adjusted offset doesn't fit within a u32, then this gets
+        // pessimized a fair amount. We can't pass the fixed-sized offset to
+        // the `heap_addr` instruction, so we need to fold the offset into the
+        // address itself. In doing so though we need to check for overflow
+        // because that would mean the address is out-of-bounds.
+        //
+        // Once we have the effective address, offset already folded in, then
+        // `heap_addr` is used to verify that the address is indeed in-bounds.
+        // The access size of the `heap_addr` is what we were passed in from
+        // above.
+        //
+        // Note that this is generating what's likely to be at least two
+        // branches, one for the overflow and one for the bounds check itself.
+        // For now though that should hopefully be ok since 4gb+ offsets are
+        // relatively odd/rare.
         Err(_) => {
-            // Note the switch from u64 offset to i64 here, but this should be
-            // ok because we're already guaranteed this won't overflow if we
-            // reach this point after the `heap_addr` instruction above.
-            let adj = builder.ins().iadd_imm(base, memarg.offset as i64);
-            (adj, 0)
+            let index_type = builder.func.heaps[heap].index_type;
+            let offset = builder.ins().iconst(index_type, memarg.offset as i64);
+            let (addr, overflow) = builder.ins().iadd_ifcout(addr, offset);
+            builder.ins().trapif(
+                environ.unsigned_add_overflow_condition(),
+                overflow,
+                ir::TrapCode::HeapOutOfBounds,
+            );
+            let base = builder
+                .ins()
+                .heap_addr(environ.pointer_type(), heap, addr, access_size);
+            (base, 0)
         }
     };
 

cranelift/wasm/src/environ/dummy.rs

@@ -652,6 +652,10 @@ impl<'dummy_environment> FuncEnvironment for DummyFuncEnvironment<'dummy_environ
     ) -> WasmResult<ir::Value> {
         Ok(pos.ins().iconst(I32, 0))
     }
+
+    fn unsigned_add_overflow_condition(&self) -> ir::condcodes::IntCC {
+        unimplemented!()
+    }
 }
 
 impl TargetEnvironment for DummyEnvironment {

cranelift/wasm/src/environ/spec.rs

@@ -697,6 +697,10 @@ pub trait FuncEnvironment: TargetEnvironment {
     ) -> WasmResult<()> {
         Ok(())
     }
+
+    /// Returns the target ISA's condition to check for unsigned addition
+    /// overflowing.
+    fn unsigned_add_overflow_condition(&self) -> ir::condcodes::IntCC;
 }
 
 /// An object satisfying the `ModuleEnvironment` trait can be passed as argument to the