From 57ab4d7b0dd53c5c4c2440d02e63d6b190c4e2c6 Mon Sep 17 00:00:00 2001
From: Thomas Lively <tlively@google.com>
Date: Tue, 11 Mar 2025 16:19:23 -0700
Subject: [PATCH] Fix optimization of casts to exact null types

The logic for optimizing ref.casts that are known to succeed did not
account for possible casts to exact null types, leading to it producing
invalid IR. Fix it and add a test.
---
 src/passes/OptimizeInstructions.cpp           | 21 +++++++++----
 .../passes/optimize-instructions-exact.wast   | 30 ++++++++++++-------
 2 files changed, 34 insertions(+), 17 deletions(-)

diff --git a/src/passes/OptimizeInstructions.cpp b/src/passes/OptimizeInstructions.cpp
index 0987ed845d6..90dcf5c5d64 100644
--- a/src/passes/OptimizeInstructions.cpp
+++ b/src/passes/OptimizeInstructions.cpp
@@ -2277,10 +2277,14 @@ struct OptimizeInstructions
           // emit a null check.
           bool needsNullCheck = ref->type.getNullability() == Nullable &&
                                 curr->type.getNullability() == NonNullable;
+          // Same with exactness.
+          bool needsExactCast = ref->type.getExactness() == Inexact &&
+                                curr->type.getExactness() == Exact;
           // If the best value to propagate is the argument to the cast, we can
           // simply remove the cast (or downgrade it to a null check if
-          // necessary).
-          if (ref == curr->ref) {
+          // necessary). This does not work if we need a cast to prove
+          // exactness.
+          if (ref == curr->ref && !needsExactCast) {
             if (needsNullCheck) {
               replaceCurrent(builder.makeRefAs(RefAsNonNull, curr->ref));
             } else {
@@ -2289,17 +2293,22 @@ struct OptimizeInstructions
             return;
           }
           // Otherwise we can't just remove the cast and replace it with `ref`
-          // because the intermediate expressions might have had side effects.
-          // We can replace the cast with a drop followed by a direct return of
-          // the value, though.
+          // because the intermediate expressions might have had side effects or
+          // we need to check exactness. We can replace the cast with a drop
+          // followed by a direct return of the value, though.
           if (ref->type.isNull()) {
+            // TODO: Remove this once we type ref.null as exact.
+            if (needsExactCast) {
+              return;
+            }
+
             // We can materialize the resulting null value directly.
             //
             // The type must be nullable for us to do that, which it normally
             // would be, aside from the interesting corner case of
             // uninhabitable types:
             //
-            //  (ref.cast func
+            //  (ref.cast (ref func)
             //    (block (result (ref nofunc))
             //      (unreachable)
             //    )
diff --git a/test/lit/passes/optimize-instructions-exact.wast b/test/lit/passes/optimize-instructions-exact.wast
index ae35380c9a1..6e03c9a4977 100644
--- a/test/lit/passes/optimize-instructions-exact.wast
+++ b/test/lit/passes/optimize-instructions-exact.wast
@@ -6,19 +6,27 @@
 ;; RUN: wasm-opt %s -all --optimize-instructions -S -o - | filecheck %s
 
 (module
- ;; CHECK:      (func $cast-to-exact-none (type $0) (param $0 anyref)
- ;; CHECK-NEXT:  (drop
- ;; CHECK-NEXT:   (ref.cast (exact nullref)
- ;; CHECK-NEXT:    (local.get $0)
- ;; CHECK-NEXT:   )
+ ;; CHECK:      (func $cast-any-to-exact-none (type $0) (param $0 anyref) (result (exact nullref))
+ ;; CHECK-NEXT:  (ref.cast (exact nullref)
+ ;; CHECK-NEXT:   (local.get $0)
  ;; CHECK-NEXT:  )
  ;; CHECK-NEXT: )
- (func $cast-to-exact-none (param anyref)
-  (drop
-   ;; This will not be changed, but should not trigger an assertion.
-   (ref.cast (exact nullref)
-    (local.get 0)
-   )
+ (func $cast-any-to-exact-none (param anyref) (result (exact nullref))
+  ;; This will not be changed, but should not trigger an assertion.
+  (ref.cast (exact nullref)
+   (local.get 0)
+  )
+ )
+ ;; CHECK:      (func $cast-null-to-exact-none (type $1) (result (exact nullref))
+ ;; CHECK-NEXT:  (local $0 nullref)
+ ;; CHECK-NEXT:  (ref.cast (exact nullref)
+ ;; CHECK-NEXT:   (local.get $0)
+ ;; CHECK-NEXT:  )
+ ;; CHECK-NEXT: )
+ (func $cast-null-to-exact-none (result (exact nullref))
+  (local nullref)
+  (ref.cast (exact nullref)
+   (local.get 0)
   )
  )
 )