5 files changed, 40 insertions, 19 deletions
diff --git a/src/passes/Memory64Lowering.cpp b/src/passes/Memory64Lowering.cpp
index 7905c944f..cee174b63 100644
--- a/src/passes/Memory64Lowering.cpp
+++ b/src/passes/Memory64Lowering.cpp
@@ -35,7 +35,7 @@ struct Memory64Lowering : public WalkerPass<PostWalker<Memory64Lowering>> {
     }
   }
 
-  void WrapAddress64(Expression*& ptr) {
+  void wrapAddress64(Expression*& ptr) {
     if (ptr->type == Type::unreachable) {
       return;
     }
@@ -46,34 +46,48 @@ struct Memory64Lowering : public WalkerPass<PostWalker<Memory64Lowering>> {
     ptr = builder.makeUnary(UnaryOp::WrapInt64, ptr);
   }
 
-  void visitLoad(Load* curr) { WrapAddress64(curr->ptr); }
+  void extendAddress64(Expression*& ptr) {
+    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);
+  }
+
+  void visitLoad(Load* curr) { wrapAddress64(curr->ptr); }
 
-  void visitStore(Store* curr) { WrapAddress64(curr->ptr); }
+  void visitStore(Store* curr) { wrapAddress64(curr->ptr); }
 
   void visitMemorySize(MemorySize* curr) {
     auto size = static_cast<Expression*>(curr);
-    WrapAddress64(size);
+    extendAddress64(size);
+    curr->ptrType = Type::i32;
     replaceCurrent(size);
   }
 
   void visitMemoryGrow(MemoryGrow* curr) {
-    WrapAddress64(curr->delta);
+    wrapAddress64(curr->delta);
     auto size = static_cast<Expression*>(curr);
-    WrapAddress64(size);
+    extendAddress64(size);
+    curr->ptrType = Type::i32;
     replaceCurrent(size);
   }
 
-  void visitMemoryInit(MemoryInit* curr) { WrapAddress64(curr->dest); }
+  void visitMemoryInit(MemoryInit* curr) { wrapAddress64(curr->dest); }
 
   void visitMemoryFill(MemoryFill* curr) {
-    WrapAddress64(curr->dest);
-    WrapAddress64(curr->size);
+    wrapAddress64(curr->dest);
+    wrapAddress64(curr->size);
   }
 
   void visitMemoryCopy(MemoryCopy* curr) {
-    WrapAddress64(curr->dest);
-    WrapAddress64(curr->source);
-    WrapAddress64(curr->size);
+    wrapAddress64(curr->dest);
+    wrapAddress64(curr->source);
+    wrapAddress64(curr->size);
   }
 
   void visitMemory(Memory* memory) {
diff --git a/src/wasm-binary.h b/src/wasm-binary.h
index 8d7abfd70..4f10ccfa1 100644
--- a/src/wasm-binary.h
+++ b/src/wasm-binary.h
@@ -1313,6 +1313,7 @@ public:
   uint64_t getU64LEB();
   int32_t getS32LEB();
   int64_t getS64LEB();
+  uint64_t getUPtrLEB();
   Type getType();
   HeapType getHeapType();
   Type getConcreteType();
diff --git a/src/wasm/wasm-binary.cpp b/src/wasm/wasm-binary.cpp
index 46ce205b7..18ae724f0 100644
--- a/src/wasm/wasm-binary.cpp
+++ b/src/wasm/wasm-binary.cpp
@@ -1215,6 +1215,10 @@ int64_t WasmBinaryBuilder::getS64LEB() {
   return ret.value;
 }
 
+uint64_t WasmBinaryBuilder::getUPtrLEB() {
+  return wasm.memory.is64() ? getU64LEB() : getU32LEB();
+}
+
 Type WasmBinaryBuilder::getType() {
   int type = getS32LEB();
   // Single value types are negative; signature indices are non-negative
@@ -3150,7 +3154,7 @@ void WasmBinaryBuilder::readMemoryAccess(Address& alignment, Address& offset) {
     throwError("Alignment must be of a reasonable size");
   }
   alignment = Bits::pow2(rawAlignment);
-  offset = getU32LEB();
+  offset = getUPtrLEB();
 }
 
 bool WasmBinaryBuilder::maybeVisitLoad(Expression*& out,
diff --git a/test/passes/memory64-lowering_enable-memory64_enable-bulk-memory.txt b/test/passes/memory64-lowering_enable-memory64_enable-bulk-memory.txt
index 3553c393c..84ded6ae5 100644
--- a/test/passes/memory64-lowering_enable-memory64_enable-bulk-memory.txt
+++ b/test/passes/memory64-lowering_enable-memory64_enable-bulk-memory.txt
@@ -3,6 +3,7 @@
  (memory $0 1 1)
  (data (i32.const 0) "\00\00\00\00\00\00\00\00\00\00")
  (func $func_1
+  (local $0 i64)
   (drop
    (i32.load
     (i32.wrap_i64
@@ -122,13 +123,13 @@
    )
    (unreachable)
   )
-  (drop
-   (i32.wrap_i64
+  (local.set $0
+   (i64.extend_i32_u
     (memory.size)
    )
   )
-  (drop
-   (i32.wrap_i64
+  (local.set $0
+   (i64.extend_i32_u
     (memory.grow
      (i32.wrap_i64
       (i64.const 1)
diff --git a/test/passes/memory64-lowering_enable-memory64_enable-bulk-memory.wast b/test/passes/memory64-lowering_enable-memory64_enable-bulk-memory.wast
index 433d509c3..f70444f49 100644
--- a/test/passes/memory64-lowering_enable-memory64_enable-bulk-memory.wast
+++ b/test/passes/memory64-lowering_enable-memory64_enable-bulk-memory.wast
@@ -2,6 +2,7 @@
  (memory $0 i64 1 1)
  (data (i32.const 0) "\00\00\00\00\00\00\00\00\00\00")
  (func $func_1
+  (local i64)
   (drop (i32.load (i64.const 4)))
   (drop (i32.load align=1 (i64.const 4)))
   (drop (i32.load align=2 (i64.const 4)))
@@ -21,8 +22,8 @@
   (i32.store offset=100 align=4 (i64.const 4) (i32.const 8))
   (i32.store offset=100 align=1 (unreachable) (i32.const 8))
   (i32.store offset=100 align=1 (i64.const 4) (unreachable))
-  (drop (memory.size))
-  (drop (memory.grow (i64.const 1)))
+  (local.set 0 (memory.size))
+  (local.set 0 (memory.grow (i64.const 1)))
   (memory.init 0 (i64.const 1) (i32.const 2) (i32.const 3))
   (memory.fill (i64.const 1) (i32.const 2) (i64.const 3))
   (memory.copy (i64.const 1) (i64.const 2) (i64.const 3))