Cranelift: Add heap_load and heap_store instructions (#5300)

* Cranelift: Define `heap_load` and `heap_store` instructions

* Cranelift: Implement interpreter support for `heap_load` and `heap_store`

* Cranelift: Add a suite runtests for `heap_{load,store}`

There are so many knobs we can twist for heaps and I wanted to exhaustively test
all of them, so I wrote a script to generate the tests. I've checked in the
script in case we want to make any changes in the future, but I don't think it
is worth adding this to CI to check that scripts are up to date or anything like
that.

* Review feedback

cranelift/codegen/src/ir/dfg.rs

@@ -4,11 +4,12 @@ use crate::entity::{self, PrimaryMap, SecondaryMap};
 use crate::ir;
 use crate::ir::builder::ReplaceBuilder;
 use crate::ir::dynamic_type::{DynamicTypeData, DynamicTypes};
+use crate::ir::immediates::HeapImmData;
 use crate::ir::instructions::{BranchInfo, CallInfo, InstructionData};
 use crate::ir::{types, ConstantData, ConstantPool, Immediate};
 use crate::ir::{
-    Block, DynamicType, FuncRef, Inst, SigRef, Signature, Type, Value, ValueLabelAssignments,
-    ValueList, ValueListPool,
+    Block, DynamicType, FuncRef, HeapImm, Inst, SigRef, Signature, Type, Value,
+    ValueLabelAssignments, ValueList, ValueListPool,
 };
 use crate::ir::{ExtFuncData, RelSourceLoc};
 use crate::packed_option::ReservedValue;
@@ -83,6 +84,9 @@ pub struct DataFlowGraph {
 
     /// Stores large immediates that otherwise will not fit on InstructionData
     pub immediates: PrimaryMap<Immediate, ConstantData>,
+
+    /// Out-of-line heap access immediates that don't fit in `InstructionData`.
+    pub heap_imms: PrimaryMap<HeapImm, HeapImmData>,
 }
 
 impl DataFlowGraph {
@@ -101,6 +105,7 @@ impl DataFlowGraph {
             values_labels: None,
             constants: ConstantPool::new(),
             immediates: PrimaryMap::new(),
+            heap_imms: PrimaryMap::new(),
         }
     }
 

cranelift/codegen/src/ir/entities.rs

@@ -394,6 +394,34 @@ impl Heap {
     }
 }
 
+/// An opaque reference to some out-of-line immediates for `heap_{load,store}`
+/// instructions.
+///
+/// These immediates are too large to store in
+/// [`InstructionData`](super::instructions::InstructionData) and therefore must
+/// be tracked separately in
+/// [`DataFlowGraph::heap_imms`](super::dfg::DataFlowGraph). `HeapImm` provides
+/// a way to reference values stored there.
+///
+/// While the order is stable, it is arbitrary.
+#[derive(Copy, Clone, PartialEq, Eq, Hash, PartialOrd, Ord)]
+#[cfg_attr(feature = "enable-serde", derive(Serialize, Deserialize))]
+pub struct HeapImm(u32);
+entity_impl!(HeapImm, "heap_imm");
+
+impl HeapImm {
+    /// Create a new `HeapImm` reference from its number.
+    ///
+    /// This method is for use by the parser.
+    pub fn with_number(n: u32) -> Option<Self> {
+        if n < u32::MAX {
+            Some(Self(n))
+        } else {
+            None
+        }
+    }
+}
+
 /// An opaque reference to a [WebAssembly
 /// table](https://developer.mozilla.org/en-US/docs/WebAssembly/Understanding_the_text_format#WebAssembly_tables).
 ///

cranelift/codegen/src/ir/immediates.rs

@@ -4,6 +4,7 @@
 //! Each type here should have a corresponding definition in the
 //! `cranelift-codegen/meta/src/shared/immediates` crate in the meta language.
 
+use crate::ir;
 use alloc::vec::Vec;
 use core::cmp::Ordering;
 use core::convert::TryFrom;
@@ -1177,6 +1178,18 @@ impl Not for Ieee64 {
     }
 }
 
