1 files changed, 60 insertions, 38 deletions

diff --git a/src/passes/Memory64Lowering.cpp b/src/passes/Memory64Lowering.cpp
index 4c2770e3a..7e581c22b 100644
--- a/src/passes/Memory64Lowering.cpp
+++ b/src/passes/Memory64Lowering.cpp
@@ -30,79 +30,101 @@ namespace wasm {
 struct Memory64Lowering : public WalkerPass<PostWalker<Memory64Lowering>> {
 
   void run(PassRunner* runner, Module* module) override {
-    if (module->memory.is64()) {
-      super::run(runner, module);
-    }
+    super::run(runner, module);
   }
 
-  void wrapAddress64(Expression*& ptr) {
+  void wrapAddress64(Expression*& ptr, Name memoryName) {
     if (ptr->type == Type::unreachable) {
       return;
     }
     auto& module = *getModule();
-    assert(module.memory.is64());
-    assert(ptr->type == Type::i64);
-    Builder builder(module);
-    ptr = builder.makeUnary(UnaryOp::WrapInt64, ptr);
+    auto memory = module.getMemory(memoryName);
+    if (memory->is64()) {
+      assert(ptr->type == Type::i64);
+      Builder builder(module);
+      ptr = builder.makeUnary(UnaryOp::WrapInt64, ptr);
+    }
   }
 
-  void extendAddress64(Expression*& ptr) {
+  void extendAddress64(Expression*& ptr, Name memoryName) {
     if (ptr->type == Type::unreachable) {
       return;
     }
     auto& module = *getModule();
-    assert(module.memory.is64());
-    assert(ptr->type == Type::i64);
-    ptr->type = Type::i32;
-    Builder builder(module);
-    ptr = builder.makeUnary(UnaryOp::ExtendUInt32, ptr);
+    auto memory = module.getMemory(memoryName);
+    if (memory->is64()) {
+      assert(ptr->type == Type::i64);
+      ptr->type = Type::i32;
+      Builder builder(module);
+      ptr = builder.makeUnary(UnaryOp::ExtendUInt32, ptr);
+    }
   }
 
-  void visitLoad(Load* curr) { wrapAddress64(curr->ptr); }
+  void visitLoad(Load* curr) { wrapAddress64(curr->ptr, curr->memory); }
 
-  void visitStore(Store* curr) { wrapAddress64(curr->ptr); }
+  void visitStore(Store* curr) { wrapAddress64(curr->ptr, curr->memory); }
 
   void visitMemorySize(MemorySize* curr) {
-    auto size = static_cast<Expression*>(curr);
-    extendAddress64(size);
-    curr->ptrType = Type::i32;
-    replaceCurrent(size);
+    auto& module = *getModule();
+    auto memory = module.getMemory(curr->memory);
+    if (memory->is64()) {
+      auto size = static_cast<Expression*>(curr);
+      extendAddress64(size, curr->memory);
+      curr->ptrType = Type::i32;
+      replaceCurrent(size);
+    }
   }
 
   void visitMemoryGrow(MemoryGrow* curr) {
-    wrapAddress64(curr->delta);
-    auto size = static_cast<Expression*>(curr);
-    extendAddress64(size);
-    curr->ptrType = Type::i32;
-    replaceCurrent(size);
+    auto& module = *getModule();
+    auto memory = module.getMemory(curr->memory);
+    if (memory->is64()) {
+      wrapAddress64(curr->delta, curr->memory);
+      auto size = static_cast<Expression*>(curr);
+      extendAddress64(size, curr->memory);
+      curr->ptrType = Type::i32;
+      replaceCurrent(size);
+    }
   }
 
-  void visitMemoryInit(MemoryInit* curr) { wrapAddress64(curr->dest); }
+  void visitMemoryInit(MemoryInit* curr) {
+    wrapAddress64(curr->dest, curr->memory);
+  }
 
   void visitMemoryFill(MemoryFill* curr) {
-    wrapAddress64(curr->dest);
-    wrapAddress64(curr->size);
+    wrapAddress64(curr->dest, curr->memory);
+    wrapAddress64(curr->size, curr->memory);
   }
 
   void visitMemoryCopy(MemoryCopy* curr) {
-    wrapAddress64(curr->dest);
-    wrapAddress64(curr->source);
-    wrapAddress64(curr->size);
+    wrapAddress64(curr->dest, curr->destMemory);
+    wrapAddress64(curr->source, curr->sourceMemory);
+    wrapAddress64(curr->size, curr->destMemory);
   }
 
-  void visitAtomicRMW(AtomicRMW* curr) { wrapAddress64(curr->ptr); }
+  void visitAtomicRMW(AtomicRMW* curr) {
+    wrapAddress64(curr->ptr, curr->memory);
+  }
 
-  void visitAtomicCmpxchg(AtomicCmpxchg* curr) { wrapAddress64(curr->ptr); }
+  void visitAtomicCmpxchg(AtomicCmpxchg* curr) {
+    wrapAddress64(curr->ptr, curr->memory);
+  }
 
-  void visitAtomicWait(AtomicWait* curr) { wrapAddress64(curr->ptr); }
+  void visitAtomicWait(AtomicWait* curr) {
+    wrapAddress64(curr->ptr, curr->memory);
+  }
 
-  void visitAtomicNotify(AtomicNotify* curr) { wrapAddress64(curr->ptr); }
+  void visitAtomicNotify(AtomicNotify* curr) {
+    wrapAddress64(curr->ptr, curr->memory);
+  }
 
   void visitMemory(Memory* memory) {
     // This is visited last.
-    memory->indexType = Type::i32;
-    if (memory->hasMax() && memory->max > Memory::kMaxSize32) {
-      memory->max = Memory::kMaxSize32;
+    if (memory->is64()) {
+      memory->indexType = Type::i32;
+      if (memory->hasMax() && memory->max > Memory::kMaxSize32) {
+        memory->max = Memory::kMaxSize32;
+      }
     }
   }