Update the walrus dependency (#2125)

This commit updates the `walrus` crate used in `wasm-bindgen`. The major
change here is how `walrus` handles element segments, exposing segments
rather than trying to keep a contiugous array of all the elements and
doing the splitting itself. That means that we need to do mroe logic
here in `wasm-bindgen` to juggle indices, segments, etc.

crates/anyref-xform/Cargo.toml

@@ -13,7 +13,7 @@ edition = '2018'
 
 [dependencies]
 anyhow = "1.0"
-walrus = "0.14.0"
+walrus = "0.16.0"
 
 [dev-dependencies]
 rayon = "1.0"

crates/anyref-xform/src/lib.rs

@@ -18,8 +18,9 @@
 use anyhow::{anyhow, bail, Error};
 use std::cmp;
 use std::collections::{BTreeMap, HashMap, HashSet};
+use std::mem;
 use walrus::ir::*;
-use walrus::{ExportId, ImportId, InstrLocId, TypeId};
+use walrus::{ElementId, ExportId, ImportId, InstrLocId, TypeId};
 use walrus::{FunctionId, GlobalId, InitExpr, Module, TableId, ValType};
 
 // must be kept in sync with src/lib.rs and ANYREF_HEAP_START
@@ -34,11 +35,19 @@ pub struct Context {
     // values in the function signature should turn into anyref.
     imports: HashMap<ImportId, Function>,
     exports: HashMap<ExportId, Function>,
-    elements: BTreeMap<u32, (u32, Function)>,
+
+    // List of functions we're transforming that are present in the function
+    // table. Each index here is an index into the function table, and the
+    // `Function` describes how we're transforming it.
+    new_elements: Vec<(u32, Function)>,
 
     // When wrapping closures with new shims, this is the index of the next
     // table entry that we'll be handing out.
-    next_element: u32,
+    new_element_offset: u32,
+
+    // Map of the existing function table, keyed by offset and contains the
+    // final offset plus the element segment used to initialized that range.
+    elements: BTreeMap<u32, ElementId>,
 
     // The anyref table we'll be using, injected after construction
     table: Option<TableId>,
@@ -93,22 +102,27 @@ impl Context {
         // Figure out what the maximum index of functions pointers are. We'll
         // be adding new entries to the function table later (maybe) so
         // precalculate this ahead of time.
-        let mut tables = module.tables.iter().filter_map(|t| match &t.kind {
-            walrus::TableKind::Function(f) => Some(f),
-            _ => None,
-        });
-        if let Some(t) = tables.next() {
-            if tables.next().is_some() {
-                bail!("more than one function table present")
+        if let Some(t) = module.tables.main_function_table()? {
+            let t = module.tables.get(t);
+            for id in t.elem_segments.iter() {
+                let elem = module.elements.get(*id);
+                let offset = match &elem.kind {
+                    walrus::ElementKind::Active { offset, .. } => offset,
+                    _ => continue,
+                };
+                let offset = match offset {
+                    walrus::InitExpr::Value(Value::I32(n)) => *n as u32,
+                    other => bail!("invalid offset for segment of function table {:?}", other),
+                };
+                let max = offset + elem.members.len() as u32;
+                self.new_element_offset = cmp::max(self.new_element_offset, max);
+                self.elements.insert(offset, *id);
             }
-            self.next_element = t.elements.len() as u32;
         }
-        drop(tables);
 
         // Add in an anyref table to the module, which we'll be using for
         // our transform below.
-        let kind = walrus::TableKind::Anyref(Default::default());
-        self.table = Some(module.tables.add_local(DEFAULT_MIN, None, kind));
+        self.table = Some(module.tables.add_local(DEFAULT_MIN, None, ValType::Anyref));
 
         Ok(())
     }
@@ -151,10 +165,8 @@ impl Context {
         ret_anyref: bool,
     ) -> Option<u32> {
         self.function(anyref, ret_anyref).map(|f| {
-            let ret = self.next_element;
-            self.next_element += 1;
-            self.elements.insert(ret, (idx, f));
-            ret
+            self.new_elements.push((idx, f));
+            self.new_elements.len() as u32 + self.new_element_offset - 1
         })
     }
 
@@ -265,7 +277,7 @@ impl Transform<'_> {
         self.process_exports(module)?;
         assert!(self.cx.exports.is_empty());
         self.process_elements(module)?;
-        assert!(self.cx.elements.is_empty());
+        assert!(self.cx.new_elements.is_empty());
 
         // If we didn't actually transform anything, no need to inject or
         // rewrite anything from below.
@@ -394,18 +406,20 @@ impl Transform<'_> {
             None => return Ok(()),
         };
         let table = module.tables.get_mut(table);
-        let kind = match &mut table.kind {
-            walrus::TableKind::Function(f) => f,
-            _ => unreachable!(),
-        };
-        if kind.relative_elements.len() > 0 {
-            bail!("not compatible with relative element initializers yet");
-        }
 
         // Create shims for all our functions and append them all to the segment
         // which places elements at the end.
-        while let Some((idx, function)) = self.cx.elements.remove(&(kind.elements.len() as u32)) {
-            let target = kind.elements[idx as usize].unwrap();
+        let mut new_segment = Vec::new();
+        for (idx, function) in mem::replace(&mut self.cx.new_elements, Vec::new()) {
+            let (&offset, &orig_element) = self
+                .cx
+                .elements
+                .range(..=idx)
+                .next_back()
+                .ok_or(anyhow!("failed to find segment defining index {}", idx))?;
+            let target = module.elements.get(orig_element).members[(idx - offset) as usize].ok_or(
+                anyhow!("function index {} not present in element segment", idx),
+            )?;
             let (shim, _anyref_ty) = self.append_shim(
                 target,
                 &format!("closure{}", idx),
@@ -414,14 +428,21 @@ impl Transform<'_> {
                 &mut module.funcs,
                 &mut module.locals,
             )?;
-            kind.elements.push(Some(shim));
+            new_segment.push(Some(shim));
         }
 
         // ... and next update the limits of the table in case any are listed.
-        table.initial = cmp::max(table.initial, kind.elements.len() as u32);
+        let new_max = self.cx.new_element_offset + new_segment.len() as u32;
+        table.initial = cmp::max(table.initial, new_max);
         if let Some(max) = table.maximum {
-            table.maximum = Some(cmp::max(max, kind.elements.len() as u32));
+            table.maximum = Some(cmp::max(max, new_max));
         }
+        let kind = walrus::ElementKind::Active {
+            table: table.id(),
+            offset: InitExpr::Value(Value::I32(self.cx.new_element_offset as i32)),
+        };
+        let segment = module.elements.add(kind, ValType::Funcref, new_segment);
+        table.elem_segments.insert(segment);
 
         Ok(())
     }

crates/anyref-xform/tests/table.wat

@@ -28,5 +28,6 @@
   (table (;0;) 2 funcref)
   (table (;1;) 32 anyref)
   (export "func" (table 0))
-  (elem (;0;) (i32.const 0) func $foo $closure0 anyref shim))
+  (elem (;0;) (i32.const 0) func $foo)
+  (elem (;1;) (i32.const 1) func $closure0 anyref shim))
 ;)

