[MemoryPacking] Emit the correct segment indices on memory.init (#3239)

This PR fixes a bug in which the segment index of a memory.init instruction was
incorrect in some circumstances. Specifically, the first segment index used in
output memory.init instructions was always the index of the first segment
created from splitting up the corresponding input segment. This was incorrect
when the input memory.init had an offset that caused it to skip over that first
emitted segment so that the first output memory.init should have referred to a
subsequent output segment.

Fixes #3225.

3 files changed, 44 insertions, 2 deletions

diff --git a/src/passes/MemoryPacking.cpp b/src/passes/MemoryPacking.cpp
index 4a183e7d7..27168942a 100644
--- a/src/passes/MemoryPacking.cpp
+++ b/src/passes/MemoryPacking.cpp
@@ -514,10 +514,16 @@ void MemoryPacking::createReplacements(Module* module,
     size_t start = init->offset->cast<Const>()->value.geti32();
     size_t end = start + init->size->cast<Const>()->value.geti32();
 
+    // Segment index used in emitted memory.init instructions
+    size_t initIndex = segmentIndex;
+
     // Index of the range from which this memory.init starts reading
     size_t firstRangeIdx = 0;
     while (firstRangeIdx < ranges.size() &&
            ranges[firstRangeIdx].end <= start) {
+      if (!ranges[firstRangeIdx].isZero) {
+        ++initIndex;
+      }
       ++firstRangeIdx;
     }
 
@@ -568,7 +574,6 @@ void MemoryPacking::createReplacements(Module* module,
 
     size_t bytesWritten = 0;
 
-    size_t initIndex = segmentIndex;
     for (size_t i = firstRangeIdx; i < ranges.size() && ranges[i].start < end;
          ++i) {
       auto& range = ranges[i];
diff --git a/test/passes/memory-packing_all-features.txt b/test/passes/memory-packing_all-features.txt
index af18ea25f..bc9139f74 100644
--- a/test/passes/memory-packing_all-features.txt
+++ b/test/passes/memory-packing_all-features.txt
@@ -1505,3 +1505,27 @@
   )
  )
 )
+(module
+ (type $none_=>_none (func))
+ (import "env" "memory" (memory $mimport$0 1 1))
+ (data passive "skipped")
+ (data passive "included")
+ (global $__mem_segment_drop_state_0 (mut i32) (i32.const 0))
+ (export "func_54" (func $0))
+ (func $0
+  (if
+   (global.get $__mem_segment_drop_state_0)
+   (unreachable)
+  )
+  (memory.fill
+   (i32.const 0)
+   (i32.const 0)
+   (i32.const 30)
+  )
+  (memory.init 1
+   (i32.const 30)
+   (i32.const 0)
+   (i32.const 8)
+  )
+ )
+)
diff --git a/test/passes/memory-packing_all-features.wast b/test/passes/memory-packing_all-features.wast
index 309ee79ad..3ee0bd5e7 100644
--- a/test/passes/memory-packing_all-features.wast
+++ b/test/passes/memory-packing_all-features.wast
@@ -494,4 +494,17 @@
     )
     (data.drop 0)
   )
-)
\ No newline at end of file
+)
+
+(module
+ (import "env" "memory" (memory 1 1))
+ (data passive "skipped\00\00\00\00\00\00\00\00\00\00\00\00\00\00\00\00\00\00\00\00\00\00\00\00\00\00\00\00\00\00included")
+ (export "func_54" (func $0))
+ (func $0
+  (memory.init 0
+   (i32.const 0)
+   (i32.const 7)
+   (i32.const 38)
+  )
+ )
+)