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,13 +18,13 @@ 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' }
 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
+            }
         })
     }