Pass operand types in to Wasm-to-CLIF translation

This lets us generate better CLIF based on what we know of the Wasm types. There
is a bunch we can do here, but this commit just lays down the initial
infrastructure and plumbs the type info through to the `ref.is_null`
implementation. We will now avoid actually performing null check when the input
to a `ref.is_null` instruction is a non-nullable ref type.

Author: fitzgen

crates/cranelift/src/func_environ.rs

@@ -2250,10 +2250,11 @@ impl FuncEnvironment<'_> {
     pub fn translate_ref_test(
         &mut self,
         builder: &mut FunctionBuilder<'_>,
-        ref_ty: WasmRefType,
+        test_ty: WasmRefType,
         gc_ref: ir::Value,
+        gc_ref_ty: WasmRefType,
     ) -> WasmResult<ir::Value> {
-        gc::translate_ref_test(self, builder, ref_ty, gc_ref)
+        gc::translate_ref_test(self, builder, test_ty, gc_ref, gc_ref_ty)
     }
 
     pub fn translate_ref_null(
@@ -2273,7 +2274,14 @@ impl FuncEnvironment<'_> {
         &mut self,
         mut pos: cranelift_codegen::cursor::FuncCursor,
         value: ir::Value,
+        ty: WasmRefType,
     ) -> WasmResult<ir::Value> {
+        // If we know the type is not nullable, then we don't actually need to
+        // check for null.
+        if !ty.nullable {
+            return Ok(pos.ins().iconst(ir::types::I32, 0));
+        }
+
         let byte_is_null =
             pos.ins()
                 .icmp_imm(cranelift_codegen::ir::condcodes::IntCC::Equal, value, 0);
@@ -3121,6 +3129,7 @@ impl FuncEnvironment<'_> {
     pub fn before_translate_operator(
         &mut self,
         op: &Operator,
+        _operand_types: Option<&[WasmValType]>,
         builder: &mut FunctionBuilder,
         state: &FuncTranslationState,
     ) -> WasmResult<()> {
@@ -3133,6 +3142,7 @@ impl FuncEnvironment<'_> {
     pub fn after_translate_operator(
         &mut self,
         op: &Operator,
+        _operand_types: Option<&[WasmValType]>,
         builder: &mut FunctionBuilder,
         state: &FuncTranslationState,
     ) -> WasmResult<()> {

crates/cranelift/src/func_environ/gc/disabled.rs

@@ -130,8 +130,9 @@ pub fn translate_array_set(
 pub fn translate_ref_test(
     _func_env: &mut FuncEnvironment<'_>,
     _builder: &mut FunctionBuilder<'_>,
-    _ref_ty: WasmRefType,
-    _gc_ref: ir::Value,
+    _test_ty: WasmRefType,
+    _val: ir::Value,
+    _val_ty: WasmRefType,
 ) -> WasmResult<ir::Value> {
     disabled()
 }

crates/cranelift/src/func_environ/gc/enabled.rs

@@ -887,17 +887,18 @@ pub fn translate_array_set(
 pub fn translate_ref_test(
     func_env: &mut FuncEnvironment<'_>,
     builder: &mut FunctionBuilder<'_>,
-    ref_ty: WasmRefType,
+    test_ty: WasmRefType,
     val: ir::Value,
+    val_ty: WasmRefType,
 ) -> WasmResult<ir::Value> {
-    log::trace!("translate_ref_test({ref_ty:?}, {val:?})");
+    log::trace!("translate_ref_test({test_ty:?}, {val:?})");
 
     // First special case: testing for references to bottom types.
-    if ref_ty.heap_type.is_bottom() {
-        let result = if ref_ty.nullable {
+    if test_ty.heap_type.is_bottom() {
+        let result = if test_ty.nullable {
             // All null references (within the same type hierarchy) match null
             // references to the bottom type.
-            func_env.translate_ref_is_null(builder.cursor(), val)?
+            func_env.translate_ref_is_null(builder.cursor(), val, val_ty)?
         } else {
             // `ref.test` is always false for non-nullable bottom types, as the
             // bottom types are uninhabited.
@@ -911,11 +912,11 @@ pub fn translate_ref_test(
     // the same type hierarchy as `heap_ty`, if `heap_ty` is its hierarchy's top
     // type, we only need to worry about whether we are testing for nullability
     // or not.
-    if ref_ty.heap_type.is_top() {
-        let result = if ref_ty.nullable {
+    if test_ty.heap_type.is_top() {
+        let result = if test_ty.nullable {
             builder.ins().iconst(ir::types::I32, 1)
         } else {
-            let is_null = func_env.translate_ref_is_null(builder.cursor(), val)?;
+            let is_null = func_env.translate_ref_is_null(builder.cursor(), val, val_ty)?;
             let zero = builder.ins().iconst(ir::types::I32, 0);
             let one = builder.ins().iconst(ir::types::I32, 1);
             builder.ins().select(is_null, zero, one)
@@ -926,14 +927,14 @@ pub fn translate_ref_test(
 
     // `i31ref`s are a little interesting because they don't point to GC
     // objects; we test the bit pattern of the reference itself.
-    if ref_ty.heap_type == WasmHeapType::I31 {
+    if test_ty.heap_type == WasmHeapType::I31 {
         let i31_mask = builder.ins().iconst(
             ir::types::I32,
             i64::from(wasmtime_environ::I31_DISCRIMINANT),
         );
         let is_i31 = builder.ins().band(val, i31_mask);
-        let result = if ref_ty.nullable {
-            let is_null = func_env.translate_ref_is_null(builder.cursor(), val)?;
+        let result = if test_ty.nullable {
+            let is_null = func_env.translate_ref_is_null(builder.cursor(), val, val_ty)?;
             builder.ins().bor(is_null, is_i31)
         } else {
             is_i31
@@ -945,15 +946,17 @@ pub fn translate_ref_test(
     // Otherwise, in the general case, we need to inspect our given object's
     // actual type, which also requires null-checking and i31-checking it.
 
-    let is_any_hierarchy = ref_ty.heap_type.top() == WasmHeapTopType::Any;
+    let is_any_hierarchy = test_ty.heap_type.top() == WasmHeapTopType::Any;
 
     let non_null_block = builder.create_block();
     let non_null_non_i31_block = builder.create_block();
     let continue_block = builder.create_block();
 
     // Current block: check if the reference is null and branch appropriately.
-    let is_null = func_env.translate_ref_is_null(builder.cursor(), val)?;
-    let result_when_is_null = builder.ins().iconst(ir::types::I32, ref_ty.nullable as i64);
+    let is_null = func_env.translate_ref_is_null(builder.cursor(), val, val_ty)?;
+    let result_when_is_null = builder
+        .ins()
+        .iconst(ir::types::I32, test_ty.nullable as i64);
     builder.ins().brif(
         is_null,
         continue_block,
@@ -977,7 +980,7 @@ pub fn translate_ref_test(
         let result_when_is_i31 = builder.ins().iconst(
             ir::types::I32,
             matches!(
-                ref_ty.heap_type,
+                test_ty.heap_type,
                 WasmHeapType::Any | WasmHeapType::Eq | WasmHeapType::I31
             ) as i64,
         );
@@ -1030,7 +1033,7 @@ pub fn translate_ref_test(
             .icmp(ir::condcodes::IntCC::Equal, and, expected_kind);
         builder.ins().uextend(ir::types::I32, kind_matches)
     };
-    let result = match ref_ty.heap_type {
+    let result = match test_ty.heap_type {
         WasmHeapType::Any
         | WasmHeapType::None
         | WasmHeapType::Extern

crates/cranelift/src/translate/code_translator.rs

@@ -94,7 +94,7 @@ use std::vec::Vec;
 use wasmparser::{FuncValidator, MemArg, Operator, WasmModuleResources};
 use wasmtime_environ::{
     DataIndex, ElemIndex, FuncIndex, GlobalIndex, MemoryIndex, Signed, TableIndex, TypeConvert,
-    TypeIndex, Unsigned, WasmRefType, WasmResult, wasm_unsupported,
+    TypeIndex, Unsigned, WasmRefType, WasmResult, WasmValType, wasm_unsupported,
 };
 
 /// Given a `Reachability<T>`, unwrap the inner `T` or, when unreachable, set
@@ -119,10 +119,13 @@ macro_rules! unwrap_or_return_unreachable_state {
 pub fn translate_operator(
     validator: &mut FuncValidator<impl WasmModuleResources>,
     op: &Operator,
+    operand_types: Option<&[WasmValType]>,
     builder: &mut FunctionBuilder,
     state: &mut FuncTranslationState,
     environ: &mut FuncEnvironment<'_>,
 ) -> WasmResult<()> {
+    log::trace!("Translating Wasm opcode: {op:?}");
+
     if !state.reachable {
         translate_unreachable_operator(validator, &op, builder, state, environ)?;
         return Ok(());
@@ -131,8 +134,11 @@ pub fn translate_operator(
     // Given that we believe the current block is reachable, the FunctionBuilder ought to agree.
     debug_assert!(!builder.is_unreachable());
 
+    let operand_types = operand_types.unwrap_or_else(|| {
+        panic!("should always have operand types available for valid, reachable ops; op = {op:?}")
+    });
+
     // This big match treats all Wasm code operators.
-    log::trace!("Translating Wasm opcode: {op:?}");
     match op {
         /********************************** Locals ****************************************
          *  `get_local` and `set_local` are treated as non-SSA variables and will completely
@@ -1258,7 +1264,10 @@ pub fn translate_operator(
         }
         Operator::RefIsNull => {
             let value = state.pop1();
-            state.push1(environ.translate_ref_is_null(builder.cursor(), value)?);
+            let [WasmValType::Ref(ty)] = operand_types else {
+                unreachable!("validation")
+            };
+            state.push1(environ.translate_ref_is_null(builder.cursor(), value, *ty)?);
         }
         Operator::RefFunc { function_index } => {
             let index = FuncIndex::from_u32(*function_index);
@@ -2448,8 +2457,11 @@ pub fn translate_operator(
 
         Operator::BrOnNull { relative_depth } => {
             let r = state.pop1();
+            let [.., WasmValType::Ref(r_ty)] = operand_types else {
+                unreachable!("validation")
+            };
             let (br_destination, inputs) = translate_br_if_args(*relative_depth, state);
-            let is_null = environ.translate_ref_is_null(builder.cursor(), r)?;
+            let is_null = environ.translate_ref_is_null(builder.cursor(), r, *r_ty)?;
             let else_block = builder.create_block();
             canonicalise_brif(builder, is_null, br_destination, inputs, else_block, &[]);
 
@@ -2464,7 +2476,11 @@ pub fn translate_operator(
             // Peek the value val from the stack.
             // If val is ref.null ht, then: pop the value val from the stack.
             // Else: Execute the instruction (br relative_depth).
-            let is_null = environ.translate_ref_is_null(builder.cursor(), state.peek1())?;
+            let r = state.peek1();
+            let [.., WasmValType::Ref(r_ty)] = operand_types else {
+                unreachable!("validation")
+            };
+            let is_null = environ.translate_ref_is_null(builder.cursor(), r, *r_ty)?;
             let (br_destination, inputs) = translate_br_if_args(*relative_depth, state);
             let else_block = builder.create_block();
             canonicalise_brif(builder, is_null, else_block, &[], br_destination, inputs);
@@ -2503,7 +2519,10 @@ pub fn translate_operator(
         }
         Operator::RefAsNonNull => {
             let r = state.pop1();
-            let is_null = environ.translate_ref_is_null(builder.cursor(), r)?;
+            let [.., WasmValType::Ref(r_ty)] = operand_types else {
+                unreachable!("validation")
+            };
+            let is_null = environ.translate_ref_is_null(builder.cursor(), r, *r_ty)?;
             environ.trapnz(builder, is_null, crate::TRAP_NULL_REFERENCE);
             state.push1(r);
         }
@@ -2770,6 +2789,9 @@ pub fn translate_operator(
         }
         Operator::RefTestNonNull { hty } => {
             let r = state.pop1();
+            let [.., WasmValType::Ref(r_ty)] = operand_types else {
+                unreachable!("validation")
+            };
             let heap_type = environ.convert_heap_type(*hty)?;
             let result = environ.translate_ref_test(
                 builder,
@@ -2778,11 +2800,15 @@ pub fn translate_operator(
                     nullable: false,
                 },
                 r,
+                *r_ty,
             )?;
             state.push1(result);
         }
         Operator::RefTestNullable { hty } => {
             let r = state.pop1();
+            let [.., WasmValType::Ref(r_ty)] = operand_types else {
+                unreachable!("validation")
+            };
             let heap_type = environ.convert_heap_type(*hty)?;
             let result = environ.translate_ref_test(
                 builder,
@@ -2791,11 +2817,15 @@ pub fn translate_operator(
                     nullable: true,
                 },
                 r,
+                *r_ty,
             )?;
             state.push1(result);
         }
         Operator::RefCastNonNull { hty } => {
             let r = state.pop1();
+            let [.., WasmValType::Ref(r_ty)] = operand_types else {
+                unreachable!("validation")
+            };
             let heap_type = environ.convert_heap_type(*hty)?;
             let cast_okay = environ.translate_ref_test(
                 builder,
@@ -2804,12 +2834,16 @@ pub fn translate_operator(
                     nullable: false,
                 },
                 r,
+                *r_ty,
             )?;
             environ.trapz(builder, cast_okay, crate::TRAP_CAST_FAILURE);
             state.push1(r);
         }
         Operator::RefCastNullable { hty } => {
             let r = state.pop1();
+            let [.., WasmValType::Ref(r_ty)] = operand_types else {
+                unreachable!("validation")
+            };
             let heap_type = environ.convert_heap_type(*hty)?;
             let cast_okay = environ.translate_ref_test(
                 builder,
@@ -2818,6 +2852,7 @@ pub fn translate_operator(
                     nullable: true,
                 },
                 r,
+                *r_ty,
             )?;
             environ.trapz(builder, cast_okay, crate::TRAP_CAST_FAILURE);
             state.push1(r);
@@ -2830,9 +2865,12 @@ pub fn translate_operator(
             from_ref_type: _,
         } => {
             let r = state.peek1();
+            let [.., WasmValType::Ref(r_ty)] = operand_types else {
+                unreachable!("validation")
+            };
 
             let to_ref_type = environ.convert_ref_type(*to_ref_type)?;
-            let cast_is_okay = environ.translate_ref_test(builder, to_ref_type, r)?;
+            let cast_is_okay = environ.translate_ref_test(builder, to_ref_type, r, *r_ty)?;
 
             let (cast_succeeds_block, inputs) = translate_br_if_args(*relative_depth, state);
             let cast_fails_block = builder.create_block();
@@ -2863,9 +2901,12 @@ pub fn translate_operator(
             from_ref_type: _,
         } => {
             let r = state.peek1();
+            let [.., WasmValType::Ref(r_ty)] = operand_types else {
+                unreachable!("validation")
+            };
 
             let to_ref_type = environ.convert_ref_type(*to_ref_type)?;
-            let cast_is_okay = environ.translate_ref_test(builder, to_ref_type, r)?;
+            let cast_is_okay = environ.translate_ref_test(builder, to_ref_type, r, *r_ty)?;
 
             let (cast_fails_block, inputs) = translate_br_if_args(*relative_depth, state);
             let cast_succeeds_block = builder.create_block();

crates/cranelift/src/translate/func_translator.rs

@@ -241,15 +241,34 @@ fn parse_function_body(
     debug_assert_eq!(state.control_stack.len(), 1, "State not initialized");
 
     environ.before_translate_function(builder, state)?;
+
     let mut reader = OperatorsReader::new(reader);
+    let mut operand_types = vec![];
+
     while !reader.eof() {
         let pos = reader.original_position();
         builder.set_srcloc(cur_srcloc(&reader.get_binary_reader()));
+
         let op = reader.read()?;
+
+        // Get the operand types for this operator.
+        let arity = op.operator_arity(&*validator);
+        operand_types.clear();
+        let operand_types = arity.and_then(|(operand_arity, _result_arity)| {
+            for i in (0..operand_arity).rev() {
+                let i = usize::try_from(i).unwrap();
+                let ty = validator.get_operand_type(i)??;
+                let ty = environ.convert_valtype(ty).ok()?;
+                operand_types.push(ty);
+            }
+            Some(&operand_types[..])
+        });
+
         validator.op(pos, &op)?;
-        environ.before_translate_operator(&op, builder, state)?;
-        translate_operator(validator, &op, builder, state, environ)?;
-        environ.after_translate_operator(&op, builder, state)?;
+
+        environ.before_translate_operator(&op, operand_types, builder, state)?;
+        translate_operator(validator, &op, operand_types, builder, state, environ)?;
+        environ.after_translate_operator(&op, operand_types, builder, state)?;
     }
     environ.after_translate_function(builder, state)?;
     reader.finish()?;