[NFC] Eagerly create segments when parsing datacount (#6958)

The purpose of the datacount section is to pre-declare how many data
segments there will be so that engines can allocate space for them
and not have to back patch subsequent instructions in the code section
that refer to them. Once we use IRBuilder in the binary parser, we will
have to have the data segments available by the time we parse
instructions that use them, so eagerly construct the data segments when
parsing the datacount section.

5 files changed, 35 insertions, 18 deletions

diff --git a/src/passes/Print.cpp b/src/passes/Print.cpp
index d49bc2974..427bff329 100644
--- a/src/passes/Print.cpp
+++ b/src/passes/Print.cpp
@@ -3238,6 +3238,11 @@ void PrintSExpression::visitMemory(Memory* curr) {
 }
 
 void PrintSExpression::visitDataSegment(DataSegment* curr) {
+  if (!curr->isPassive && !curr->offset) {
+    // This data segment must have been created from the datacount section but
+    // not parsed yet. Skip it.
+    return;
+  }
   doIndent(o, indent);
   o << '(';
   printMajor(o, "data ");
diff --git a/src/wasm/wasm-binary.cpp b/src/wasm/wasm-binary.cpp
index d999141fd..cb9ea3731 100644
--- a/src/wasm/wasm-binary.cpp
+++ b/src/wasm/wasm-binary.cpp
@@ -3273,18 +3273,37 @@ void WasmBinaryReader::processNames() {
 void WasmBinaryReader::readDataSegmentCount() {
   hasDataCount = true;
   dataCount = getU32LEB();
+  // Eagerly create the data segments so they are available during parsing of
+  // the code section.
+  for (size_t i = 0; i < dataCount; ++i) {
+    auto curr = Builder::makeDataSegment();
+    curr->setName(Name::fromInt(i), false);
+    wasm.addDataSegment(std::move(curr));
+  }
 }
 
 void WasmBinaryReader::readDataSegments() {
   auto num = getU32LEB();
+  if (hasDataCount) {
+    if (num != dataCount) {
+      throwError("data count and data sections disagree on size");
+    }
+  } else {
+    // We haven't already created the data segments, so create them now.
+    for (size_t i = 0; i < num; ++i) {
+      auto curr = Builder::makeDataSegment();
+      curr->setName(Name::fromInt(i), false);
+      wasm.addDataSegment(std::move(curr));
+    }
+  }
+  assert(wasm.dataSegments.size() == num);
   for (size_t i = 0; i < num; i++) {
-    auto curr = Builder::makeDataSegment();
+    auto& curr = wasm.dataSegments[i];
     uint32_t flags = getU32LEB();
     if (flags > 2) {
       throwError("bad segment flags, must be 0, 1, or 2, not " +
                  std::to_string(flags));
     }
-    curr->setName(Name::fromInt(i), false);
     curr->isPassive = flags & BinaryConsts::IsPassive;
     if (curr->isPassive) {
       curr->memory = Name();
@@ -3300,7 +3319,6 @@ void WasmBinaryReader::readDataSegments() {
     auto size = getU32LEB();
     auto data = getByteView(size);
     curr->data = {data.begin(), data.end()};
-    wasm.addDataSegment(std::move(curr));
   }
 }
 
diff --git a/test/lit/binary/bad-datacount.test b/test/lit/binary/bad-datacount.test
new file mode 100644
index 000000000..27e0c72a3
--- /dev/null
+++ b/test/lit/binary/bad-datacount.test
@@ -0,0 +1,9 @@
+RUN: not wasm-opt --debug %s.wasm 2>&1 | filecheck %s
+
+;; Check that we get the expected error.
+
+;; CHECK: data count and data sections disagree on size
+
+;; Check that we print out the module rather than hitting an assertion error.
+
+;; CHECK: (module
diff --git a/test/unit/input/bulkmem_bad_datacount.wasm b/test/lit/binary/bad-datacount.test.wasm
index e7b6e4bb7..e7b6e4bb7 100755
--- a/test/unit/input/bulkmem_bad_datacount.wasm
+++ b/test/lit/binary/bad-datacount.test.wasm
diff --git a/test/unit/test_datacount.py b/test/unit/test_datacount.py
deleted file mode 100644
index b996865a5..000000000
--- a/test/unit/test_datacount.py
+++ /dev/null
@@ -1,15 +0,0 @@
-from scripts.test import shared
-from . import utils
-
-
-class DataCountTest(utils.BinaryenTestCase):
-    def test_datacount(self):
-        self.roundtrip('bulkmem_data.wasm')
-
-    def test_bad_datacount(self):
-        path = self.input_path('bulkmem_bad_datacount.wasm')
-        p = shared.run_process(shared.WASM_OPT + ['-g', '-o', '-', path],
-                               check=False, capture_output=True)
-        self.assertNotEqual(p.returncode, 0)
-        self.assertIn('Number of segments does not agree with DataCount section',
-                      p.stderr)