crates/cli-support/Cargo.toml

@@ -18,7 +18,7 @@ log = "0.4"
 rustc-demangle = "0.1.13"
 serde_json = "1.0"
 tempfile = "3.0"
-walrus = "0.14.0"
+walrus = "0.16.1"
 wasm-bindgen-anyref-xform = { path = '../anyref-xform', version = '=0.2.62' }
 wasm-bindgen-multi-value-xform = { path = '../multi-value-xform', version = '=0.2.62' }
 wasm-bindgen-shared = { path = "../shared", version = '=0.2.62' }
@@ -26,5 +26,5 @@ wasm-bindgen-threads-xform = { path = '../threads-xform', version = '=0.2.62' }
 wasm-bindgen-wasm-conventions = { path = '../wasm-conventions', version = '=0.2.62' }
 wasm-bindgen-wasm-interpreter = { path = "../wasm-interpreter", version = '=0.2.62' }
 wit-text = "0.1.1"
-wit-walrus = "0.1.0"
+wit-walrus = "0.2.0"
 wit-validator = "0.1.0"

crates/cli-support/src/anyref.rs

@@ -3,12 +3,12 @@ use crate::intrinsic::Intrinsic;
 use crate::wit::AuxImport;
 use crate::wit::{AdapterKind, Instruction, NonstandardWitSection};
 use crate::wit::{AdapterType, InstructionData, StackChange, WasmBindgenAux};
