Wasmtime(gc): Support (ref.i31 (global.get $g)) const expressions

We had something hacked together to support `(ref.i31 (i32.const N))`. It wasn't
a long-term solution. This is the first time that we have to really deal with
multi-instruction const expressions.

This commit introduces a tiny interpreter to evaluate const expressions.

cranelift/wasm/src/environ/spec.rs

@@ -8,8 +8,8 @@
 
 use crate::state::FuncTranslationState;
 use crate::{
-    DataIndex, ElemIndex, FuncIndex, Global, GlobalIndex, GlobalInit, Heap, HeapData, Memory,
-    MemoryIndex, Table, TableIndex, Tag, TagIndex, TypeConvert, TypeIndex, WasmError, WasmFuncType,
+    DataIndex, ElemIndex, FuncIndex, Global, GlobalIndex, Heap, HeapData, Memory, MemoryIndex,
+    Table, TableIndex, Tag, TagIndex, TypeConvert, TypeIndex, WasmError, WasmFuncType,
     WasmHeapType, WasmResult,
 };
 use cranelift_codegen::cursor::FuncCursor;
@@ -21,7 +21,7 @@ use cranelift_frontend::FunctionBuilder;
 use std::boxed::Box;
 use std::string::ToString;
 use wasmparser::{FuncValidator, FunctionBody, Operator, ValidatorResources, WasmFeatures};
-use wasmtime_types::ModuleInternedTypeIndex;
+use wasmtime_types::{ConstExpr, ModuleInternedTypeIndex};
 
 /// The value of a WebAssembly global variable.
 #[derive(Clone, Copy)]
@@ -802,7 +802,7 @@ pub trait ModuleEnvironment<'data>: TypeConvert {
     }
 
     /// Declares a global to the environment.
-    fn declare_global(&mut self, global: Global, init: GlobalInit) -> WasmResult<()>;
+    fn declare_global(&mut self, global: Global, init: ConstExpr) -> WasmResult<()>;
 
     /// Provides the number of exports up front. By default this does nothing, but
     /// implementations can use this to preallocate memory if desired.

cranelift/wasm/src/sections_translator.rs

@@ -10,8 +10,8 @@
 use crate::environ::ModuleEnvironment;
 use crate::wasm_unsupported;
 use crate::{
-    DataIndex, ElemIndex, FuncIndex, GlobalIndex, GlobalInit, Memory, MemoryIndex, TableIndex, Tag,
-    TagIndex, TypeIndex, WasmError, WasmResult,
+    DataIndex, ElemIndex, FuncIndex, GlobalIndex, Memory, MemoryIndex, TableIndex, Tag, TagIndex,
+    TypeIndex, WasmError, WasmResult,
 };
 use cranelift_entity::packed_option::ReservedValue;
 use cranelift_entity::EntityRef;
@@ -23,6 +23,7 @@ use wasmparser::{
     ImportSectionReader, MemorySectionReader, MemoryType, NameSectionReader, Naming, Operator,
     TableSectionReader, TagSectionReader, TagType, TypeRef, TypeSectionReader,
 };
