Preserve element segment encoding more often (#1807)

This commit fixes a small regression in reencoding from #1794 where
element segments subtly changed encoding by accident. This commit
additionally ensures that there's a text format for all element
encodings and updates printing to respect the same defaults. It should
now be possible to represent all formats for element segments in the
text format and have them round-trippable.

Author: alexcrichton

crates/wasm-encoder/src/reencode.rs

@@ -1408,14 +1408,13 @@ pub mod utils {
                 // example which wants to track uses to know when it's ok to
                 // remove a table.
                 //
-                // If the table index is still zero though go ahead and pass
-                // `None` here since that implicitly references table 0. Note
-                // that this means that this does not round-trip the encoding of
-                // `Some(0)` since that reencodes to `None`, but that's seen as
-                // hopefully ok.
-                match reencoder.table_index(table_index.unwrap_or(0)) {
-                    0 => None,
-                    i => Some(i),
+                // If the table index started at `None` and is still zero then
+                // preserve this encoding and keep it at `None`. Otherwise if
+                // the result is nonzero or it was previously nonzero then keep
+                // that encoding too.
+                match (table_index, reencoder.table_index(table_index.unwrap_or(0))) {
+                    (None, 0) => None,
+                    (_, n) => Some(n),
                 },
                 &reencoder.const_expr(offset_expr)?,
                 elems,

crates/wasmprinter/src/lib.rs

@@ -1531,8 +1531,7 @@ impl Printer<'_, '_> {
                     table_index,
                     offset_expr,
                 } => {
-                    let table_index = table_index.unwrap_or(0);
-                    if table_index != 0 {
+                    if let Some(table_index) = *table_index {
                         self.result.write_str(" ")?;
                         self.start_group("table ")?;
                         self.print_idx(&state.core.table_names, table_index)?;

crates/wast/src/core/binary.rs

@@ -703,7 +703,7 @@ impl Encode for Elem<'_> {
         match (&self.kind, &self.payload) {
             (
                 ElemKind::Active {
-                    table: Index::Num(0, _),
+                    table: None,
                     offset,
                 },
                 ElemPayload::Indices(_),
@@ -715,7 +715,13 @@ impl Encode for Elem<'_> {
                 e.push(0x01); // flags
                 e.push(0x00); // extern_kind
             }
-            (ElemKind::Active { table, offset }, ElemPayload::Indices(_)) => {
+            (
+                ElemKind::Active {
+                    table: Some(table),
+                    offset,
+                },
+                ElemPayload::Indices(_),
+            ) => {
                 e.push(0x02); // flags
                 table.encode(e);
                 offset.encode(e, None);
@@ -727,7 +733,7 @@ impl Encode for Elem<'_> {
             }
             (
                 ElemKind::Active {
-                    table: Index::Num(0, _),
+                    table: None,
                     offset,
                 },
                 ElemPayload::Exprs {
@@ -752,7 +758,7 @@ impl Encode for Elem<'_> {
             }
             (ElemKind::Active { table, offset }, ElemPayload::Exprs { ty, .. }) => {
                 e.push(0x06);
-                table.encode(e);
+                table.map(|t| t.unwrap_u32()).unwrap_or(0).encode(e);
                 offset.encode(e, None);
                 ty.encode(e);
             }

crates/wast/src/core/resolve/deinline_import_export.rs

@@ -139,7 +139,7 @@ pub fn run(fields: &mut Vec<ModuleField>) {
                             id: None,
                             name: None,
                             kind: ElemKind::Active {
-                                table: Index::Id(id),
+                                table: Some(Index::Id(id)),
                                 offset: Expression::one(if is64 {
                                     Instruction::I64Const(0)
                                 } else {

crates/wast/src/core/resolve/names.rs

@@ -184,7 +184,9 @@ impl<'a> Resolver<'a> {
             ModuleField::Elem(e) => {
                 match &mut e.kind {
                     ElemKind::Active { table, offset } => {
-                        self.resolve(table, Ns::Table)?;
+                        if let Some(table) = table {
+                            self.resolve(table, Ns::Table)?;
+                        }
                         self.resolve_expr(offset)?;
                     }
                     ElemKind::Passive { .. } | ElemKind::Declared { .. } => {}

crates/wast/src/core/table.rs

@@ -149,7 +149,7 @@ pub enum ElemKind<'a> {
     /// An active segment associated with a table.
     Active {
         /// The table this `elem` is initializing.
-        table: Index<'a>,
+        table: Option<Index<'a>>,
         /// The offset within `table` that we'll initialize at.
         offset: Expression<'a>,
     },
@@ -197,15 +197,15 @@ impl<'a> Parse<'a> for Elem<'a> {
                 // time, this probably should get removed when the threads
                 // proposal is rebased on the current spec.
                 table_omitted = true;
-                Index::Num(parser.parse()?, span)
+                Some(Index::Num(parser.parse()?, span))
             } else if parser.peek2::<kw::table>()? {
-                parser.parens(|p| {
+                Some(parser.parens(|p| {
                     p.parse::<kw::table>()?;
                     p.parse()
-                })?
+                })?)
             } else {
                 table_omitted = true;
-                Index::Num(0, span)
+                None
             };
 
             let offset = parse_expr_or_single_instr::<kw::offset>(parser)?;

fuzz/src/lib.rs

@@ -1,5 +1,6 @@
 use libfuzzer_sys::arbitrary::{Result, Unstructured};
 use std::fmt::Debug;
+use std::sync::atomic::{AtomicUsize, Ordering::SeqCst};
 use wasm_smith::{Component, Config, Module};
 use wasmparser::WasmFeatures;
 
@@ -120,15 +121,25 @@ pub fn validator_for_config(config: &Config) -> wasmparser::Validator {
 pub fn log_wasm(wasm: &[u8], config: impl Debug) {
     drop(env_logger::try_init());
 
-    if log::log_enabled!(log::Level::Debug) {
-        log::debug!("writing test case to `test.wasm` ...");
-        std::fs::write("test.wasm", wasm).unwrap();
-        std::fs::write("test.config", format!("{:#?}", config)).unwrap();
-        if let Ok(wat) = wasmprinter::print_bytes(wasm) {
-            log::debug!("writing text format to `test.wat` ...");
-            std::fs::write("test.wat", wat).unwrap();
-        } else {
-            drop(std::fs::remove_file("test.wat"));
-        }
+    if !log::log_enabled!(log::Level::Debug) {
+        return;
+    }
+
+    static CNT: AtomicUsize = AtomicUsize::new(0);
+
+    let i = CNT.fetch_add(1, SeqCst);
+
+    let wasm_file = format!("test{i}.wasm");
+    let config_file = format!("test{i}.config");
+    let wat_file = format!("test{i}.wat");
+
+    log::debug!("writing test case to `{wasm_file}` ...");
+    std::fs::write(&wasm_file, wasm).unwrap();
+    std::fs::write(&config_file, format!("{:#?}", config)).unwrap();
+    if let Ok(wat) = wasmprinter::print_bytes(wasm) {
+        log::debug!("writing text format to `{wat_file}` ...");
+        std::fs::write(&wat_file, wat).unwrap();
+    } else {
+        drop(std::fs::remove_file(&wat_file));
     }
 }

fuzz/src/reencode.rs

@@ -2,14 +2,19 @@ use arbitrary::{Result, Unstructured};
 use wasm_encoder::reencode::{Reencode, RoundtripReencoder};
 
 pub fn run(u: &mut Unstructured<'_>) -> Result<()> {
-    let (module1, _) = super::generate_valid_module(u, |_, _| Ok(()))?;
+    let (module1, config) = super::generate_valid_module(u, |_, _| Ok(()))?;
 
     let mut module2 = Default::default();
     RoundtripReencoder
         .parse_core_module(&mut module2, wasmparser::Parser::new(0), &module1)
         .unwrap();
 
     let module2 = module2.finish();
+    if module1 == module2 {
+        return Ok(());
+    }
+    crate::log_wasm(&module1, &config);
+    crate::log_wasm(&module2, &config);
     assert_eq!(module1, module2);
 
     Ok(())

tests/cli/dump-elem-segments.wat

@@ -0,0 +1,18 @@
+;; RUN: print % | dump
+
+;; test that all forms of element segments can be round-tripped from
+;; text-to-binary.
+(module
+  (table 1 1 funcref)
+
+  (func $f)
+
+  (elem (i32.const 0) func)                   ;; 0x00
+  (elem func)                                 ;; 0x01
+  (elem (table 0) (i32.const 0) func)         ;; 0x02
+  (elem declare func)                         ;; 0x03
+  (elem (i32.const 0) funcref)                ;; 0x04
+  (elem funcref)                              ;; 0x05
+  (elem (table 0) (i32.const 0) funcref)      ;; 0x06
+  (elem declare funcref)                      ;; 0x07
+)

tests/cli/dump-elem-segments.wat.stdout

@@ -0,0 +1,46 @@
+  0x0 | 00 61 73 6d | version 1 (Module)
+      | 01 00 00 00
+  0x8 | 01 04       | type section
+  0xa | 01          | 1 count
+--- rec group 0 (implicit) ---
+  0xb | 60 00 00    | [type 0] SubType { is_final: true, supertype_idx: None, composite_type: CompositeType { inner: Func(FuncType { params: [], results: [] }), shared: false } }
+  0xe | 03 02       | func section
+ 0x10 | 01          | 1 count
+ 0x11 | 00          | [func 0] type 0
+ 0x12 | 04 05       | table section
+ 0x14 | 01          | 1 count
+ 0x15 | 70 01 01 01 | [table 0] Table { ty: TableType { element_type: funcref, table64: false, initial: 1, maximum: Some(1), shared: false }, init: RefNull }
+ 0x19 | 09 25       | element section
+ 0x1b | 08          | 8 count
+ 0x1c | 00          | element table[None]
+ 0x1d | 41 00       | i32_const value:0
+ 0x1f | 0b          | end
+ 0x20 | 00          | 0 items [indices]
+ 0x21 | 01 00 00    | element passive, 0 items [indices]
+ 0x24 | 02 00       | element table[Some(0)]
+ 0x26 | 41 00       | i32_const value:0
+ 0x28 | 0b          | end
+ 0x29 | 00 00       | 0 items [indices]
+ 0x2b | 03 00 00    | element declared 0 items [indices]
+ 0x2e | 04          | element table[None]
+ 0x2f | 41 00       | i32_const value:0
+ 0x31 | 0b          | end
+ 0x32 | 00          | 0 items [exprs funcref]
+ 0x33 | 05 70 00    | element passive, 0 items [exprs funcref]
+ 0x36 | 06 00       | element table[Some(0)]
+ 0x38 | 41 00       | i32_const value:0
+ 0x3a | 0b          | end
+ 0x3b | 70 00       | 0 items [exprs funcref]
+ 0x3d | 07 70 00    | element declared 0 items [exprs funcref]
+ 0x40 | 0a 04       | code section
+ 0x42 | 01          | 1 count
+============== func 0 ====================
+ 0x43 | 02          | size of function
+ 0x44 | 00          | 0 local blocks
+ 0x45 | 0b          | end
+ 0x46 | 00 0b       | custom section
+ 0x48 | 04 6e 61 6d | name: "name"
+      | 65         
+ 0x4d | 01 04       | function name section
+ 0x4f | 01          | 1 count
+ 0x50 | 00 01 66    | Naming { index: 0, name: "f" }