[Wasm GC] Optimize away null arms that would trap (#5358)

E.g.

(struct.get
  (select
    (ref.null ..)
    (something)
    (condition)
  )
)

If traps-never-happen then this can be

(drop (condition))
(struct.get
  (something)
)

That is, we can remove the arm that is null, as it would trap but traps are
assumed to not happen.

Also fix a bug this uncovers on struct.set on a null type.

2 files changed, 361 insertions, 8 deletions

diff --git a/src/passes/OptimizeInstructions.cpp b/src/passes/OptimizeInstructions.cpp
index f96d4d30a..d218f5627 100644
--- a/src/passes/OptimizeInstructions.cpp
+++ b/src/passes/OptimizeInstructions.cpp
@@ -1513,10 +1513,75 @@ struct OptimizeInstructions
     return getDroppedChildrenAndAppend(curr, result);
   }
 
-  bool trapOnNull(Expression* curr, Expression* ref) {
+  Expression* getResultOfFirst(Expression* first, Expression* second) {
+    return wasm::getResultOfFirst(
+      first, second, getFunction(), getModule(), getPassOptions());
+  }
+
+  // Optimize an instruction and the reference it operates on, under the
+  // assumption that if the reference is a null then we will trap. Returns true
+  // if we replaced the expression with something simpler. Returns false if we
+  // found nothing to optimize, or if we just modified or replaced the ref (but
+  // not the expression itself).
+  bool trapOnNull(Expression* curr, Expression*& ref) {
+    Builder builder(*getModule());
+
+    if (getPassOptions().trapsNeverHappen) {
+      // We can ignore the possibility of the reference being an input, so
+      //
+      //    (if
+      //      (condition)
+      //      (null)
+      //      (other))
+      // =>
+      //    (drop
+      //      (condition))
+      //    (other)
+      //
+      // That is, we will by assumption not read from the null, so remove that
+      // arm.
+      //
+      // TODO We could recurse here.
+      // TODO We could do similar things for casts (rule out an impossible arm).
+      // TODO Worth thinking about an 'assume' instrinsic of some form that
+      //      annotates knowledge about a value, or another mechanism to allow
+      //      that information to be passed around.
+      if (auto* iff = ref->dynCast<If>()) {
+        if (iff->ifFalse) {
+          if (iff->ifTrue->type.isNull()) {
+            ref = builder.makeSequence(builder.makeDrop(iff->condition),
+                                       iff->ifFalse);
+            return false;
+          }
+          if (iff->ifFalse->type.isNull()) {
+            ref = builder.makeSequence(builder.makeDrop(iff->condition),
+                                       iff->ifTrue);
+            return false;
+          }
+        }
+      }
+
+      if (auto* select = ref->dynCast<Select>()) {
+        if (select->ifTrue->type.isNull()) {
+          ref = builder.makeSequence(
+            builder.makeDrop(select->ifTrue),
+            getResultOfFirst(select->ifFalse,
+                             builder.makeDrop(select->condition)));
+          return false;
+        }
+        if (select->ifFalse->type.isNull()) {
+          ref = getResultOfFirst(
+            select->ifTrue,
+            builder.makeSequence(builder.makeDrop(select->ifFalse),
+                                 builder.makeDrop(select->condition)));
+          return false;
+        }
+      }
+    }
+
     if (ref->type.isNull()) {
-      replaceCurrent(getDroppedChildrenAndAppend(
-        curr, Builder(*getModule()).makeUnreachable()));
+      replaceCurrent(
+        getDroppedChildrenAndAppend(curr, builder.makeUnreachable()));
       // Propagate the unreachability.
       refinalize = true;
       return true;
@@ -1591,8 +1656,12 @@ struct OptimizeInstructions
     }
 
     if (curr->ref->type != Type::unreachable && curr->value->type.isInteger()) {
-      const auto& fields = curr->ref->type.getHeapType().getStruct().fields;
-      optimizeStoredValue(curr->value, fields[curr->index].getByteSize());
+      // We must avoid the case of a null type.
+      auto heapType = curr->ref->type.getHeapType();
+      if (heapType.isStruct()) {
+        const auto& fields = heapType.getStruct().fields;
+        optimizeStoredValue(curr->value, fields[curr->index].getByteSize());
+      }
     }
 
     // If our reference is a tee of a struct.new, we may be able to fold the
diff --git a/test/lit/passes/optimize-instructions-gc-tnh.wast b/test/lit/passes/optimize-instructions-gc-tnh.wast
index e1bbe65cd..d9520b18c 100644
--- a/test/lit/passes/optimize-instructions-gc-tnh.wast
+++ b/test/lit/passes/optimize-instructions-gc-tnh.wast
@@ -3,9 +3,9 @@
 ;; RUN: wasm-opt %s --optimize-instructions                      -all --nominal -S -o - | filecheck %s --check-prefix NO_TNH
 
 (module
-  ;; TNH:      (type $struct (struct ))
-  ;; NO_TNH:      (type $struct (struct ))
-  (type $struct (struct_subtype data))
+  ;; TNH:      (type $struct (struct (field (mut i32))))
+  ;; NO_TNH:      (type $struct (struct (field (mut i32))))
+  (type $struct (struct_subtype (field (mut i32)) data))
 
   ;; TNH:      (func $ref.eq (type $eqref_eqref_=>_i32) (param $a eqref) (param $b eqref) (result i32)
   ;; TNH-NEXT:  (ref.eq
@@ -191,4 +191,288 @@
       )
     )
   )
+
+  ;; TNH:      (func $if.arm.null (type $i32_ref|$struct|_=>_none) (param $x i32) (param $ref (ref $struct))
+  ;; TNH-NEXT:  (struct.set $struct 0
+  ;; TNH-NEXT:   (block (result (ref $struct))
+  ;; TNH-NEXT:    (drop
+  ;; TNH-NEXT:     (local.get $x)
+  ;; TNH-NEXT:    )
+  ;; TNH-NEXT:    (local.get $ref)
+  ;; TNH-NEXT:   )
+  ;; TNH-NEXT:   (i32.const 1)
+  ;; TNH-NEXT:  )
+  ;; TNH-NEXT:  (struct.set $struct 0
+  ;; TNH-NEXT:   (block (result (ref $struct))
+  ;; TNH-NEXT:    (drop
+  ;; TNH-NEXT:     (local.get $x)
+  ;; TNH-NEXT:    )
+  ;; TNH-NEXT:    (local.get $ref)
+  ;; TNH-NEXT:   )
+  ;; TNH-NEXT:   (i32.const 2)
+  ;; TNH-NEXT:  )
+  ;; TNH-NEXT: )
+  ;; NO_TNH:      (func $if.arm.null (type $i32_ref|$struct|_=>_none) (param $x i32) (param $ref (ref $struct))
+  ;; NO_TNH-NEXT:  (struct.set $struct 0
+  ;; NO_TNH-NEXT:   (if (result (ref null $struct))
+  ;; NO_TNH-NEXT:    (local.get $x)
+  ;; NO_TNH-NEXT:    (local.get $ref)
+  ;; NO_TNH-NEXT:    (ref.null none)
+  ;; NO_TNH-NEXT:   )
+  ;; NO_TNH-NEXT:   (i32.const 1)
+  ;; NO_TNH-NEXT:  )
+  ;; NO_TNH-NEXT:  (struct.set $struct 0
+  ;; NO_TNH-NEXT:   (if (result (ref null $struct))
+  ;; NO_TNH-NEXT:    (local.get $x)
+  ;; NO_TNH-NEXT:    (ref.null none)
+  ;; NO_TNH-NEXT:    (local.get $ref)
+  ;; NO_TNH-NEXT:   )
+  ;; NO_TNH-NEXT:   (i32.const 2)
+  ;; NO_TNH-NEXT:  )
+  ;; NO_TNH-NEXT: )
+  (func $if.arm.null (param $x i32) (param $ref (ref $struct))
+    ;; A set will trap on a null, so in tnh mode we know the null arm is not
+    ;; executed, and the other one is.
+    (struct.set $struct 0
+      (if (result (ref null $struct))
+        (local.get $x)
+        (local.get $ref)
+        (ref.null none)
+      )
+      (i32.const 1)
+    )
+    (struct.set $struct 0
+      (if (result (ref null $struct))
+        (local.get $x)
+        (ref.null none)
+        (local.get $ref)
+      )
+      (i32.const 2)
+    )
+  )
+
+  ;; TNH:      (func $select.arm.null (type $i32_ref|$struct|_=>_none) (param $x i32) (param $ref (ref $struct))
+  ;; TNH-NEXT:  (struct.set $struct 0
+  ;; TNH-NEXT:   (block (result (ref $struct))
+  ;; TNH-NEXT:    (block
+  ;; TNH-NEXT:     (drop
+  ;; TNH-NEXT:      (ref.null none)
+  ;; TNH-NEXT:     )
+  ;; TNH-NEXT:     (drop
+  ;; TNH-NEXT:      (local.get $x)
+  ;; TNH-NEXT:     )
+  ;; TNH-NEXT:    )
+  ;; TNH-NEXT:    (local.get $ref)
+  ;; TNH-NEXT:   )
+  ;; TNH-NEXT:   (i32.const 1)
+  ;; TNH-NEXT:  )
+  ;; TNH-NEXT:  (struct.set $struct 0
+  ;; TNH-NEXT:   (block (result (ref $struct))
+  ;; TNH-NEXT:    (drop
+  ;; TNH-NEXT:     (ref.null none)
+  ;; TNH-NEXT:    )
+  ;; TNH-NEXT:    (block (result (ref $struct))
+  ;; TNH-NEXT:     (drop
+  ;; TNH-NEXT:      (local.get $x)
+  ;; TNH-NEXT:     )
+  ;; TNH-NEXT:     (local.get $ref)
+  ;; TNH-NEXT:    )
+  ;; TNH-NEXT:   )
+  ;; TNH-NEXT:   (i32.const 2)
+  ;; TNH-NEXT:  )
+  ;; TNH-NEXT: )
+  ;; NO_TNH:      (func $select.arm.null (type $i32_ref|$struct|_=>_none) (param $x i32) (param $ref (ref $struct))
+  ;; NO_TNH-NEXT:  (struct.set $struct 0
+  ;; NO_TNH-NEXT:   (select (result (ref null $struct))
+  ;; NO_TNH-NEXT:    (local.get $ref)
+  ;; NO_TNH-NEXT:    (ref.null none)
+  ;; NO_TNH-NEXT:    (local.get $x)
+  ;; NO_TNH-NEXT:   )
+  ;; NO_TNH-NEXT:   (i32.const 1)
+  ;; NO_TNH-NEXT:  )
+  ;; NO_TNH-NEXT:  (struct.set $struct 0
+  ;; NO_TNH-NEXT:   (select (result (ref null $struct))
+  ;; NO_TNH-NEXT:    (ref.null none)
+  ;; NO_TNH-NEXT:    (local.get $ref)
+  ;; NO_TNH-NEXT:    (local.get $x)
+  ;; NO_TNH-NEXT:   )
+  ;; NO_TNH-NEXT:   (i32.const 2)
+  ;; NO_TNH-NEXT:  )
+  ;; NO_TNH-NEXT: )
+  (func $select.arm.null (param $x i32) (param $ref (ref $struct))
+    ;; As above but with a select.
+    (struct.set $struct 0
+      (select (result (ref null $struct))
+        (local.get $ref)
+        (ref.null none)
+        (local.get $x)
+      )
+      (i32.const 1)
+    )
+    (struct.set $struct 0
+      (select (result (ref null $struct))
+        (ref.null none)
+        (local.get $ref)
+        (local.get $x)
+      )
+      (i32.const 2)
+    )
+  )
+
+  ;; TNH:      (func $select.arm.null.effects (type $none_=>_none)
+  ;; TNH-NEXT:  (local $0 (ref $struct))
+  ;; TNH-NEXT:  (local $1 (ref $struct))
+  ;; TNH-NEXT:  (struct.set $struct 0
+  ;; TNH-NEXT:   (block (result (ref $struct))
+  ;; TNH-NEXT:    (local.set $0
+  ;; TNH-NEXT:     (call $get-ref)
+  ;; TNH-NEXT:    )
+  ;; TNH-NEXT:    (block
+  ;; TNH-NEXT:     (drop
+  ;; TNH-NEXT:      (call $get-null)
+  ;; TNH-NEXT:     )
+  ;; TNH-NEXT:     (drop
+  ;; TNH-NEXT:      (call $get-i32)
+  ;; TNH-NEXT:     )
+  ;; TNH-NEXT:    )
+  ;; TNH-NEXT:    (local.get $0)
+  ;; TNH-NEXT:   )
+  ;; TNH-NEXT:   (i32.const 1)
+  ;; TNH-NEXT:  )
+  ;; TNH-NEXT:  (struct.set $struct 0
+  ;; TNH-NEXT:   (block (result (ref $struct))
+  ;; TNH-NEXT:    (drop
+  ;; TNH-NEXT:     (call $get-null)
+  ;; TNH-NEXT:    )
+  ;; TNH-NEXT:    (block (result (ref $struct))
+  ;; TNH-NEXT:     (local.set $1
+  ;; TNH-NEXT:      (call $get-ref)
+  ;; TNH-NEXT:     )
+  ;; TNH-NEXT:     (drop
+  ;; TNH-NEXT:      (call $get-i32)
+  ;; TNH-NEXT:     )
+  ;; TNH-NEXT:     (local.get $1)
+  ;; TNH-NEXT:    )
+  ;; TNH-NEXT:   )
+  ;; TNH-NEXT:   (i32.const 2)
+  ;; TNH-NEXT:  )
+  ;; TNH-NEXT: )
+  ;; NO_TNH:      (func $select.arm.null.effects (type $none_=>_none)
+  ;; NO_TNH-NEXT:  (struct.set $struct 0
+  ;; NO_TNH-NEXT:   (select (result (ref null $struct))
+  ;; NO_TNH-NEXT:    (call $get-ref)
+  ;; NO_TNH-NEXT:    (call $get-null)
+  ;; NO_TNH-NEXT:    (call $get-i32)
+  ;; NO_TNH-NEXT:   )
+  ;; NO_TNH-NEXT:   (i32.const 1)
+  ;; NO_TNH-NEXT:  )
+  ;; NO_TNH-NEXT:  (struct.set $struct 0
+  ;; NO_TNH-NEXT:   (select (result (ref null $struct))
+  ;; NO_TNH-NEXT:    (call $get-null)
+  ;; NO_TNH-NEXT:    (call $get-ref)
+  ;; NO_TNH-NEXT:    (call $get-i32)
+  ;; NO_TNH-NEXT:   )
+  ;; NO_TNH-NEXT:   (i32.const 2)
+  ;; NO_TNH-NEXT:  )
+  ;; NO_TNH-NEXT: )
+  (func $select.arm.null.effects
+    ;; As above but there are conflicting effects and we must add a local when
+    ;; we optimize.
+    (struct.set $struct 0
+      (select (result (ref null $struct))
+        (call $get-ref)
+        (call $get-null)
+        (call $get-i32)
+      )
+      (i32.const 1)
+    )
+    (struct.set $struct 0
+      (select (result (ref null $struct))
+        (call $get-null)
+        (call $get-ref)
+        (call $get-i32)
+      )
+      (i32.const 2)
+    )
+  )
+
+  ;; TNH:      (func $null.arm.null.effects (type $none_=>_none)
+  ;; TNH-NEXT:  (block ;; (replaces something unreachable we can't emit)
+  ;; TNH-NEXT:   (drop
+  ;; TNH-NEXT:    (block (result nullref)
+  ;; TNH-NEXT:     (drop
+  ;; TNH-NEXT:      (ref.as_non_null
+  ;; TNH-NEXT:       (ref.null none)
+  ;; TNH-NEXT:      )
+  ;; TNH-NEXT:     )
+  ;; TNH-NEXT:     (block (result nullref)
+  ;; TNH-NEXT:      (drop
+  ;; TNH-NEXT:       (call $get-i32)
+  ;; TNH-NEXT:      )
+  ;; TNH-NEXT:      (ref.null none)
+  ;; TNH-NEXT:     )
+  ;; TNH-NEXT:    )
+  ;; TNH-NEXT:   )
+  ;; TNH-NEXT:   (drop
+  ;; TNH-NEXT:    (i32.const 1)
+  ;; TNH-NEXT:   )
+  ;; TNH-NEXT:   (unreachable)
+  ;; TNH-NEXT:  )
+  ;; TNH-NEXT: )
+  ;; NO_TNH:      (func $null.arm.null.effects (type $none_=>_none)
+  ;; NO_TNH-NEXT:  (struct.set $struct 0
+  ;; NO_TNH-NEXT:   (select (result (ref null $struct))
+  ;; NO_TNH-NEXT:    (ref.as_non_null
+  ;; NO_TNH-NEXT:     (ref.null none)
+  ;; NO_TNH-NEXT:    )
+  ;; NO_TNH-NEXT:    (ref.null none)
+  ;; NO_TNH-NEXT:    (call $get-i32)
+  ;; NO_TNH-NEXT:   )
+  ;; NO_TNH-NEXT:   (i32.const 1)
+  ;; NO_TNH-NEXT:  )
+  ;; NO_TNH-NEXT: )
+  (func $null.arm.null.effects
+    ;; Verify we do not error on a null reference in a select, even if cast to
+    ;; non-null.
+    (struct.set $struct 0
+      (select (result (ref null $struct))
+        (ref.as_non_null
+          (ref.null none)
+        )
+        (ref.null none)
+        (call $get-i32)
+      )
+      (i32.const 1)
+    )
+  )
+
+  ;; Helper functions.
+
+  ;; TNH:      (func $get-i32 (type $none_=>_i32) (result i32)
+  ;; TNH-NEXT:  (unreachable)
+  ;; TNH-NEXT: )
+  ;; NO_TNH:      (func $get-i32 (type $none_=>_i32) (result i32)
+  ;; NO_TNH-NEXT:  (unreachable)
+  ;; NO_TNH-NEXT: )
+  (func $get-i32 (result i32)
+    (unreachable)
+  )
+  ;; TNH:      (func $get-ref (type $none_=>_ref|$struct|) (result (ref $struct))
+  ;; TNH-NEXT:  (unreachable)
+  ;; TNH-NEXT: )
+  ;; NO_TNH:      (func $get-ref (type $none_=>_ref|$struct|) (result (ref $struct))
+  ;; NO_TNH-NEXT:  (unreachable)
+  ;; NO_TNH-NEXT: )
+  (func $get-ref (result (ref $struct))
+    (unreachable)
+  )
+  ;; TNH:      (func $get-null (type $none_=>_nullref) (result nullref)
+  ;; TNH-NEXT:  (unreachable)
+  ;; TNH-NEXT: )
+  ;; NO_TNH:      (func $get-null (type $none_=>_nullref) (result nullref)
+  ;; NO_TNH-NEXT:  (unreachable)
+  ;; NO_TNH-NEXT: )
+  (func $get-null (result (ref null none))
+    (unreachable)
+  )
 )