Register target features in custom section when using xforms (#3967)

Author: daxpedda

Diff for: CHANGELOG.md

@@ -45,9 +45,12 @@
 * Fix MSRV compilation.
   [#3927](https://github.com/rustwasm/wasm-bindgen/pull/3927)
 
-* Fixed `clippy::empty_docs` lint.
+* Fix `clippy::empty_docs` lint.
   [#3946](https://github.com/rustwasm/wasm-bindgen/pull/3946)
 
+* Fix missing target features in module when enabling reference types or multi-value transformation.
+  [#3967](https://github.com/rustwasm/wasm-bindgen/pull/3967)
+
 --------------------------------------------------------------------------------
 
 ## [0.2.92](https://github.com/rustwasm/wasm-bindgen/compare/0.2.91...0.2.92)

Diff for: crates/externref-xform/Cargo.toml

@@ -15,6 +15,7 @@ rust-version = "1.57"
 [dependencies]
 anyhow = "1.0"
 walrus = "0.20.2"
+wasm-bindgen-wasm-conventions = { path = "../wasm-conventions", version = "=0.2.92" }
 
 [dev-dependencies]
 rayon = "1.0"

Diff for: crates/externref-xform/src/lib.rs

@@ -15,7 +15,7 @@
 //! goal at least is to have valid wasm modules coming in that don't use
 //! `externref` and valid wasm modules going out which use `externref` at the fringes.
 
-use anyhow::{anyhow, bail, Error};
+use anyhow::{anyhow, bail, Context as _, Error};
 use std::cmp;
 use std::collections::{BTreeMap, HashMap, HashSet};
 use std::mem;
@@ -101,6 +101,10 @@ impl Context {
     /// large the function table is so we know what indexes to hand out when
     /// we're appending entries.
     pub fn prepare(&mut self, module: &mut Module) -> Result<(), Error> {
+        // Insert reference types to the target features section.
+        wasm_bindgen_wasm_conventions::insert_target_feature(module, "reference-types")
+            .context("failed to parse `target_features` custom section")?;
+
         // 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.

Diff for: crates/multi-value-xform/Cargo.toml

@@ -15,6 +15,7 @@ rust-version = "1.57"
 [dependencies]
 anyhow = "1.0"
 walrus = "0.20.2"
+wasm-bindgen-wasm-conventions = { path = "../wasm-conventions", version = "=0.2.92" }
 
 [dev-dependencies]
 rayon = "1.0"

Diff for: crates/multi-value-xform/src/lib.rs

@@ -92,6 +92,8 @@
 
 #![deny(missing_docs, missing_debug_implementations)]
 
+use anyhow::Context;
+
 /// Run the transformation.
 ///
 /// See the module-level docs for details on the transformation.
@@ -117,6 +119,10 @@ pub fn run(
     shadow_stack_pointer: walrus::GlobalId,
     to_xform: &[(walrus::FunctionId, usize, Vec<walrus::ValType>)],
 ) -> Result<Vec<walrus::FunctionId>, anyhow::Error> {
+    // Insert multi-value to the target features section.
+    wasm_bindgen_wasm_conventions::insert_target_feature(module, "multivalue")
+        .context("failed to parse `target_features` custom section")?;
+
     let mut wrappers = Vec::new();
     for (func, return_pointer_index, results) in to_xform {
         wrappers.push(xform_one(

Diff for: crates/wasm-conventions/Cargo.toml

@@ -11,5 +11,8 @@ edition = "2018"
 rust-version = "1.57"
 
 [dependencies]
+leb128 = "0.2"
 walrus = "0.20.2"
+# Matching the version `walrus` depends on.
+wasmparser = "0.80"
 anyhow = "1.0"

Diff for: crates/wasm-conventions/src/lib.rs

@@ -6,11 +6,14 @@
 //! * The shadow stack pointer
 //! * The canonical linear memory that contains the shadow stack
 
-use anyhow::{anyhow, bail, Result};
+use std::io::Cursor;
+
+use anyhow::{anyhow, bail, Context, Result};
 use walrus::{
     ir::Value, ElementId, FunctionBuilder, FunctionId, FunctionKind, GlobalId, GlobalKind,
-    InitExpr, MemoryId, Module, ValType,
+    InitExpr, MemoryId, Module, RawCustomSection, ValType,
 };
+use wasmparser::BinaryReader;
 
 /// Get a Wasm module's canonical linear memory.
 pub fn get_memory(module: &Module) -> Result<MemoryId> {
@@ -148,3 +151,72 @@ pub fn get_or_insert_start_builder(module: &mut Module) -> &mut FunctionBuilder
         .unwrap_local_mut()
         .builder_mut()
 }
+
+pub fn insert_target_feature(module: &mut Module, new_feature: &str) -> Result<()> {
+    // Taken from <https://github.com/bytecodealliance/wasm-tools/blob/f1898f46bb9d96f0f09682415cb6ccfd6a4dca79/crates/wasmparser/src/limits.rs#L27>.
+    anyhow::ensure!(new_feature.len() <= 100_000, "feature name too long");
+
+    // Try to find an existing section.
+    let section = module
+        .customs
+        .iter_mut()
+        .find(|(_, custom)| custom.name() == "target_features");
+
+    // If one exists, check if the target feature is already present.
+    let section = if let Some((_, section)) = section {
+        let section: &mut RawCustomSection = section
+            .as_any_mut()
+            .downcast_mut()
+            .context("failed to read section")?;
+        let mut reader = BinaryReader::new(&section.data);
+        // The first integer contains the target feature count.
+        let count = reader.read_var_u32()?;
+
+        // Try to find if the target feature is already present.
+        for _ in 0..count {
+            // First byte is the prefix.
+            let prefix_index = reader.current_position();
+            let prefix = reader.read_u8()? as u8;
+            // Read the feature.
+            let length = reader.read_var_u32()?;
+            let feature = reader.read_bytes(length as usize)?;
+
+            // If we found the target feature, we are done here.
+            if feature == new_feature.as_bytes() {
+                // Make sure we set any existing prefix to "enabled".
+                if prefix == b'-' {
+                    section.data[prefix_index] = b'+';
+                }
+
+                return Ok(());
+            }
+        }
+
+        section
+    } else {
+        let mut data = Vec::new();
+        leb128::write::unsigned(&mut data, 0).unwrap();
+        let id = module.customs.add(RawCustomSection {
+            name: String::from("target_features"),
+            data,
+        });
+        module.customs.get_mut(id).unwrap()
+    };
+
+    // If we couldn't find the target feature, insert it.
+
+    // The first byte contains an integer describing the target feature count, which we increase by one.
+    let mut data = Cursor::new(&section.data);
+    let count = leb128::read::unsigned(&mut data).unwrap();
+    let mut new_count = Vec::new();
+    leb128::write::unsigned(&mut new_count, count + 1).unwrap();
+    section.data.splice(0..data.position() as usize, new_count);
+    // Then we insert the "enabled" prefix at the end.
+    section.data.push(b'+');
+    // The next byte contains the length of the target feature string.
+    leb128::write::unsigned(&mut section.data, new_feature.len() as u64).unwrap();
+    // Lastly the target feature string is inserted.
+    section.data.extend(new_feature.as_bytes());
+
+    Ok(())
+}