+use wasmtime_types::ConstExpr;
 
 fn memory(ty: MemoryType) -> Memory {
     Memory {
@@ -169,29 +170,7 @@ pub fn parse_global_section(
 
     for entry in globals {
         let wasmparser::Global { ty, init_expr } = entry?;
-        let mut init_expr_reader = init_expr.get_binary_reader();
-        let initializer = match init_expr_reader.read_operator()? {
-            Operator::I32Const { value } => GlobalInit::I32Const(value),
-            Operator::I64Const { value } => GlobalInit::I64Const(value),
-            Operator::F32Const { value } => GlobalInit::F32Const(value.bits()),
-            Operator::F64Const { value } => GlobalInit::F64Const(value.bits()),
-            Operator::V128Const { value } => {
-                GlobalInit::V128Const(u128::from_le_bytes(*value.bytes()))
-            }
-            Operator::RefNull { hty: _ } => GlobalInit::RefNullConst,
-            Operator::RefFunc { function_index } => {
-                GlobalInit::RefFunc(FuncIndex::from_u32(function_index))
-            }
-            Operator::GlobalGet { global_index } => {
-                GlobalInit::GetGlobal(GlobalIndex::from_u32(global_index))
-            }
-            ref s => {
-                return Err(wasm_unsupported!(
-                    "unsupported init expr in global section: {:?}",
-                    s
-                ));
-            }
-        };
+        let (initializer, _escaped) = ConstExpr::from_wasmparser(init_expr)?;
         let ty = environ.convert_global_type(&ty);
         environ.declare_global(ty, initializer)?;
     }

crates/cranelift/src/func_environ.rs

@@ -759,7 +759,7 @@ impl<'module_environment> FuncEnvironment<'module_environment> {
         };
 
         let table = &self.module.table_plans[index].table;
-        let element_size = if table.wasm_ty.is_gc_heap_type() {
+        let element_size = if table.wasm_ty.is_vmgcref_type() {
             // For GC-managed references, tables store `Option<VMGcRef>`s.
             ir::types::I32.bytes()
         } else {

crates/cranelift/src/gc.rs

@@ -74,7 +74,7 @@ pub fn gc_ref_table_grow_builtin(
     func_env: &mut FuncEnvironment<'_>,
     func: &mut ir::Function,
 ) -> WasmResult<ir::FuncRef> {
-    debug_assert!(ty.is_gc_heap_type());
+    debug_assert!(ty.is_vmgcref_type_and_not_i31());
     imp::gc_ref_table_grow_builtin(ty, func_env, func)
 }
 
@@ -85,7 +85,7 @@ pub fn gc_ref_table_fill_builtin(
     func_env: &mut FuncEnvironment<'_>,
     func: &mut ir::Function,
 ) -> WasmResult<ir::FuncRef> {
-    debug_assert!(ty.is_gc_heap_type());
+    debug_assert!(ty.is_vmgcref_type_and_not_i31());
     imp::gc_ref_table_fill_builtin(ty, func_env, func)
 }
 

crates/cranelift/src/gc/enabled.rs

@@ -55,7 +55,7 @@ pub fn gc_ref_table_grow_builtin(
     func_env: &mut FuncEnvironment<'_>,
     func: &mut ir::Function,
 ) -> WasmResult<ir::FuncRef> {
-    debug_assert!(ty.is_gc_heap_type());
+    debug_assert!(ty.is_vmgcref_type_and_not_i31());
     Ok(func_env.builtin_functions.table_grow_gc_ref(func))
 }
 
@@ -64,7 +64,7 @@ pub fn gc_ref_table_fill_builtin(
     func_env: &mut FuncEnvironment<'_>,
     func: &mut ir::Function,
 ) -> WasmResult<ir::FuncRef> {
-    debug_assert!(ty.is_gc_heap_type());
+    debug_assert!(ty.is_vmgcref_type_and_not_i31());
     Ok(func_env.builtin_functions.table_fill_gc_ref(func))
 }
 
@@ -169,7 +169,7 @@ impl FuncEnvironment<'_> {
         ty: WasmRefType,
         gc_ref: ir::Value,
     ) -> ir::Value {
-        assert!(ty.is_gc_heap_type());
+        assert!(ty.is_vmgcref_type_and_not_i31());
 
         let might_be_i31 = match ty.heap_type {
             WasmHeapType::Any => true,
@@ -317,7 +317,7 @@ impl GcCompiler for DrcCompiler {
         src: ir::Value,
         flags: ir::MemFlags,
     ) -> WasmResult<ir::Value> {
-        assert!(ty.is_gc_heap_type());
+        assert!(ty.is_vmgcref_type_and_not_i31());
 
         let reference_type = func_env.reference_type(ty.heap_type);
 
@@ -460,7 +460,7 @@ impl GcCompiler for DrcCompiler {
         new_val: ir::Value,
         flags: ir::MemFlags,
     ) -> WasmResult<()> {
-        assert!(ty.is_gc_heap_type());
+        assert!(ty.is_vmgcref_type_and_not_i31());
 
         let ref_ty = func_env.reference_type(ty.heap_type);
 

crates/environ/src/module.rs

@@ -418,17 +418,8 @@ pub enum TableInitialValue {
         /// case the elements are initialized to null.
         precomputed: Vec<FuncIndex>,
     },
-
-    /// Initialize each table element to the function reference given
-    /// by the `FuncIndex`.
-    FuncRef(FuncIndex),
-
-    /// At instantiation time this global is loaded and the funcref value is
-    /// used to initialize the table.
-    GlobalGet(GlobalIndex),
-
-    /// Initialize the table element to an `i31ref` of the given value.
-    I31Ref(i32),
+    /// An arbitrary const expression.
+    Expr(ConstExpr),
 }
 
 /// A WebAssembly table initializer segment.
@@ -452,7 +443,7 @@ pub enum TableSegmentElements {
     /// indicates a null function.
     Functions(Box<[FuncIndex]>),
     /// Arbitrary expressions, aka either functions, null or a load of a global.
-    Expressions(Box<[TableElementExpression]>),
+    Expressions(Box<[ConstExpr]>),
 }
 
 impl TableSegmentElements {
@@ -465,17 +456,6 @@ impl TableSegmentElements {
     }
 }
 
-/// Different kinds of expression that can initialize table elements.
-#[derive(Clone, Debug, Serialize, Deserialize)]
-pub enum TableElementExpression {
-    /// `ref.func $f`
-    Function(FuncIndex),
-    /// `global.get $g`
-    GlobalGet(GlobalIndex),
-    /// `ref.null $ty`
-    Null,
-}
-
 /// A translated WebAssembly module, excluding the function bodies and
 /// memory initializers.
 #[derive(Default, Debug, Serialize, Deserialize)]
@@ -542,7 +522,7 @@ pub struct Module {
     pub globals: PrimaryMap<GlobalIndex, Global>,
 
     /// WebAssembly global initializers for locally-defined globals.
-    pub global_initializers: PrimaryMap<DefinedGlobalIndex, GlobalInit>,
+    pub global_initializers: PrimaryMap<DefinedGlobalIndex, ConstExpr>,
 }
 
 /// Initialization routines for creating an instance, encompassing imports,
@@ -720,18 +700,6 @@ impl Module {
             func_ref: FuncRefIndex::reserved_value(),
         })
     }
-
-    /// Appends a new function to this module with the given type information.
-    pub fn push_escaped_function(
-        &mut self,
-        signature: ModuleInternedTypeIndex,
-        func_ref: FuncRefIndex,
-    ) -> FuncIndex {
-        self.functions.push(FunctionType {
-            signature,
-            func_ref,
-        })
-    }
 }
 
 /// Type information about functions in a wasm module.