Do not flatten memory when array.new_data is present (#5650)

Like data.drop etc., it notices data segment identity.

2 files changed, 63 insertions, 12 deletions

diff --git a/src/ir/memory-utils.cpp b/src/ir/memory-utils.cpp
index b21d8f0ad..9ab0e33fc 100644
--- a/src/ir/memory-utils.cpp
+++ b/src/ir/memory-utils.cpp
@@ -27,20 +27,33 @@ bool flatten(Module& wasm) {
   // The presence of any instruction that cares about segment identity is a
   // problem because flattening gets rid of that (when it merges them all into
   // one big segment).
-  ModuleUtils::ParallelFunctionAnalysis<bool> analysis(
-    wasm, [&](Function* func, bool& noticesSegmentIdentity) {
-      if (func->imported()) {
-        return;
-      }
-      noticesSegmentIdentity =
-        FindAll<MemoryInit>(func->body).list.size() > 0 ||
-        FindAll<DataDrop>(func->body).list.size() > 0;
-    });
+  struct Scanner : public WalkerPass<PostWalker<Scanner>> {
+    std::atomic<bool>& noticesSegmentIdentity;
 
-  for (auto& [func, noticesSegmentIdentity] : analysis.map) {
-    if (noticesSegmentIdentity) {
-      return false;
+    Scanner(std::atomic<bool>& noticesSegmentIdentity)
+      : noticesSegmentIdentity(noticesSegmentIdentity) {}
+
+    std::unique_ptr<Pass> create() override {
+      return std::make_unique<Scanner>(noticesSegmentIdentity);
+    }
+
+    void visitMemoryInit(MemoryInit* curr) { noticesSegmentIdentity = true; }
+    void visitDataDrop(DataDrop* curr) { noticesSegmentIdentity = true; }
+    void visitArrayNewSeg(ArrayNewSeg* curr) {
+      if (curr->op == NewData) {
+        noticesSegmentIdentity = true;
+      }
     }
+  };
+
+  std::atomic<bool> noticesSegmentIdentity = false;
+  PassRunner runner(&wasm);
+  Scanner scanner(noticesSegmentIdentity);
+  scanner.setPassRunner(&runner);
+  scanner.run(&wasm);
+  scanner.runOnModuleCode(&runner, &wasm);
+  if (noticesSegmentIdentity) {
+    return false;
   }
 
   auto& dataSegments = wasm.dataSegments;
diff --git a/test/lit/ctor-eval/array_new_data.wast b/test/lit/ctor-eval/array_new_data.wast
new file mode 100644
index 000000000..3dca60207
--- /dev/null
+++ b/test/lit/ctor-eval/array_new_data.wast
@@ -0,0 +1,38 @@
+;; NOTE: Assertions have been generated by update_lit_checks.py --all-items and should not be edited.
+;; RUN: wasm-ctor-eval %s --ctors=test --kept-exports=test --quiet -all -S -o - | filecheck %s
+
+(module
+ ;; CHECK:      (type $none_=>_none (func))
+
+ ;; CHECK:      (type $[i8] (array i8))
+ (type $[i8] (array i8))
+
+ ;; CHECK:      (memory $0 (shared 16 17))
+ (memory $0 (shared 16 17))
+ ;; CHECK:      (data $0 (i32.const 40) "")
+ (data $0 (i32.const 40) "")
+ ;; CHECK:      (data $1 (i32.const 0) "")
+ (data $1 (i32.const 0) "")
+
+ (func "test"
+  ;; An array.new_data cannot be evalled since ctor-eval flattens memory segments
+  ;; atm. In fact the module would not validate as we refer to segment 1 here
+  ;; but after flattening only segment 0 exists.
+  (drop
+   (array.new_data $[i8] $1
+    (i32.const 16)
+    (i32.const 8)
+   )
+  )
+ )
+)
+;; CHECK:      (export "test" (func $0))
+
+;; CHECK:      (func $0 (type $none_=>_none)
+;; CHECK-NEXT:  (drop
+;; CHECK-NEXT:   (array.new_data $[i8] $1
+;; CHECK-NEXT:    (i32.const 16)
+;; CHECK-NEXT:    (i32.const 8)
+;; CHECK-NEXT:   )
+;; CHECK-NEXT:  )
+;; CHECK-NEXT: )