[Table64Lowering] Don't assume that all segments are from 64-bit tabl…

…es (#6599)

This allows modules to contains both 32-bit and 64-bit segment.

In order to check the table/memory state when visiting segments we need
to ensure that memories/tables are visited only after their segments.

The comments in visitTable/visitMemory already assumed this but it
wasn't true in practice.

src/passes/Memory64Lowering.cpp

@@ -118,6 +118,7 @@ struct Memory64Lowering : public WalkerPass<PostWalker<Memory64Lowering>> {
 
   void visitMemory(Memory* memory) {
     // This is visited last.
+    seenMemory = true;
     if (memory->is64()) {
       memory->indexType = Type::i32;
       if (memory->hasMax() && memory->max > Memory::kMaxSize32) {
@@ -127,16 +128,19 @@ struct Memory64Lowering : public WalkerPass<PostWalker<Memory64Lowering>> {
   }
 
   void visitDataSegment(DataSegment* segment) {
-    if (segment->isPassive) {
-      // passive segments don't have any offset to adjust
+    // We assume that memories are visited after segments, so assert that here.
+    assert(!seenMemory);
+    auto& module = *getModule();
+
+    // passive segments don't have any offset to adjust
+    if (segment->isPassive || !module.getMemory(segment->memory)->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& module = *getModule();
       auto* g = module.getGlobal(get->name);
       if (g->imported() && g->base == MEMORY_BASE) {
         ImportInfo info(module);
@@ -170,6 +174,8 @@ struct Memory64Lowering : public WalkerPass<PostWalker<Memory64Lowering>> {
     super::run(module);
     module->features.disable(FeatureSet::Memory64);
   }
+
+  bool seenMemory = false;
 };
 
 Pass* createMemory64LoweringPass() { return new Memory64Lowering(); }

src/passes/Table64Lowering.cpp

@@ -97,22 +97,26 @@ struct Table64Lowering : public WalkerPass<PostWalker<Table64Lowering>> {
 
   void visitTable(Table* table) {
     // This is visited last.
+    seenTable = true;
     if (table->is64()) {
       table->indexType = Type::i32;
     }
   }
 
   void visitElementSegment(ElementSegment* segment) {
-    if (segment->table.isNull()) {
-      // Passive segments don't have any offset to update.
+    // We assume that tables are visited after segments, so assert that here.
+    assert(!seenTable);
+    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& module = *getModule();
       auto* g = module.getGlobal(get->name);
       if (g->imported() && g->base == TABLE_BASE) {
         ImportInfo info(module);
@@ -138,6 +142,8 @@ struct Table64Lowering : public WalkerPass<PostWalker<Table64Lowering>> {
       WASM_UNREACHABLE("unexpected elem offset");
     }
   }
+
+  bool seenTable = false;
 };
 
 Pass* createTable64LoweringPass() { return new Table64Lowering(); }

src/wasm-traversal.h

@@ -256,18 +256,18 @@ struct Walker : public VisitorType {
         self->walkTag(curr.get());
       }
     }
-    for (auto& curr : module->tables) {
-      self->walkTable(curr.get());
-    }
     for (auto& curr : module->elementSegments) {
       self->walkElementSegment(curr.get());
     }
-    for (auto& curr : module->memories) {
-      self->walkMemory(curr.get());
+    for (auto& curr : module->tables) {
+      self->walkTable(curr.get());
     }
     for (auto& curr : module->dataSegments) {
       self->walkDataSegment(curr.get());
     }
+    for (auto& curr : module->memories) {
+      self->walkMemory(curr.get());
+    }
   }
 
   // Walks module-level code, that is, code that is not in functions.

test/lit/passes/table64-lowering-features.wast → test/lit/passes/table64-lowering.wast

@@ -7,14 +7,20 @@
 
   ;; CHECK:      (type $1 (func (result i64)))
 
-  ;; CHECK:      (table $t 10 100 funcref)
-  (table $t i64 10 100 funcref)
+  ;; CHECK:      (table $t64 10 100 funcref)
+  (table $t64 i64 10 100 funcref)
 
-  ;; CHECK:      (elem $elem (table $t) (i32.const 0) funcref (ref.null nofunc))
-  (elem $elem (table $t) (i64.const 0) funcref (ref.null func))
+  ;; CHECK:      (table $t32 10 100 funcref)
+
+  ;; CHECK:      (elem $elem64 (table $t64) (i32.const 0) funcref (ref.null nofunc))
+  (elem $elem64 (table $t64) (i64.const 0) funcref (ref.null func))
+
+  (table $t32 10 100 funcref)
+  ;; CHECK:      (elem $elem32 (table $t32) (i32.const 0) funcref (ref.null nofunc))
+  (elem $elem32 (table $t32) (i32.const 0) funcref (ref.null func))
 
   ;; CHECK:      (func $test_call_indirect
-  ;; CHECK-NEXT:  (call_indirect $t (type $0)
+  ;; CHECK-NEXT:  (call_indirect $t64 (type $0)
   ;; CHECK-NEXT:   (i32.wrap_i64
   ;; CHECK-NEXT:    (i64.const 0)
   ;; CHECK-NEXT:   )
@@ -26,16 +32,16 @@
 
   ;; CHECK:      (func $test_table_size (result i64)
   ;; CHECK-NEXT:  (i64.extend_i32_u
-  ;; CHECK-NEXT:   (table.size $t)
+  ;; CHECK-NEXT:   (table.size $t64)
   ;; CHECK-NEXT:  )
   ;; CHECK-NEXT: )
   (func $test_table_size (result i64)
-    (table.size $t)
+    (table.size $t64)
   )
 
   ;; CHECK:      (func $test_table_grow (result i64)
   ;; CHECK-NEXT:  (i64.extend_i32_u
-  ;; CHECK-NEXT:   (table.grow $t
+  ;; CHECK-NEXT:   (table.grow $t64
   ;; CHECK-NEXT:    (ref.null nofunc)
   ;; CHECK-NEXT:    (i32.wrap_i64
   ;; CHECK-NEXT:     (i64.const 10)
@@ -44,11 +50,11 @@
   ;; CHECK-NEXT:  )
   ;; CHECK-NEXT: )
   (func $test_table_grow (result i64)
-    (table.grow $t (ref.null func) (i64.const 10))
+    (table.grow $t64 (ref.null func) (i64.const 10))
   )
 
   ;; CHECK:      (func $test_table_fill
-  ;; CHECK-NEXT:  (table.fill $t
+  ;; CHECK-NEXT:  (table.fill $t64
   ;; CHECK-NEXT:   (i32.wrap_i64
   ;; CHECK-NEXT:    (i64.const 0)
   ;; CHECK-NEXT:   )
@@ -59,6 +65,6 @@
   ;; CHECK-NEXT:  )
   ;; CHECK-NEXT: )
   (func $test_table_fill
-    (table.fill $t (i64.const 0) (ref.null func) (i64.const 10))
+    (table.fill $t64 (i64.const 0) (ref.null func) (i64.const 10))
   )
 )