[Wasm GC] Handle non-nullable locals in Flatten pass (#3720)

That pass adds lots of new locals, and we need to handle non-nullable ones.

3 files changed, 30 insertions, 0 deletions

diff --git a/src/passes/Flatten.cpp b/src/passes/Flatten.cpp
index 20d7cef23..39c97b133 100644
--- a/src/passes/Flatten.cpp
+++ b/src/passes/Flatten.cpp
@@ -24,6 +24,7 @@
 #include <ir/effects.h>
 #include <ir/flat.h>
 #include <ir/properties.h>
+#include <ir/type-updating.h>
 #include <ir/utils.h>
 #include <pass.h>
 #include <wasm-builder.h>
@@ -316,6 +317,8 @@ struct Flatten
     }
     // the body may have preludes
     curr->body = getPreludesWithExpression(originalBody, curr->body);
+    // New locals we added may be non-nullable.
+    TypeUpdating::handleNonNullableLocals(curr, *getModule());
   }
 
 private:
diff --git a/test/passes/flatten_all-features.txt b/test/passes/flatten_all-features.txt
index efcce6797..99c8cb4bb 100644
--- a/test/passes/flatten_all-features.txt
+++ b/test/passes/flatten_all-features.txt
@@ -2443,3 +2443,20 @@
   (unreachable)
  )
 )
+(module
+ (type $none_=>_none (func))
+ (type $none_=>_funcref (func (result funcref)))
+ (func $0 (result funcref)
+  (local $0 (ref null $none_=>_none))
+  (local.set $0
+   (ref.as_non_null
+    (ref.null $none_=>_none)
+   )
+  )
+  (return
+   (ref.as_non_null
+    (local.get $0)
+   )
+  )
+ )
+)
diff --git a/test/passes/flatten_all-features.wast b/test/passes/flatten_all-features.wast
index 14da24a74..874e91572 100644
--- a/test/passes/flatten_all-features.wast
+++ b/test/passes/flatten_all-features.wast
@@ -1055,3 +1055,13 @@
   )
  )
 )
+;; non-nullable temp vars we add must be handled properly, as non-nullable
+;; locals are not allowed
+(module
+ (type $none_=>_none (func))
+ (func $0 (result funcref)
+  (ref.as_non_null
+   (ref.null $none_=>_none)
+  )
+ )
+)