[Memory64] (#3302)

Fixed bug in memory64-lowering pass for memory.size/grow

3 files changed, 32 insertions, 13 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,