Flat IR: Allow ref.as_non_null in nested positions (#3732)

We can't disallow it, as its result is non-null which we can't spill to a local.

This may cause issues eventually in the combination of GC + flatten, but
I don't expect it to. If it does we may need to revisit.

2 files changed, 53 insertions, 3 deletions

diff --git a/src/ir/flat.h b/src/ir/flat.h
index 706ca5a86..273acc42f 100644
--- a/src/ir/flat.h
+++ b/src/ir/flat.h
@@ -43,8 +43,8 @@
 // making the AST have these properties:
 //
 //  1. Aside from a local.set, the operands of an instruction must be a
-//     local.get, a const, or an unreachable. Anything else is written
-//     to a local earlier.
+//     local.get, a const, an unreachable, or a ref.as_non_null. Anything else
+//     is written to a local earlier.
 //  2. Disallow control flow (block, loop, if, and try) return values, and do
 //     not allow the function body to have a concrete type, i.e., do not use
 //     control flow to pass around values.
@@ -54,6 +54,10 @@
 //     values is prohibited already, but e.g. a block ending in unreachable,
 //     which can normally be nested, is also disallowed).
 //
+// Note: ref.as_non_null must be allowed in a nested position because we cannot
+// spill it to a local - the result is non-null, which is not allowable in a
+// local.
+//
 
 #ifndef wasm_ir_flat_h
 #define wasm_ir_flat_h
@@ -82,8 +86,11 @@ inline void verifyFlatness(Function* func) {
                "set values cannot be control flow");
       } else {
         for (auto* child : ChildIterator(curr)) {
+          bool isRefAsNonNull =
+            child->is<RefAs>() && child->cast<RefAs>()->op == RefAsNonNull;
           verify(Properties::isConstantExpression(child) ||
-                   child->is<LocalGet>() || child->is<Unreachable>(),
+                   child->is<LocalGet>() || child->is<Unreachable>() ||
+                   isRefAsNonNull,
                  "instructions must only have constant expressions, local.get, "
                  "or unreachable as children");
         }
diff --git a/test/lit/passes/opt_flatten.wast b/test/lit/passes/opt_flatten.wast
new file mode 100644
index 000000000..e3cf103f3
--- /dev/null
+++ b/test/lit/passes/opt_flatten.wast
@@ -0,0 +1,43 @@
+;; NOTE: Assertions have been generated by update_lit_checks.py and should not be edited.
+
+;; Optimize the code, flatten, and run a pass that will verify it is properly
+;; flat.
+;; RUN: wasm-opt %s -all -O1 --flatten --rereloop -S -o - | filecheck %s
+
+(module
+ (export "foo" (func $foo))
+ ;; CHECK:      (func $foo (result funcref)
+ ;; CHECK-NEXT:  (local $0 funcref)
+ ;; CHECK-NEXT:  (local $1 (ref null $none_=>_funcref))
+ ;; CHECK-NEXT:  (local $2 (ref null $none_=>_funcref))
+ ;; CHECK-NEXT:  (local $3 i32)
+ ;; CHECK-NEXT:  (block
+ ;; CHECK-NEXT:   (local.set $0
+ ;; CHECK-NEXT:    (call $foo)
+ ;; CHECK-NEXT:   )
+ ;; CHECK-NEXT:   (drop
+ ;; CHECK-NEXT:    (local.get $0)
+ ;; CHECK-NEXT:   )
+ ;; CHECK-NEXT:   (local.set $1
+ ;; CHECK-NEXT:    (ref.func $foo)
+ ;; CHECK-NEXT:   )
+ ;; CHECK-NEXT:   (local.set $2
+ ;; CHECK-NEXT:    (ref.as_non_null
+ ;; CHECK-NEXT:     (local.get $1)
+ ;; CHECK-NEXT:    )
+ ;; CHECK-NEXT:   )
+ ;; CHECK-NEXT:   (return
+ ;; CHECK-NEXT:    (ref.as_non_null
+ ;; CHECK-NEXT:     (local.get $2)
+ ;; CHECK-NEXT:    )
+ ;; CHECK-NEXT:   )
+ ;; CHECK-NEXT:  )
+ ;; CHECK-NEXT: )
+ (func $foo (result funcref)
+  ;; the resulting ref.as_non_nulls here must be validated as ok in flat mode.
+  (drop
+   (call $foo)
+  )
+  (ref.func $foo)
+ )
+)