-use anyhow::Error;
+use anyhow::Result;
 use std::collections::HashMap;
-use walrus::Module;
+use walrus::{ir::Value, ElementKind, InitExpr, Module};
 use wasm_bindgen_anyref_xform::Context;
 
-pub fn process(module: &mut Module) -> Result<(), Error> {
+pub fn process(module: &mut Module) -> Result<()> {
     let mut cfg = Context::default();
     cfg.prepare(module)?;
     let section = module
@@ -382,3 +382,86 @@ fn module_needs_anyref_metadata(aux: &WasmBindgenAux, section: &NonstandardWitSe
         })
     })
 }
+
+/// In MVP wasm all element segments must be contiguous lists of function
+/// indices. Post-MVP with reference types element segments can have holes.
+/// While `walrus` will select the encoding that fits, this function forces the
+/// listing of segments to be MVP-compatible.
+pub fn force_contiguous_elements(module: &mut Module) -> Result<()> {
+    // List of new element segments we're going to be adding.
+    let mut new_segments = Vec::new();
+
+    // Here we take a look at all element segments in the module to see if we
+    // need to split them.
+    for segment in module.elements.iter_mut() {
+        // If this segment has all-`Some` members then it's alrady contiguous
+        // and we can skip it.
+        if segment.members.iter().all(|m| m.is_some()) {
+            continue;
+        }
+
+        // For now active segments are all we're interested in since
+        // passive/declared have no hope of being MVP-compatible anyway.
+        // Additionally we only handle active segments with i32 offsets, since
+        // global offsets get funky since we'd need to add an offset.
+        let (table, offset) = match &segment.kind {
+            ElementKind::Active {
+                table,
+                offset: InitExpr::Value(Value::I32(n)),
+            } => (*table, *n),
+            _ => continue,
+        };
+
+        // `block` keeps track of a block of contiguous segment of functions
+        let mut block = None;
+        // This keeps track of where we're going to truncate the current segment
+        // after we split out all the blocks.
+        let mut truncate = 0;
+        // This commits a block of contiguous functions into the `new_segments`
+        // list, accounting for the new offset which is relative to the old
+        // offset.
+        let mut commit = |last_idx: usize, block: Vec<_>| {
+            let new_offset = offset + (last_idx - block.len()) as i32;
+            let new_offset = InitExpr::Value(Value::I32(new_offset));
+            new_segments.push((table, new_offset, segment.ty, block));
+        };
+        for (i, id) in segment.members.iter().enumerate() {
+            match id {
+                // If we find a function, then we either start a new block or
+                // push it onto the existing block.
+                Some(id) => block.get_or_insert(Vec::new()).push(Some(*id)),
+                None => {
+                    let block = match block.take() {
+                        Some(b) => b,
+                        None => continue,
+                    };
+                    // If this is the first block (truncate isn't set and the
+                    // length of the block means it starts from the beginning),
+                    // then we leave it in the original list and don't commit
+                    // anything, we'll just edit where to truncate later.
+                    // Otherwise we commit this block to the new segment list.
+                    if truncate == 0 && block.len() == i {
+                        truncate = i;
+                    } else {
+                        commit(i, block);
+                    }
+                }
+            }
+        }
+
+        // If there's no trailing empty slots then we commit the last block onto
+        // the new segment list.
+        if let Some(block) = block {
+            commit(segment.members.len(), block);
+        }
+        segment.members.truncate(truncate);
+    }
+
+    for (table, offset, ty, members) in new_segments {
+        let id = module
+            .elements
+            .add(ElementKind::Active { table, offset }, ty, members);
+        module.tables.get_mut(table).elem_segments.insert(id);
+    }
+    Ok(())
+}

crates/cli-support/src/descriptors.rs

