Deduplicate const-eval fast paths in Wasmtime (#12406)

Closes #12243

Author: alexcrichton

crates/cranelift/src/func_environ.rs

@@ -3,8 +3,8 @@ pub(crate) mod stack_switching;
 
 use crate::compiler::Compiler;
 use crate::translate::{
-    FuncTranslationStacks, GlobalConstValue, GlobalVariable, Heap, HeapData, StructFieldsVec,
-    TableData, TableSize, TargetEnvironment,
+    FuncTranslationStacks, GlobalVariable, Heap, HeapData, StructFieldsVec, TableData, TableSize,
+    TargetEnvironment,
 };
 use crate::{BuiltinFunctionSignatures, TRAP_INTERNAL_ASSERT};
 use cranelift_codegen::cursor::FuncCursor;
@@ -25,11 +25,11 @@ use wasmparser::{FuncValidator, Operator, WasmFeatures, WasmModuleResources};
 use wasmtime_core::math::f64_cvt_to_int_bounds;
 use wasmtime_environ::{
     BuiltinFunctionIndex, DataIndex, DefinedFuncIndex, ElemIndex, EngineOrModuleTypeIndex,
-    FrameStateSlotBuilder, FrameValType, FuncIndex, FuncKey, GlobalIndex, IndexType, Memory,
-    MemoryIndex, Module, ModuleInternedTypeIndex, ModuleTranslation, ModuleTypesBuilder, PtrSize,
-    Table, TableIndex, TagIndex, TripleExt, Tunables, TypeConvert, TypeIndex, VMOffsets,
-    WasmCompositeInnerType, WasmFuncType, WasmHeapTopType, WasmHeapType, WasmRefType, WasmResult,
-    WasmValType,
+    FrameStateSlotBuilder, FrameValType, FuncIndex, FuncKey, GlobalConstValue, GlobalIndex,
+    IndexType, Memory, MemoryIndex, Module, ModuleInternedTypeIndex, ModuleTranslation,
+    ModuleTypesBuilder, PtrSize, Table, TableIndex, TagIndex, TripleExt, Tunables, TypeConvert,
+    TypeIndex, VMOffsets, WasmCompositeInnerType, WasmFuncType, WasmHeapTopType, WasmHeapType,
+    WasmRefType, WasmResult, WasmValType,
 };
 use wasmtime_environ::{FUNCREF_INIT_BIT, FUNCREF_MASK};
 
@@ -1429,7 +1429,7 @@ impl FuncEnvironment<'_> {
         if !self.module.globals[index].mutability {
             if let Some(index) = self.module.defined_global_index(index) {
                 let init = &self.module.global_initializers[index];
-                if let Some(value) = GlobalConstValue::const_eval(init) {
+                if let Some(value) = init.const_eval() {
                     return GlobalVariable::Constant { value };
                 }
             }

crates/cranelift/src/translate/environ/mod.rs

@@ -3,6 +3,4 @@
 #[macro_use]
 mod spec;
 
-pub use crate::translate::environ::spec::{
-    GlobalConstValue, GlobalVariable, StructFieldsVec, TargetEnvironment,
-};
+pub use crate::translate::environ::spec::{GlobalVariable, StructFieldsVec, TargetEnvironment};

crates/cranelift/src/translate/environ/spec.rs

@@ -10,7 +10,7 @@ use cranelift_codegen::ir;
 use cranelift_codegen::ir::immediates::Offset32;
 use cranelift_codegen::isa::TargetFrontendConfig;
 use smallvec::SmallVec;
-use wasmtime_environ::{ConstExpr, ConstOp, Tunables, TypeConvert, WasmHeapType};
+use wasmtime_environ::{GlobalConstValue, Tunables, TypeConvert, WasmHeapType};
 
 /// The value of a WebAssembly global variable.
 #[derive(Clone, Copy)]
@@ -35,32 +35,6 @@ pub enum GlobalVariable {
     Custom,
 }
 
-/// A global's constant value, known at compile time.
-#[derive(Clone, Copy)]
-pub enum GlobalConstValue {
-    I32(i32),
-    I64(i64),
-    F32(u32),
-    F64(u64),
-    V128(u128),
-}
-
-impl GlobalConstValue {
-    /// Attempt to evaluate the given const-expr at compile time.
-    pub fn const_eval(init: &ConstExpr) -> Option<GlobalConstValue> {
-        // TODO: Actually maintain an evaluation stack and handle `i32.add`,
-        // `i32.sub`, etc... const ops.
-        match init.ops() {
-            [ConstOp::I32Const(x)] => Some(Self::I32(*x)),
-            [ConstOp::I64Const(x)] => Some(Self::I64(*x)),
-            [ConstOp::F32Const(x)] => Some(Self::F32(*x)),
-            [ConstOp::F64Const(x)] => Some(Self::F64(*x)),
-            [ConstOp::V128Const(x)] => Some(Self::V128(*x)),
-            _ => None,
-        }
-    }
-}
-
 /// Environment affecting the translation of a WebAssembly.
 pub trait TargetEnvironment: TypeConvert {
     /// Get the information needed to produce Cranelift IR for the given target.

crates/cranelift/src/translate/mod.rs

@@ -18,7 +18,7 @@ mod stack;
 mod table;
 mod translation_utils;
 
-pub use self::environ::{GlobalConstValue, GlobalVariable, StructFieldsVec, TargetEnvironment};
+pub use self::environ::{GlobalVariable, StructFieldsVec, TargetEnvironment};
 pub use self::func_translator::FuncTranslator;
 pub use self::heap::{Heap, HeapData};
 pub use self::stack::FuncTranslationStacks;

crates/environ/src/types.rs

@@ -1849,23 +1849,39 @@ impl ConstExpr {
     /// We use this for certain table optimizations that rely on
     /// knowing for sure that index 0 is not referenced.
     pub fn provably_nonzero_i32(&self) -> bool {
-        assert!(self.ops.len() > 0);
-        if self.ops.len() > 1 {
-            // Compound expressions not yet supported: conservatively
-            // return `false` (we can't prove nonzero).
-            return false;
-        }
-        // Exactly one op at this point.
-        match self.ops[0] {
-            // An actual zero value -- definitely not nonzero!
-            ConstOp::I32Const(0) => false,
-            // Any other constant value -- provably nonzero, if above
-            // did not match.
-            ConstOp::I32Const(_) => true,
-            // Anything else: we can't prove anything.
+        match self.const_eval() {
+            Some(GlobalConstValue::I32(x)) => x != 0,
+
+            // Conservatively return `false` for non-const-eval-able expressions
+            // as well as everything else.
             _ => false,
         }
     }
+
+    /// Attempt to evaluate the given const-expr at compile time.
+    pub fn const_eval(&self) -> Option<GlobalConstValue> {
+        // TODO: Actually maintain an evaluation stack and handle `i32.add`,
+        // `i32.sub`, etc... const ops.
+        match self.ops() {
+            [ConstOp::I32Const(x)] => Some(GlobalConstValue::I32(*x)),
+            [ConstOp::I64Const(x)] => Some(GlobalConstValue::I64(*x)),
+            [ConstOp::F32Const(x)] => Some(GlobalConstValue::F32(*x)),
+            [ConstOp::F64Const(x)] => Some(GlobalConstValue::F64(*x)),
+            [ConstOp::V128Const(x)] => Some(GlobalConstValue::V128(*x)),
+            _ => None,
+        }
+    }
+}
+
+/// A global's constant value, known at compile time.
+#[expect(missing_docs, reason = "self-describing variants")]
+#[derive(Clone, Copy)]
+pub enum GlobalConstValue {
+    I32(i32),
+    I64(i64),
+    F32(u32),
+    F64(u64),
+    V128(u128),
 }
 
 /// The subset of Wasm opcodes that are constant.

crates/wasmtime/src/runtime/vm/const_expr.rs

@@ -8,7 +8,7 @@ use crate::{
     AnyRef, ArrayRef, ArrayRefPre, ArrayType, ExternRef, I31, StructRef, StructRefPre, StructType,
 };
 use crate::{OpaqueRootScope, Val};
-use wasmtime_environ::{ConstExpr, ConstOp, FuncIndex, GlobalIndex};
+use wasmtime_environ::{ConstExpr, ConstOp, FuncIndex, GlobalConstValue, GlobalIndex};
 #[cfg(feature = "gc")]
 use wasmtime_environ::{VMSharedTypeIndex, WasmCompositeInnerType, WasmCompositeType, WasmSubType};
 
@@ -168,12 +168,12 @@ impl ConstExprEvaluator {
     /// for `i32.const N`.
     #[inline]
     pub fn try_simple(&mut self, expr: &ConstExpr) -> Option<&Val> {
-        match expr.ops() {
-            [ConstOp::I32Const(i)] => Some(self.return_one(Val::I32(*i))),
-            [ConstOp::I64Const(i)] => Some(self.return_one(Val::I64(*i))),
-            [ConstOp::F32Const(f)] => Some(self.return_one(Val::F32(*f))),
-            [ConstOp::F64Const(f)] => Some(self.return_one(Val::F64(*f))),
-            _ => None,
+        match expr.const_eval()? {
+            GlobalConstValue::I32(i) => Some(self.return_one(Val::I32(i))),
+            GlobalConstValue::I64(i) => Some(self.return_one(Val::I64(i))),
+            GlobalConstValue::F32(f) => Some(self.return_one(Val::F32(f))),
+            GlobalConstValue::F64(f) => Some(self.return_one(Val::F64(f))),
+            GlobalConstValue::V128(i) => Some(self.return_one(Val::V128(i.into()))),
         }
     }