+/// Out-of-line heap access immediates.
+#[derive(Copy, Clone, Debug, PartialEq, Eq, Hash)]
+#[cfg_attr(feature = "enable-serde", derive(Serialize, Deserialize))]
+pub struct HeapImmData {
+    /// The memory flags for the heap access.
+    pub flags: ir::MemFlags,
+    /// The heap being accessed.
+    pub heap: ir::Heap,
+    /// The static offset added to the heap access's index.
+    pub offset: Uimm32,
+}
+
 #[cfg(test)]
 mod tests {
     use super::*;

cranelift/codegen/src/ir/mod.rs

@@ -37,8 +37,8 @@ pub use crate::ir::constant::{ConstantData, ConstantPool};
 pub use crate::ir::dfg::{DataFlowGraph, ValueDef};
 pub use crate::ir::dynamic_type::{dynamic_to_fixed, DynamicTypeData, DynamicTypes};
 pub use crate::ir::entities::{
-    Block, Constant, DynamicStackSlot, DynamicType, FuncRef, GlobalValue, Heap, Immediate, Inst,
-    JumpTable, SigRef, StackSlot, Table, UserExternalNameRef, Value,
+    Block, Constant, DynamicStackSlot, DynamicType, FuncRef, GlobalValue, Heap, HeapImm, Immediate,
+    Inst, JumpTable, SigRef, StackSlot, Table, UserExternalNameRef, Value,
 };
 pub use crate::ir::extfunc::{
     AbiParam, ArgumentExtension, ArgumentPurpose, ExtFuncData, Signature,

cranelift/codegen/src/isa/aarch64/lower_inst.rs

@@ -139,8 +139,8 @@ pub(crate) fn lower_insn_to_regs(
             panic!("Direct stack memory access not supported; should not be used by Wasm");
         }
 
-        Opcode::HeapAddr => {
-            panic!("heap_addr should have been removed by legalization!");
+        Opcode::HeapLoad | Opcode::HeapStore | Opcode::HeapAddr => {
+            panic!("heap access instructions should have been removed by legalization!");
         }
 
         Opcode::TableAddr => {

cranelift/codegen/src/isa/s390x/lower.rs

@@ -212,8 +212,8 @@ impl LowerBackend for S390xBackend {
             Opcode::StackLoad | Opcode::StackStore => {
                 panic!("Direct stack memory access not supported; should not be used by Wasm");
             }
-            Opcode::HeapAddr => {
-                panic!("heap_addr should have been removed by legalization!");
+            Opcode::HeapLoad | Opcode::HeapStore | Opcode::HeapAddr => {
+                panic!("heap access instructions should have been removed by legalization!");
             }
             Opcode::TableAddr => {
                 panic!("table_addr should have been removed by legalization!");

cranelift/codegen/src/isa/x64/lower.rs

@@ -549,8 +549,8 @@ fn lower_insn_to_regs(
             panic!("global_value should have been removed by legalization!");
         }
 
-        Opcode::HeapAddr => {
-            panic!("heap_addr should have been removed by legalization!");
+        Opcode::HeapLoad | Opcode::HeapStore | Opcode::HeapAddr => {
+            panic!("heap access instructions should have been removed by legalization!");
         }
 
         Opcode::TableAddr => {

cranelift/codegen/src/opts.rs

@@ -9,8 +9,8 @@ pub use crate::ir::immediates::{Ieee32, Ieee64, Imm64, Offset32, Uimm32, Uimm64,
 pub use crate::ir::types::*;
 pub use crate::ir::{
     dynamic_to_fixed, AtomicRmwOp, Block, Constant, DynamicStackSlot, FuncRef, GlobalValue, Heap,
-    Immediate, InstructionImms, JumpTable, MemFlags, Opcode, StackSlot, Table, TrapCode, Type,
-    Value,
+    HeapImm, Immediate, InstructionImms, JumpTable, MemFlags, Opcode, StackSlot, Table, TrapCode,
+    Type, Value,
 };
 use crate::isle_common_prelude_methods;
 use crate::machinst::isle::*;

cranelift/codegen/src/verifier/mod.rs

@@ -64,6 +64,7 @@ use crate::entity::SparseSet;
 use crate::flowgraph::{BlockPredecessor, ControlFlowGraph};
 use crate::ir;
 use crate::ir::entities::AnyEntity;
+use crate::ir::immediates::HeapImmData;
 use crate::ir::instructions::{BranchInfo, CallInfo, InstructionFormat, ResolvedConstraint};
 use crate::ir::{
     types, ArgumentPurpose, Block, Constant, DynamicStackSlot, FuncRef, Function, GlobalValue,
@@ -678,6 +679,10 @@ impl<'a> Verifier<'a> {
             UnaryGlobalValue { global_value, .. } => {
                 self.verify_global_value(inst, global_value, errors)?;
             }
+            HeapLoad { heap_imm, .. } | HeapStore { heap_imm, .. } => {
+                let HeapImmData { heap, .. } = self.func.dfg.heap_imms[heap_imm];
+                self.verify_heap(inst, heap, errors)?;
+            }
             HeapAddr { heap, .. } => {
                 self.verify_heap(inst, heap, errors)?;
             }

cranelift/codegen/src/write.rs

@@ -5,6 +5,7 @@
 
 use crate::entity::SecondaryMap;
 use crate::ir::entities::AnyEntity;
+use crate::ir::immediates::{HeapImmData, Uimm32};
 use crate::ir::{Block, DataFlowGraph, Function, Inst, SigRef, Type, Value, ValueDef};
 use crate::packed_option::ReservedValue;
 use alloc::string::{String, ToString};
@@ -476,6 +477,47 @@ pub fn write_operands(w: &mut dyn Write, dfg: &DataFlowGraph, inst: Inst) -> fmt
             dynamic_stack_slot,
             ..
         } => write!(w, " {}, {}", arg, dynamic_stack_slot),
+        HeapLoad {
+            opcode: _,
+            heap_imm,
+            arg,
+        } => {
+            let HeapImmData {
+                flags,
+                heap,
+                offset,
+            } = dfg.heap_imms[heap_imm];
+            write!(
+                w,
+                " {heap} {flags} {arg}{optional_offset}",
+                optional_offset = if offset == Uimm32::from(0) {
+                    "".to_string()
+                } else {
+                    format!("+{offset}")
+                }
+            )
+        }
+        HeapStore {
+            opcode: _,
+            heap_imm,
+            args,
+        } => {
+            let HeapImmData {
+                flags,
+                heap,
+                offset,
+            } = dfg.heap_imms[heap_imm];
+            let [index, value] = args;
+            write!(
+                w,
+                " {heap} {flags} {index}{optional_offset}, {value}",
+                optional_offset = if offset == Uimm32::from(0) {
+                    "".to_string()
+                } else {
+                    format!("+{offset}")
+                }
+            )
+        }
         HeapAddr {
             heap,
             arg,