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)