@@ -107,19 +107,9 @@ impl WasmBindgenDescriptorsSection {
         // For all indirect functions that were closure descriptors, delete them
         // from the function table since we've executed them and they're not
         // necessary in the final binary.
-        let table_id = match interpreter.function_table_id() {
-            Some(id) => id,
-            None => return Ok(()),
-        };
-        let table = module.tables.get_mut(table_id);
-        let table = match &mut table.kind {
-            walrus::TableKind::Function(f) => f,
-            _ => unreachable!(),
-        };
-        for idx in element_removal_list {
+        for (segment, idx) in element_removal_list {
             log::trace!("delete element {}", idx);
-            assert!(table.elements[idx].is_some());
-            table.elements[idx] = None;
+            module.elements.get_mut(segment).members[idx] = None;
         }
 
         // And finally replace all calls of `wbindgen_describe_closure` with a

crates/cli-support/src/js/binding.rs

@@ -1145,18 +1145,9 @@ impl Invocation {
             // The function table never changes right now, so we can statically
             // look up the desired function.
             CallTableElement(idx) => {
-                let table = module
-                    .tables
-                    .main_function_table()?
-                    .ok_or_else(|| anyhow!("no function table found"))?;
-                let functions = match &module.tables.get(table).kind {
-                    walrus::TableKind::Function(f) => f,
-                    _ => bail!("should have found a function table"),
-                };
-                let id = functions
-                    .elements
-                    .get(*idx as usize)
-                    .and_then(|id| *id)
+                let entry = wasm_bindgen_wasm_conventions::get_function_table_entry(module, *idx)?;
+                let id = entry
+                    .func
                     .ok_or_else(|| anyhow!("function table wasn't filled in a {}", idx))?;
                 Invocation::Core { id, defer: false }
             }

crates/cli-support/src/lib.rs

@@ -373,6 +373,11 @@ impl Bindgen {
             for id in ids {
                 module.exports.delete(id);
             }
+            // Clean up element segments as well if they have holes in them
+            // after some of our transformations, because non-anyref engines
+            // only support contiguous arrays of function references in element
+            // segments.
+            anyref::force_contiguous_elements(&mut module)?;
         }
 
         // If wasm interface types are enabled then the `__wbindgen_throw`
@@ -564,13 +569,6 @@ impl OutputMode {
         }
     }
 
-    fn bundler(&self) -> bool {
-        match self {
-            OutputMode::Bundler { .. } => true,
-            _ => false,
-        }
-    }
-
     fn esm_integration(&self) -> bool {
         match self {
             OutputMode::Bundler { .. }

crates/cli-support/src/wit/section.rs

@@ -220,18 +220,11 @@ fn translate_instruction(
             _ => bail!("can only call exported functions"),
         },
         CallTableElement(e) => {
-            let table = module
-                .tables
-                .main_function_table()?
-                .ok_or_else(|| anyhow!("no function table found in module"))?;
-            let functions = match &module.tables.get(table).kind {
-                walrus::TableKind::Function(f) => f,
-                _ => unreachable!(),
-            };
-            match functions.elements.get(*e as usize) {
-                Some(Some(f)) => Ok(wit_walrus::Instruction::CallCore(*f)),
-                _ => bail!("expected to find an element of the function table"),
-            }
+            let entry = wasm_bindgen_wasm_conventions::get_function_table_entry(module, *e)?;
+            let id = entry
+                .func
+                .ok_or_else(|| anyhow!("function table wasn't filled in a {}", e))?;
+            Ok(wit_walrus::Instruction::CallCore(id))
         }
         StringToMemory {
             mem,

crates/cli-support/src/wit/standard.rs

@@ -313,7 +313,9 @@ impl AdapterType {
             walrus::ValType::F32 => AdapterType::F32,
             walrus::ValType::F64 => AdapterType::F64,
             walrus::ValType::Anyref => AdapterType::Anyref,
-            walrus::ValType::V128 => return None,
+            walrus::ValType::Funcref | walrus::ValType::Nullref | walrus::ValType::V128 => {
+                return None
+            }
         })
     }
 

crates/cli/Cargo.toml

@@ -24,7 +24,7 @@ rouille = { version = "3.0.0", default-features = false }
 serde = { version = "1.0", features = ['derive'] }
 serde_derive = "1.0"
 serde_json = "1.0"
-walrus = { version = "0.14.0", features = ['parallel'] }
+walrus = { version = "0.16.0", features = ['parallel'] }
 wasm-bindgen-cli-support = { path = "../cli-support", version = "=0.2.62" }
 wasm-bindgen-shared = { path = "../shared", version = "=0.2.62" }
 
@@ -34,11 +34,11 @@ diff = "0.1"
 predicates = "1.0.0"
 rayon = "1.0"
 tempfile = "3.0"
-walrus = "0.14"
+walrus = "0.16"
 wit-printer = "0.1"
 wit-text = "0.1"
 wit-validator = "0.1"
-wit-walrus = "0.1"
+wit-walrus = "0.2"
 
 [[test]]
 name = "reference"

crates/cli/tests/reference.rs

@@ -152,6 +152,13 @@ fn sanitize_wasm(wasm: &Path) -> Result<String> {
     for mem in module.memories.iter_mut() {
         mem.data_segments.drain();
     }
+    let ids = module.elements.iter().map(|d| d.id()).collect::<Vec<_>>();
+    for id in ids {
+        module.elements.delete(id);
+    }
+    for table in module.tables.iter_mut() {
+        table.elem_segments.drain();
+    }
     let ids = module
         .exports
         .iter()

crates/multi-value-xform/Cargo.toml

@@ -13,12 +13,12 @@ edition = "2018"
 
 [dependencies]
 anyhow = "1.0"
-walrus = "0.14.0"
+walrus = "0.16.0"
 
 [dev-dependencies]
 rayon = "1.0"
 wasmprinter = "0.2"
-wast = "3.0"
+wast = "15.0"
 wat = "1.0"
 
 [[test]]

crates/multi-value-xform/src/lib.rs

@@ -162,10 +162,12 @@ fn xform_one(
                 round_up_to_alignment(results_size, 8) + 8
             }
             walrus::ValType::V128 => round_up_to_alignment(results_size, 16) + 16,
-            walrus::ValType::Anyref => anyhow::bail!(
-                "cannot multi-value transform functions that return \
-                 anyref, since they can't go into linear memory"
-            ),
+            walrus::ValType::Anyref | walrus::ValType::Funcref | walrus::ValType::Nullref => {
+                anyhow::bail!(
+                    "cannot multi-value transform functions that return \
+                     reference types, since they can't go into linear memory"
+                )
+            }
         };
     }
     // Round up to 16-byte alignment, since that's what LLVM's emitted Wasm code
@@ -284,7 +286,9 @@ fn xform_one(
                 );
                 offset += 16;
             }
-            walrus::ValType::Anyref => unreachable!(),
+            walrus::ValType::Anyref | walrus::ValType::Funcref | walrus::ValType::Nullref => {
+                unreachable!()
+            }
         }
     }
 

