Fix roundtripping specialized element segments of table zero (#4212)

Before this fix, the first table (index 0) is counted as its element segment
having "no table index" even when its type is not funcref, which could break
things if that table had a more specialized type.

2 files changed, 48 insertions, 1 deletions

diff --git a/src/wasm/wasm-binary.cpp b/src/wasm/wasm-binary.cpp
index a8eea2808..df7d8c114 100644
--- a/src/wasm/wasm-binary.cpp
+++ b/src/wasm/wasm-binary.cpp
@@ -587,10 +587,15 @@ void WasmBinaryWriter::writeElementSegments() {
     // If the segment is MVP, we can use the shorter form.
     bool usesExpressions = TableUtils::usesExpressions(segment.get(), wasm);
 
+    // The table index can and should be elided for active segments of table 0
+    // when table 0 has type funcref. This was the only type of segment
+    // supported by the MVP, which also did not support table indices in the
+    // segment encoding.
     bool hasTableIndex = false;
     if (!isPassive) {
       tableIdx = getTableIndex(segment->table);
-      hasTableIndex = tableIdx > 0;
+      hasTableIndex =
+        tableIdx > 0 || wasm->getTable(segment->table)->type != Type::funcref;
     }
 
     uint32_t flags = 0;
diff --git a/test/lit/table-first-special.wast b/test/lit/table-first-special.wast
new file mode 100644
index 000000000..0ddc730b9
--- /dev/null
+++ b/test/lit/table-first-special.wast
@@ -0,0 +1,42 @@
+;; NOTE: Assertions have been generated by update_lit_checks.py --all-items and should not be edited.
+
+;; RUN: wasm-as %s -all -g -o %t.wasm
+;; RUN: wasm-dis %t.wasm -all -o %t.wast
+;; RUN: wasm-as %s -all -o %t.nodebug.wasm
+;; RUN: wasm-dis %t.nodebug.wasm -all -o %t.nodebug.wast
+;; RUN: wasm-opt %t.wast -all -o %t.text.wast -g -S
+;; RUN: cat %t.wast | filecheck %s --check-prefix=CHECK-BINARY
+;; RUN: cat %t.nodebug.wast | filecheck %s --check-prefix=CHECK-NODEBUG
+;; RUN: cat %t.text.wast | filecheck %s --check-prefix=CHECK-TEXT
+
+;; The very first table has a specialized type, as does its element segment.
+;; Verify that we can roundtrip that.
+(module
+ ;; CHECK-BINARY:      (type $vii (func (param i32 i32)))
+ ;; CHECK-TEXT:      (type $vii (func (param i32 i32)))
+ (type $vii (func (param i32 i32)))
+ ;; CHECK-BINARY:      (table $table-1 10 10 (ref null $vii))
+ ;; CHECK-TEXT:      (table $table-1 10 10 (ref null $vii))
+ (table $table-1 10 10 (ref null $vii))
+ ;; CHECK-BINARY:      (elem $elem-1 (table $table-1) (i32.const 0) (ref null $vii))
+ ;; CHECK-TEXT:      (elem $elem-1 (table $table-1) (i32.const 0) (ref null $vii))
+ (elem $elem-1 (table $table-1) (i32.const 0) (ref null $vii))
+ ;; CHECK-BINARY:      (func $foo (param $0 i32) (param $1 i32)
+ ;; CHECK-BINARY-NEXT:  (nop)
+ ;; CHECK-BINARY-NEXT: )
+ ;; CHECK-TEXT:      (func $foo (param $0 i32) (param $1 i32)
+ ;; CHECK-TEXT-NEXT:  (nop)
+ ;; CHECK-TEXT-NEXT: )
+ (func $foo (param $0 i32) (param $1 i32)
+  (nop)
+ )
+)
+;; CHECK-NODEBUG:      (type $i32_i32_=>_none (func (param i32 i32)))
+
+;; CHECK-NODEBUG:      (table $0 10 10 (ref null $i32_i32_=>_none))
+
+;; CHECK-NODEBUG:      (elem (table $0) (i32.const 0) (ref null $i32_i32_=>_none))
+
+;; CHECK-NODEBUG:      (func $0 (param $0 i32) (param $1 i32)
+;; CHECK-NODEBUG-NEXT:  (nop)
+;; CHECK-NODEBUG-NEXT: )