4 files changed, 119 insertions, 167 deletions
diff --git a/src/passes/CMakeLists.txt b/src/passes/CMakeLists.txt
index 6b78e487d..61c1b2f46 100644
--- a/src/passes/CMakeLists.txt
+++ b/src/passes/CMakeLists.txt
@@ -112,7 +112,6 @@ set(passes_SOURCES
   ReorderGlobals.cpp
   ReorderLocals.cpp
   ReReloop.cpp
-  Table64Lowering.cpp
   TrapMode.cpp
   TypeGeneralizing.cpp
   TypeRefining.cpp
diff --git a/src/passes/Memory64Lowering.cpp b/src/passes/Memory64Lowering.cpp
index a3913c759..714f3aa5b 100644
--- a/src/passes/Memory64Lowering.cpp
+++ b/src/passes/Memory64Lowering.cpp
@@ -31,33 +31,58 @@ namespace wasm {
 static Name MEMORY_BASE("__memory_base");
 static Name MEMORY_BASE32("__memory_base32");
 
+static Name TABLE_BASE("__table_base");
+static Name TABLE_BASE32("__table_base32");
+
 struct Memory64Lowering : public WalkerPass<PostWalker<Memory64Lowering>> {
 
-  void wrapAddress64(Expression*& ptr, Name memoryName) {
+  void wrapAddress64(Expression*& ptr,
+                     Name memoryOrTableName,
+                     bool isTable = false) {
     if (ptr->type == Type::unreachable) {
       return;
     }
     auto& module = *getModule();
-    auto* memory = module.getMemory(memoryName);
-    if (memory->is64()) {
+    bool is64 = false;
+    if (isTable) {
+      is64 = module.getTable(memoryOrTableName)->is64();
+    } else {
+      is64 = module.getMemory(memoryOrTableName)->is64();
+    }
+    if (is64) {
       assert(ptr->type == Type::i64);
       ptr = Builder(module).makeUnary(UnaryOp::WrapInt64, ptr);
     }
   }
 
-  void extendAddress64(Expression*& ptr, Name memoryName) {
+  void extendAddress64(Expression*& ptr,
+                       Name memoryOrTableName,
+                       bool isTable = false) {
     if (ptr->type == Type::unreachable) {
       return;
     }
     auto& module = *getModule();
-    auto* memory = module.getMemory(memoryName);
-    if (memory->is64()) {
+    bool is64 = false;
+    if (isTable) {
+      is64 = module.getTable(memoryOrTableName)->is64();
+    } else {
+      is64 = module.getMemory(memoryOrTableName)->is64();
+    }
+    if (is64) {
       assert(ptr->type == Type::i64);
       ptr->type = Type::i32;
       ptr = Builder(module).makeUnary(UnaryOp::ExtendUInt32, ptr);
     }
   }
 
+  void wrapTableAddress64(Expression*& ptr, Name tableName) {
+    return wrapAddress64(ptr, tableName, true);
+  }
+
+  void extendTableAddress64(Expression*& ptr, Name tableName) {
+    return extendAddress64(ptr, tableName, true);
+  }
+
   void visitLoad(Load* curr) { wrapAddress64(curr->ptr, curr->memory); }
 
   void visitStore(Store* curr) { wrapAddress64(curr->ptr, curr->memory); }
@@ -177,14 +202,92 @@ struct Memory64Lowering : public WalkerPass<PostWalker<Memory64Lowering>> {
     }
   }
 
+  void visitTableSize(TableSize* curr) {
+    auto& module = *getModule();
+    auto* table = module.getTable(curr->table);
+    if (table->is64()) {
+      auto* size = static_cast<Expression*>(curr);
+      extendTableAddress64(size, curr->table);
+      replaceCurrent(size);
+    }
+  }
+
+  void visitTableGrow(TableGrow* curr) {
+    auto& module = *getModule();
+    auto* table = module.getTable(curr->table);
+    if (table->is64()) {
+      wrapTableAddress64(curr->delta, curr->table);
+      auto* size = static_cast<Expression*>(curr);
+      extendTableAddress64(size, curr->table);
+      replaceCurrent(size);
+    }
+  }
+
+  void visitTableFill(TableFill* curr) {
+    wrapTableAddress64(curr->dest, curr->table);
+    wrapTableAddress64(curr->size, curr->table);
+  }
+
+  void visitTableCopy(TableCopy* curr) {
+    wrapTableAddress64(curr->dest, curr->destTable);
+    wrapTableAddress64(curr->source, curr->sourceTable);
+    wrapTableAddress64(curr->size, curr->destTable);
+  }
+
+  void visitTableInit(TableInit* curr) {
+    wrapTableAddress64(curr->dest, curr->table);
+  }
+
+  void visitCallIndirect(CallIndirect* curr) {
+    wrapTableAddress64(curr->target, curr->table);
+  }
+
+  void visitElementSegment(ElementSegment* segment) {
+    auto& module = *getModule();
+
+    // Passive segments don't have any offset to update.
+    if (segment->table.isNull() || !module.getTable(segment->table)->is64()) {
+      return;
+    }
+
+    if (auto* c = segment->offset->dynCast<Const>()) {
+      c->value = Literal(static_cast<uint32_t>(c->value.geti64()));
+      c->type = Type::i32;
+    } else if (auto* get = segment->offset->dynCast<GlobalGet>()) {
+      auto* g = module.getGlobal(get->name);
+      if (g->imported() && g->base == TABLE_BASE) {
+        ImportInfo info(module);
+        auto* memoryBase32 = info.getImportedGlobal(g->module, TABLE_BASE32);
+        if (!memoryBase32) {
+          Builder builder(module);
+          memoryBase32 = builder
+                           .makeGlobal(TABLE_BASE32,
+                                       Type::i32,
+                                       builder.makeConst(int32_t(0)),
+                                       Builder::Immutable)
+                           .release();
+          memoryBase32->module = g->module;
+          memoryBase32->base = TABLE_BASE32;
+          module.addGlobal(memoryBase32);
+        }
+        // Use this alternative import when initializing the segment.
+        assert(memoryBase32);
+        get->type = Type::i32;
+        get->name = memoryBase32->name;
+      }
+    } else {
+      WASM_UNREACHABLE("unexpected elem offset");
+    }
+  }
+
   void run(Module* module) override {
     if (!module->features.has(FeatureSet::Memory64)) {
       return;
     }
     Super::run(module);
-    // Don't modify the memories themselves until after the traversal since we
-    // that would require memories to be the last thing that get visited, and
-    // we don't want to depend on that specific ordering.
+    // Don't modify the memories or tables themselves until after the traversal
+    // since we that would require memories to be the last thing that get
+    // visited, and we don't want to depend on that specific ordering.
     for (auto& memory : module->memories) {
       if (memory->is64()) {
         memory->addressType = Type::i32;
@@ -193,6 +296,11 @@ struct Memory64Lowering : public WalkerPass<PostWalker<Memory64Lowering>> {
         }
       }
     }
+    for (auto& table : module->tables) {
+      if (table->is64()) {
+        table->addressType = Type::i32;
+      }
+    }
     module->features.disable(FeatureSet::Memory64);
   }
 };
diff --git a/src/passes/Table64Lowering.cpp b/src/passes/Table64Lowering.cpp
deleted file mode 100644
index 1167515de..000000000
--- a/src/passes/Table64Lowering.cpp
+++ /dev/null
@@ -1,155 +0,0 @@
-/*
- * Copyright 2024 WebAssembly Community Group participants
- *
- * Licensed under the Apache License, Version 2.0 (the "License");
- * you may not use this file except in compliance with the License.
- * You may obtain a copy of the License at
- *
- *     http://www.apache.org/licenses/LICENSE-2.0
- *
- * Unless required by applicable law or agreed to in writing, software
- * distributed under the License is distributed on an "AS IS" BASIS,
- * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
- * See the License for the specific language governing permissions and
- * limitations under the License.
- */
-
-//
-// Lowers a module with a 64-bit table to one with a 32-bit table.
-//
-// This pass can be deleted once table64 is implemented in Wasm engines:
-// https://github.com/WebAssembly/memory64/issues/51
-//
-
-#include "ir/bits.h"
-#include "ir/import-utils.h"
-#include "pass.h"
-#include "wasm-builder.h"
-#include "wasm.h"
-
-namespace wasm {
-
-static Name TABLE_BASE("__table_base");
-static Name TABLE_BASE32("__table_base32");
-
-struct Table64Lowering : public WalkerPass<PostWalker<Table64Lowering>> {
-
-  void wrapAddress64(Expression*& ptr, Name tableName) {
-    if (ptr->type == Type::unreachable) {
-      return;
-    }
-    auto& module = *getModule();
-    auto* table = module.getTable(tableName);
-    if (table->is64()) {
-      assert(ptr->type == Type::i64);
-      ptr = Builder(module).makeUnary(UnaryOp::WrapInt64, ptr);
-    }
-  }
-
-  void extendAddress64(Expression*& ptr, Name tableName) {
-    if (ptr->type == Type::unreachable) {
-      return;
-    }
-    auto& module = *getModule();
-    auto* table = module.getTable(tableName);
-    if (table->is64()) {
-      assert(ptr->type == Type::i64);
-      ptr->type = Type::i32;
-      ptr = Builder(module).makeUnary(UnaryOp::ExtendUInt32, ptr);
-    }
-  }
-
-  void visitTableSize(TableSize* curr) {
-    auto& module = *getModule();
-    auto* table = module.getTable(curr->table);
-    if (table->is64()) {
-      auto* size = static_cast<Expression*>(curr);
-      extendAddress64(size, curr->table);
-      replaceCurrent(size);
-    }
-  }
-
-  void visitTableGrow(TableGrow* curr) {
-    auto& module = *getModule();
-    auto* table = module.getTable(curr->table);
-    if (table->is64()) {
-      wrapAddress64(curr->delta, curr->table);
-      auto* size = static_cast<Expression*>(curr);
-      extendAddress64(size, curr->table);
-      replaceCurrent(size);
-    }
-  }
-
-  void visitTableFill(TableFill* curr) {
-    wrapAddress64(curr->dest, curr->table);
-    wrapAddress64(curr->size, curr->table);
-  }
-
-  void visitTableCopy(TableCopy* curr) {
-    wrapAddress64(curr->dest, curr->destTable);
-    wrapAddress64(curr->source, curr->sourceTable);
-    wrapAddress64(curr->size, curr->destTable);
-  }
-
-  void visitTableInit(TableInit* curr) {
-    wrapAddress64(curr->dest, curr->table);
-  }
-
-  void visitCallIndirect(CallIndirect* curr) {
-    wrapAddress64(curr->target, curr->table);
-  }
-
-  void visitElementSegment(ElementSegment* segment) {
-    auto& module = *getModule();
-
-    // Passive segments don't have any offset to update.
-    if (segment->table.isNull() || !module.getTable(segment->table)->is64()) {
-      return;
-    }
-
-    if (auto* c = segment->offset->dynCast<Const>()) {
-      c->value = Literal(static_cast<uint32_t>(c->value.geti64()));
-      c->type = Type::i32;
-    } else if (auto* get = segment->offset->dynCast<GlobalGet>()) {
-      auto* g = module.getGlobal(get->name);
-      if (g->imported() && g->base == TABLE_BASE) {
-        ImportInfo info(module);
-        auto* memoryBase32 = info.getImportedGlobal(g->module, TABLE_BASE32);
-        if (!memoryBase32) {
-          Builder builder(module);
-          memoryBase32 = builder
-                           .makeGlobal(TABLE_BASE32,
-                                       Type::i32,
-                                       builder.makeConst(int32_t(0)),
-                                       Builder::Immutable)
-                           .release();
-          memoryBase32->module = g->module;
-          memoryBase32->base = TABLE_BASE32;
-          module.addGlobal(memoryBase32);
-        }
-        // Use this alternative import when initializing the segment.
-        assert(memoryBase32);
-        get->type = Type::i32;
-        get->name = memoryBase32->name;
-      }
-    } else {
-      WASM_UNREACHABLE("unexpected elem offset");
-    }
-  }
-
-  void run(Module* module) override {
-    Super::run(module);
-    // Don't modify the tables themselves until after the traversal since we
-    // that would require tables to be the last thing that get visited, and
-    // we don't want to depend on that specific ordering.
-    for (auto& table : module->tables) {
-      if (table->is64()) {
-        table->addressType = Type::i32;
-      }
-    }
-  }
-};
-
-Pass* createTable64LoweringPass() { return new Table64Lowering(); }
-
-} // namespace wasm
diff --git a/src/passes/pass.cpp b/src/passes/pass.cpp
index 7f6985e0f..ab95300d0 100644
--- a/src/passes/pass.cpp
+++ b/src/passes/pass.cpp
@@ -273,8 +273,8 @@ void PassRegistry::registerPasses() {
                "32-bit one",
                createMemory64LoweringPass);
   registerPass("table64-lowering",
-               "lower 64-bit tables 32-bit ones",
-               createTable64LoweringPass);
+               "alias for memory64-lowering",
+               createMemory64LoweringPass);
   registerPass("llvm-memory-copy-fill-lowering",
                "Lower memory.copy and memory.fill to wasm mvp and disable "
                "the bulk-memory feature.",