crates/threads-xform/Cargo.toml

@@ -13,5 +13,5 @@ edition = "2018"
 
 [dependencies]
 anyhow = "1.0"
-walrus = "0.14.0"
+walrus = "0.16.0"
 wasm-bindgen-wasm-conventions = { path = "../wasm-conventions", version = "=0.2.62" }

crates/threads-xform/src/lib.rs

@@ -401,6 +401,9 @@ fn inject_start(
                         match segment.offset {
                             InitExpr::Global(id) => body.global_get(id),
                             InitExpr::Value(v) => body.const_(v),
+                            InitExpr::RefNull | InitExpr::RefFunc(_) => {
+                                panic!("not a valid i32 initializer")
+                            }
                         };
                         body.i32_const(0)
                             .i32_const(segment.len as i32)

crates/wasm-conventions/Cargo.toml

@@ -10,5 +10,5 @@ description = "Utilities for working with Wasm codegen conventions (usually esta
 edition = "2018"
 
 [dependencies]
-walrus = "0.14.0"
+walrus = "0.16.0"
 anyhow = "1.0"

crates/wasm-conventions/src/lib.rs

@@ -6,13 +6,13 @@
 //! * The shadow stack pointer
 //! * The canonical linear memory that contains the shadow stack
 
-#![deny(missing_docs, missing_debug_implementations)]
-
-use anyhow::{anyhow, bail, Error};
-use walrus::{ir::Value, GlobalId, GlobalKind, InitExpr, MemoryId, Module, ValType};
+use anyhow::{anyhow, bail, Result};
+use walrus::{
+    ir::Value, ElementId, FunctionId, GlobalId, GlobalKind, InitExpr, MemoryId, Module, ValType,
+};
 
 /// Get a Wasm module's canonical linear memory.
-pub fn get_memory(module: &Module) -> Result<MemoryId, Error> {
+pub fn get_memory(module: &Module) -> Result<MemoryId> {
     let mut memories = module.memories.iter().map(|m| m.id());
     let memory = memories.next();
     if memories.next().is_some() {
@@ -55,3 +55,40 @@ pub fn get_shadow_stack_pointer(module: &Module) -> Option<GlobalId> {
         _ => None,
     }
 }
+
+pub struct FunctionTableEntry {
+    pub element: ElementId,
+    pub idx: usize,
+    pub func: Option<FunctionId>,
+}
+
+/// Looks up a function table entry by index in the main function table.
+pub fn get_function_table_entry(module: &Module, idx: u32) -> Result<FunctionTableEntry> {
+    let table = module
+        .tables
+        .main_function_table()?
+        .ok_or_else(|| anyhow!("no function table found in module"))?;
+    let table = module.tables.get(table);
+    for &segment in table.elem_segments.iter() {
+        let segment = module.elements.get(segment);
+        let offset = match &segment.kind {
+            walrus::ElementKind::Active {
+                offset: InitExpr::Value(Value::I32(n)),
+                ..
+            } => *n as u32,
+            _ => continue,
+        };
+        let idx = (idx - offset) as usize;
+        match segment.members.get(idx) {
+            Some(slot) => {
+                return Ok(FunctionTableEntry {
+                    element: segment.id(),
+                    idx,
+                    func: slot.clone(),
+                })
+            }
+            None => continue,
+        }
+    }
+    bail!("failed to find `{}` in function table", idx);
+}

crates/wasm-interpreter/Cargo.toml

@@ -14,7 +14,8 @@ edition = '2018'
 [dependencies]
 anyhow = "1.0"
 log = "0.4"
-walrus = "0.14.0"
+walrus = "0.16.0"
+wasm-bindgen-wasm-conventions = { path = "../wasm-conventions", version = "0.2.62" }
 
 [dev-dependencies]
 tempfile = "3"

crates/wasm-interpreter/src/lib.rs

@@ -20,7 +20,7 @@
 
 use std::collections::{BTreeMap, HashMap, HashSet};
 use walrus::ir::Instr;
-use walrus::{FunctionId, LocalId, Module, TableId};
+use walrus::{ElementId, FunctionId, LocalId, Module, TableId};
 
 /// A ready-to-go interpreter of a wasm module.
 ///
@@ -159,7 +159,7 @@ impl Interpreter {
         &mut self,
         id: FunctionId,
         module: &Module,
-        entry_removal_list: &mut HashSet<usize>,
+        entry_removal_list: &mut HashSet<(ElementId, usize)>,
     ) -> Option<&[u32]> {
         // Call the `id` function. This is an internal `#[inline(never)]`
         // whose code is completely controlled by the `wasm-bindgen` crate, so
@@ -184,28 +184,19 @@ impl Interpreter {
 
         let args = vec![0; num_params];
         self.call(id, module, &args);
-        let descriptor_table_idx =
-            self.descriptor_table_idx
-                .take()
-                .expect("descriptor function should return index") as usize;
+        let descriptor_table_idx = self
+            .descriptor_table_idx
+            .take()
+            .expect("descriptor function should return index");
 
         // After we've got the table index of the descriptor function we're
         // interested go take a look in the function table to find what the
         // actual index of the function is.
-        let functions = self.functions.expect("function table should be present");
-        let functions = match &module.tables.get(functions).kind {
-            walrus::TableKind::Function(f) => f,
-            _ => unreachable!(),
-        };
-        let descriptor_id = functions
-            .elements
-            .get(descriptor_table_idx)
-            .expect("out of bounds read of function table")
-            .expect("attempting to execute null function");
-
-        // This is used later to actually remove the entry from the table, but
-        // we don't do the removal just yet
-        entry_removal_list.insert(descriptor_table_idx);
+        let entry =
+            wasm_bindgen_wasm_conventions::get_function_table_entry(module, descriptor_table_idx)
+                .expect("failed to find entry in function table");
+        let descriptor_id = entry.func.expect("element segment slot wasn't set");
+        entry_removal_list.insert((entry.element, entry.idx));
 
         // And now execute the descriptor!
         self.interpret_descriptor(descriptor_id, module)