Fix MemoryPacking handling of array.init_data (WebAssembly#5644)

tlively · web-flow · commit 0578862dc556 · 2023-04-07T14:06:10.000-07:00
Do not optimize out or split segments that are referred to array.init_data
instructions. Fixes a bug where segments could get optimized out, producing
invalid modules. Doing the work to actually split segments used by
array.init_data is left for the future.

Also fix a latent UBSan failure revealed by the new test case.
diff --git a/src/passes/MemoryPacking.cpp b/src/passes/MemoryPacking.cpp
@@ -278,8 +278,8 @@ bool MemoryPacking::canSplit(const std::unique_ptr<DataSegment>& segment,
           return false;
         }
       }
-    } else if (referrer->is<ArrayNewSeg>()) {
-      // TODO: Split segments referenced by array.new_data instructions.
+    } else if (referrer->is<ArrayNewSeg>() || referrer->is<ArrayInit>()) {
+      // TODO: Split segments referenced by GC instructions.
       return false;
     }
   }
@@ -479,8 +479,10 @@ void MemoryPacking::getSegmentReferrers(Module* module,
           referrers[curr->segment].push_back(curr);
         }
       }
-      void doWalkFunction(Function* func) {
-        super::doWalkFunction(func);
+      void visitArrayInit(ArrayInit* curr) {
+        if (curr->op == InitData) {
+          referrers[curr->segment].push_back(curr);
+        }
       }
     } collector(referrers);
     collector.walkFunctionInModule(func, module);
@@ -586,7 +588,7 @@ void MemoryPacking::createSplitSegments(
                                          segment->memory,
                                          segment->isPassive,
                                          offset,
-                                         &segment->data[range.start],
+                                         segment->data.data() + range.start,
                                          range.end - range.start);
     curr->hasExplicitName = hasExplicitName;
     packed.push_back(std::move(curr));
diff --git a/test/lit/passes/memory-packing_all-features.wast b/test/lit/passes/memory-packing_all-features.wast
@@ -2389,3 +2389,31 @@
   )
  )
 )
+(module
+ ;; CHECK:      (type $array (array (mut i32)))
+ (type $array (array (mut i32)))
+ ;; CHECK:      (type $ref|$array|_i32_i32_i32_=>_none (func (param (ref $array) i32 i32 i32)))
+
+ ;; CHECK:      (memory $0 (shared 16 17))
+ (memory $0 (shared 16 17))
+ ;; CHECK:      (data $0 "")
+ (data $0 "")
+ ;; CHECK:      (func $0 (type $ref|$array|_i32_i32_i32_=>_none) (param $0 (ref $array)) (param $1 i32) (param $2 i32) (param $3 i32)
+ ;; CHECK-NEXT:  (array.init_data $array $0
+ ;; CHECK-NEXT:   (local.get $0)
+ ;; CHECK-NEXT:   (i32.const 0)
+ ;; CHECK-NEXT:   (i32.const 0)
+ ;; CHECK-NEXT:   (i32.const 0)
+ ;; CHECK-NEXT:  )
+ ;; CHECK-NEXT: )
+ (func $0 (param $0 (ref $array)) (param $1 i32) (param $2 i32) (param $3 i32)
+  ;; test that we do not improperly optimize out segments referred to by
+  ;; array.init_data instructions.
+  (array.init_data $array $0
+   (local.get $0)
+   (i32.const 0)
+   (i32.const 0)
+   (i32.const 0)
+  )
+ )
+)

Original file line number	Diff line number	Diff line change
`@@ -278,8 +278,8 @@ bool MemoryPacking::canSplit(const std::unique_ptr<DataSegment>& segment,`
`278`	`278`	`return false;`
`279`	`279`	`}`
`280`	`280`	`}`
`281`		`- } else if (referrer->is<ArrayNewSeg>()) {`
`282`		`- // TODO: Split segments referenced by array.new_data instructions.`
	`281`	`+ } else if (referrer->is<ArrayNewSeg>() \|\| referrer->is<ArrayInit>()) {`
	`282`	`+ // TODO: Split segments referenced by GC instructions.`
`283`	`283`	`return false;`
`284`	`284`	`}`
`285`	`285`	`}`
`@@ -479,8 +479,10 @@ void MemoryPacking::getSegmentReferrers(Module* module,`
`479`	`479`	`referrers[curr->segment].push_back(curr);`
`480`	`480`	`}`
`481`	`481`	`}`
`482`		`- void doWalkFunction(Function* func) {`
`483`		`- super::doWalkFunction(func);`
	`482`	`+ void visitArrayInit(ArrayInit* curr) {`
	`483`	`+ if (curr->op == InitData) {`
	`484`	`+ referrers[curr->segment].push_back(curr);`
	`485`	`+ }`
`484`	`486`	`}`
`485`	`487`	`} collector(referrers);`
`486`	`488`	`collector.walkFunctionInModule(func, module);`
`@@ -586,7 +588,7 @@ void MemoryPacking::createSplitSegments(`
`586`	`588`	`segment->memory,`
`587`	`589`	`segment->isPassive,`
`588`	`590`	`offset,`
`589`		`- &segment->data[range.start],`
	`591`	`+ segment->data.data() + range.start,`
`590`	`592`	`range.end - range.start);`
`591`	`593`	`curr->hasExplicitName = hasExplicitName;`
`592`	`594`	`packed.push_back(std::move(curr));`