3 files changed, 24 insertions, 19 deletions
diff --git a/src/passes/Memory64Lowering.cpp b/src/passes/Memory64Lowering.cpp
index 7e581c22b..15f3a7719 100644
--- a/src/passes/Memory64Lowering.cpp
+++ b/src/passes/Memory64Lowering.cpp
@@ -29,10 +29,6 @@ namespace wasm {
 
 struct Memory64Lowering : public WalkerPass<PostWalker<Memory64Lowering>> {
 
-  void run(PassRunner* runner, Module* module) override {
-    super::run(runner, module);
-  }
-
   void wrapAddress64(Expression*& ptr, Name memoryName) {
     if (ptr->type == Type::unreachable) {
       return;
diff --git a/src/wasm-binary.h b/src/wasm-binary.h
index aef1767bb..c229e4666 100644
--- a/src/wasm-binary.h
+++ b/src/wasm-binary.h
@@ -1482,6 +1482,9 @@ public:
   Name getGlobalName(Index index);
   Name getTagName(Index index);
 
+  // gets a memory in the combined import+defined space
+  Memory* getMemory(Index index);
+
   void getResizableLimits(Address& initial,
                           Address& max,
                           bool& shared,
diff --git a/src/wasm/wasm-binary.cpp b/src/wasm/wasm-binary.cpp
index fe827f291..87fb0e293 100644
--- a/src/wasm/wasm-binary.cpp
+++ b/src/wasm/wasm-binary.cpp
@@ -2220,6 +2220,16 @@ Name WasmBinaryBuilder::getTagName(Index index) {
   return wasm.tags[index]->name;
 }
 
+Memory* WasmBinaryBuilder::getMemory(Index index) {
+  Index numMemoryImports = memoryImports.size();
+  if (index < numMemoryImports) {
+    return memoryImports[index];
+  } else if (index - numMemoryImports < memories.size()) {
+    return memories[index - numMemoryImports].get();
+  }
+  throwError("Memory index out of range.");
+}
+
 void WasmBinaryBuilder::getResizableLimits(Address& initial,
                                            Address& max,
                                            bool& shared,
@@ -3047,16 +3057,7 @@ void WasmBinaryBuilder::readDataSegments() {
     if (flags & BinaryConsts::HasIndex) {
       memIdx = getU32LEB();
     }
-    Memory* memory = nullptr;
-    Index numMemoryImports = memoryImports.size();
-    if (memIdx < numMemoryImports) {
-      memory = memoryImports[memIdx];
-    } else if (memIdx - numMemoryImports < memories.size()) {
-      memory = memories[memIdx - numMemoryImports].get();
-    }
-    if (!memory) {
-      throwError("Memory index out of range while reading data segments.");
-    }
+    auto* memory = getMemory(memIdx);
     curr->memory = memory->name;
     if (!curr->isPassive) {
       curr->offset = readExpression();
@@ -6515,17 +6516,22 @@ void WasmBinaryBuilder::visitReturn(Return* curr) {
 
 void WasmBinaryBuilder::visitMemorySize(MemorySize* curr) {
   BYN_TRACE("zz node: MemorySize\n");
-  Index memIdx = getU32LEB();
+  Index index = getU32LEB();
+  if (getMemory(index)->is64()) {
+    curr->make64();
+  }
   curr->finalize();
-  memoryRefs[memIdx].push_back(&curr->memory);
+  memoryRefs[index].push_back(&curr->memory);
 }
 
 void WasmBinaryBuilder::visitMemoryGrow(MemoryGrow* curr) {
   BYN_TRACE("zz node: MemoryGrow\n");
   curr->delta = popNonVoidExpression();
-  Index memIdx = getU32LEB();
-  curr->finalize();
-  memoryRefs[memIdx].push_back(&curr->memory);
+  Index index = getU32LEB();
+  if (getMemory(index)->is64()) {
+    curr->make64();
+  }
+  memoryRefs[index].push_back(&curr->memory);
 }
 
 void WasmBinaryBuilder::visitNop(Nop* curr) { BYN_TRACE("zz node: Nop\n"); }