js-component-bindgen 2.1.1

//! Support for transpiling core modules using multi-memory to those that don't
//! use multi-memory.
//!
//! Wasmtime's implementation of adapter modules between components requires the
//! usage of multi-memory for copying data back and forth between two
//! components. The multi-memory proposal is, at this time, not stable in any JS
//! engine. This module is an attempt to polyfill this until at such a time that
//! multi-memory can be used natively.
//!
//! The purpose of this module is to identify core wasms which require
//! multi-memory coming out of Wasmtime. These wasms are rewritten to not
//! actually use more than one memory. The implementation here is to replace all
//! memory instructions operating on memory index 1 or greater with function
//! calls where JS is the one that does the load/store/etc. This is not expected
//! to be fast at runtime but is intended to be just enough to get this working
//! in JS environments at this time. The true speed is expected to come with the
//! multi-memory proposal.
//!
//! This module exports a [`Translation`] which wraps a [`ModuleTranslation`]
//! either as a pass-through "normal" or an "augmented" version where
//! "augmented" means that the original wasm is not used but instead a
//! recompiled copy without multiple memories is used. When calculating the
//! imports for the "augmented" module the arguments for JS functions that
//! read/write memory are automatically injected and handled.
//!
//! Callers of this module need to have an implementation in JS for all of the
//! entries listed in [`AugmentedOp`], likely through the `DataView` class in
//! JS.
//!
//! Note that at this time this module is not intended to be a complete and
//! general purpose method of compiling multiple memories to single-memory
//! modules. This does not handle all instructions that use memory for example,
//! but only those that Wasmtime's adapter modules emits. It's possible to add
//! support for more instructions but such support isn't required at this time.
//! Examples of unsupported instructions are `i64.load8_u` and `memory.copy`.
//! Additionally core wasm sections such as data sections and tables are not
//! supported because, again, Wasmtime doesn't use it at this time.

use std::collections::{HashMap, HashSet};

use anyhow::{Result, bail};
use wasm_encoder::{
    CodeSection, EntityType, ExportKind, ExportSection, Function, FunctionSection, ImportSection,
    Module, TypeSection, reencode::Reencode,
};
use wasmparser::collections::IndexMap;
use wasmparser::{
    Export, ExternalKind, FunctionBody, Import, Parser, Payload, TypeRef, Validator, VisitOperator,
    VisitSimdOperator, WasmFeatures,
};
use wasmtime_environ::component::CoreDef;
use wasmtime_environ::{EntityIndex, MemoryIndex, ModuleTranslation, PrimaryMap};

pub enum Translation<'a> {
    Normal(ModuleTranslation<'a>),
    Augmented {
        original: ModuleTranslation<'a>,
        wasm: Vec<u8>,
        imports_removed: HashSet<(String, String)>,
        imports_added: Vec<(String, String, MemoryIndex, AugmentedOp)>,
    },
}

#[derive(Debug)]
pub enum AugmentedImport<'a> {
    CoreDef(&'a CoreDef),
    Memory { mem: &'a CoreDef, op: AugmentedOp },
}

#[derive(Debug, Copy, Clone, PartialEq, Eq, Hash)]
pub enum AugmentedOp {
    I32Load,
    I32Load8U,
    I32Load8S,
    I32Load16U,
    I32Load16S,
    I64Load,
    F32Load,
    F64Load,
    I32Store,
    I32Store8,
    I32Store16,
    I64Store,
    F32Store,
    F64Store,
    MemorySize,
}

impl<'a> Translation<'a> {
    /// Create a new translation
    pub fn new(translation: ModuleTranslation<'a>, multi_memory: bool) -> Result<Translation<'a>> {
        if multi_memory {
            return Ok(Translation::Normal(translation));
        }

        // Set up wasm features to use
        let mut features = WasmFeatures::default();
        features.set(WasmFeatures::MULTI_MEMORY, false);

        match Validator::new_with_features(features).validate_all(translation.wasm) {
            // This module validates without multi-memory, no need to augment
            // it
            Ok(_) => return Ok(Translation::Normal(translation)),
            Err(e) => {
                features.set(WasmFeatures::MULTI_MEMORY, true);
                match Validator::new_with_features(features).validate_all(translation.wasm) {
                    // This module validates with multi-memory, so fall through
                    // to augmentation.
                    Ok(_) => {}

                    // This appears to not validate at all.
                    Err(_) => return Err(e.into()),
                }
            }
        }

        let mut augmenter = Augmenter {
            translation: &translation,
            imports_removed: Default::default(),
            imports_added: Default::default(),
            imported_funcs: Default::default(),
            imported_memories: Default::default(),
            imports: Default::default(),
            exports: Default::default(),
            local_func_tys: Default::default(),
            local_funcs: Default::default(),
            types: Default::default(),
            augments: Default::default(),
        };
        let wasm = augmenter.run()?;
        Ok(Translation::Augmented {
            wasm,
            imports_removed: augmenter.imports_removed,
            imports_added: augmenter.imports_added,
            original: translation,
        })
    }

    /// Returns the encoded wasm that represents this module, automatically
    /// returning the augmented version if multi-memory augmentation was
    /// required.
    pub fn wasm(&self) -> &[u8] {
        match self {
            Translation::Normal(translation) => translation.wasm,
            Translation::Augmented { wasm, .. } => wasm,
        }
    }

    /// Returns an iterator over the imports for this module using the `args` as
    /// supplied to the original module.
    ///
    /// The returned imports are either those within `args` or augmented
    /// versions based on `args` that perform an `AugmentedOp`.
    pub fn imports<'b>(
        &'b self,
        args: &'b [CoreDef],
    ) -> Vec<(&'b str, &'b str, AugmentedImport<'b>)> {
        match self {
            Translation::Normal(translation) => {
                assert_eq!(translation.module.imports().len(), args.len());
                translation
                    .module
                    .imports()
                    .zip(args)
                    .map(|((module, name, _), arg)| (module, name, AugmentedImport::CoreDef(arg)))
                    .collect()
            }
            Translation::Augmented {
                original,
                imports_removed,
                imports_added,
                ..
            } => {
                let mut ret = Vec::new();
                let mut memories: PrimaryMap<MemoryIndex, &'b CoreDef> = PrimaryMap::new();
                for ((module, name, _), arg) in original.module.imports().zip(args) {
                    if imports_removed.contains(&(module.to_string(), name.to_string())) {
                        memories.push(arg);
                    } else {
                        ret.push((module, name, AugmentedImport::CoreDef(arg)));
                    }
                }
                for (module, name, index, op) in imports_added {
                    ret.push((
                        module,
                        name,
                        AugmentedImport::Memory {
                            mem: memories[*index],
                            op: *op,
                        },
                    ));
                }
                ret
            }
        }
    }

    /// Returns the exports of this module, which are not modified by
    /// augmentation.
    pub fn exports(&self) -> IndexMap<String, EntityIndex> {
        let (translation, exports) = match self {
            Translation::Normal(translation) => (translation, &translation.module.exports),
            Translation::Augmented { original, .. } => (original, &original.module.exports),
        };

        exports
            .iter()
            .map(|(atom, entity_idx)| {
                (String::from(&translation.module.strings[atom]), *entity_idx)
            })
            .collect()
    }
}

pub struct Augmenter<'a> {
    translation: &'a ModuleTranslation<'a>,
    imports_removed: HashSet<(String, String)>,
    imports_added: Vec<(String, String, MemoryIndex, AugmentedOp)>,
    augments: HashMap<(MemoryIndex, AugmentedOp), u32>,

    types: Vec<wasmparser::FuncType>,
    imports: Vec<Import<'a>>,
    imported_funcs: u32,
    imported_memories: u32,
    exports: Vec<Export<'a>>,
    local_funcs: Vec<FunctionBody<'a>>,
    local_func_tys: Vec<u32>,
}

impl Augmenter<'_> {
    fn run(&mut self) -> Result<Vec<u8>> {
        // The first step is to parse the input original wasm and learn about
        // its structure. This validates that all the sections are supported and
        // records various bits of information about the module within `self`.
        for payload in Parser::new(0).parse_all(self.translation.wasm) {
            match payload? {
                Payload::TypeSection(s) => {
                    for grp in s.into_iter_err_on_gc_types() {
                        self.types.push(grp?);
                    }
                }
                Payload::ImportSection(section) => {
                    for import in section.into_imports() {
                        let i = import?;
                        match i.ty {
                            TypeRef::Func(_) => self.imported_funcs += 1,
                            TypeRef::Memory(_) => {
                                if self.imported_memories > 0 {
                                    let ok = self
                                        .imports_removed
                                        .insert((i.module.to_string(), i.name.to_string()));
                                    assert!(ok);
                                    continue;
                                }
                                self.imported_memories += 1;
                            }
                            _ => {}
                        }
                        self.imports.push(i);
                    }
                }
                Payload::ExportSection(s) => {
                    for e in s {
                        let e = e?;
                        self.exports.push(e);
                    }
                }
                Payload::FunctionSection(s) => {
                    for ty in s {
                        let ty = ty?;
                        self.local_func_tys.push(ty);
                    }
                }
                Payload::CodeSectionEntry(body) => {
                    self.local_funcs.push(body);
                }

                // NB: these sections are theoretically possible to handle but
                // are not required at this time.
                Payload::DataCountSection { .. }
                | Payload::GlobalSection(_)
                | Payload::TableSection(_)
                | Payload::MemorySection(_)
                | Payload::ElementSection(_)
                | Payload::DataSection(_)
                | Payload::StartSection { .. }
                | Payload::TagSection(_)
                | Payload::UnknownSection { .. } => {
                    bail!("unsupported section found in module using multiple memories")
                }

                // component-model related things that shouldn't show up
                Payload::ModuleSection { .. }
                | Payload::ComponentSection { .. }
                | Payload::InstanceSection(_)
                | Payload::ComponentInstanceSection(_)
                | Payload::ComponentAliasSection(_)
                | Payload::ComponentCanonicalSection(_)
                | Payload::ComponentStartSection { .. }
                | Payload::ComponentImportSection(_)
                | Payload::CoreTypeSection(_)
                | Payload::ComponentExportSection(_)
                | Payload::ComponentTypeSection(_) => {
                    bail!("component section found in module using multiple memories")
                }

                _ => {}
            }
        }

        // After the module has been parsed next the set of adapter functions is
        // determined. This is done by parsing all instructions in the module
        // and looking for anything that operates on memory index 1 or greater.
        //
        // This will fill out `self.augments` which is a list of functionality
        // that must be provided by JS to mutate non-index-0 memories.
        for body in self.local_funcs.clone() {
            let mut reader = body.get_operators_reader()?;
            while !reader.eof() {
                reader.visit_operator(&mut CollectMemOps(self))?;
            }
        }

        // And now at the end we've got all the information for encoding so
        // begin that process.
        self.encode()
    }

    fn augment_op(&mut self, mem: u32, op: AugmentedOp) {
        // Memory 0 stays in the module and isn't removed, so no need to
        // register an augmentation.
        if mem == 0 {
            return;
        }
        let index = MemoryIndex::from_u32(mem - 1);
        self.augments.entry((index, op)).or_insert_with(|| {
            let idx = self.imported_funcs + self.imports_added.len() as u32;
            self.imports_added.push((
                "augments".to_string(),
                format!("mem{mem} {op:?}"),
                index,
                op,
            ));
            idx
        });
    }

    fn encode(&self) -> Result<Vec<u8>> {
        let mut module = Module::new();
        let mut reencoder = MultiMemoryCoreReencoder { augmenter: self };

        // Types are all passed through as-is to retain the same type section as
        // before.
        let mut types = TypeSection::new();
        for ty in &self.types {
            let params = ty
                .params()
                .iter()
                .map(|ty| reencoder.val_type(*ty))
                .collect::<std::result::Result<Vec<_>, _>>()?;
            let results = ty
                .results()
                .iter()
                .map(|ty| reencoder.val_type(*ty))
                .collect::<std::result::Result<Vec<_>, _>>()?;
            types.ty().function(params, results);
        }

        // Pass through all of `self.imports` into the import section. This will
        // already have imports of multiple memories removed so this will import
        // at most one memory.
        let mut imports = ImportSection::new();
        for import in self.imports.iter() {
            let ty = reencoder.entity_type(import.ty)?;
            imports.import(import.module, import.name, ty);
        }

        // After the normal imports are all registered next the
        // memory-modification-functions are all imported. This is a new
        // addition to this module which shifts all functions in the index
        // space, hence the rewriting of all function bodies below.
        //
        // Each augmentation function declares its type signature in the type
        // section at the end of the type section to avoid tampering with the
        // type section's original index spaces. It would be more efficient to
        // not redeclare function signatures and reuse existing function
        // signatures, but that's left as an optimization for a later date.
        for (module, name, _, op) in self.imports_added.iter() {
            let cnt = types.len();
            op.encode_type(&mut types);
            imports.import(module, name, EntityType::Function(cnt));
        }

        // The function section remains the same as we're not tampering with the
        // count or types of all local functions.
        let mut funcs = FunctionSection::new();
        for ty in self.local_func_tys.iter() {
            funcs.function(*ty);
        }

        // Exports all remain the same with the one caveat that the function
        // index space has changed so those indices are remapped.
        let mut exports = ExportSection::new();
        for e in self.exports.iter() {
            let (kind, index) = match e.kind {
                ExternalKind::Func => (ExportKind::Func, self.remap_func(e.index)),
                ExternalKind::Table => (ExportKind::Table, e.index),
                ExternalKind::Global => (ExportKind::Global, e.index),
                ExternalKind::Memory => {
                    assert!(e.index < 1);
                    (ExportKind::Memory, e.index)
                }
                ExternalKind::Tag => (ExportKind::Tag, e.index),
                ExternalKind::FuncExact => (ExportKind::Func, self.remap_func(e.index)),
            };
            exports.export(e.name, kind, index);
        }

        // Finally the code section is remapped. This is done by translating
        // operator-by-operator from `wasmparser` to `wasm-encoder`. This
        // is where instructions like `i32.load 1` will become `call
        // $i32_load_memory_1`.
        let mut code = CodeSection::new();
        for body in self.local_funcs.iter() {
            let mut locals = Vec::new();

            for local in body.get_locals_reader()? {
                let (cnt, ty) = local?;
                locals.push((cnt, reencoder.val_type(ty)?));
            }

            let mut f = Function::new(locals);

            let mut ops = body.get_operators_reader()?;
            while !ops.eof() {
                reencoder.translate(&mut f, ops.read()?)?;
            }

            code.function(&f);
        }

        module.section(&types);
        module.section(&imports);
        module.section(&funcs);
        module.section(&exports);
        module.section(&code);

        Ok(module.finish())
    }

    fn remap_func(&self, index: u32) -> u32 {
        if index < self.imported_funcs {
            index
        } else {
            index + self.imports_added.len() as u32
        }
    }

    fn remap_memory(&self, index: u32) -> u32 {
        assert!(index < 1);
        index
    }
}
struct CollectMemOps<'a, 'b>(&'a mut Augmenter<'b>);

macro_rules! define_visit {
    ($( @$proposal:ident $op:ident $({ $($arg:ident: $argty:ty),* })? => $visit:ident ($($ann:tt)*))*) => {
        $(
            #[allow(unreachable_code)]
            #[allow(unused)]
            fn $visit(&mut self $( $( ,$arg: $argty)* )?) {
                define_visit!(augment self $op $($($arg)*)?);
            }
        )*
    };

    // List of instructions that are augmented which register the memory index
    // and the relevant augmentation operation.
    (augment $self:ident I32Load $memarg:ident) => {
        $self.0.augment_op($memarg.memory, AugmentedOp::I32Load);
    };
    (augment $self:ident I64Load $memarg:ident) => {
        $self.0.augment_op($memarg.memory, AugmentedOp::I64Load);
    };
    (augment $self:ident F32Load $memarg:ident) => {
        $self.0.augment_op($memarg.memory, AugmentedOp::F32Load);
    };
    (augment $self:ident F64Load $memarg:ident) => {
        $self.0.augment_op($memarg.memory, AugmentedOp::F64Load);
    };
    (augment $self:ident I32Load8U $memarg:ident) => {
        $self.0.augment_op($memarg.memory, AugmentedOp::I32Load8U);
    };
    (augment $self:ident I32Load8S $memarg:ident) => {
        $self.0.augment_op($memarg.memory, AugmentedOp::I32Load8S);
    };
    (augment $self:ident I32Load16U $memarg:ident) => {
        $self.0.augment_op($memarg.memory, AugmentedOp::I32Load16U);
    };
    (augment $self:ident I32Load16S $memarg:ident) => {
        $self.0.augment_op($memarg.memory, AugmentedOp::I32Load16S);
    };
    (augment $self:ident I32Store $memarg:ident) => {
        $self.0.augment_op($memarg.memory, AugmentedOp::I32Store);
    };
    (augment $self:ident I64Store $memarg:ident) => {
        $self.0.augment_op($memarg.memory, AugmentedOp::I64Store);
    };
    (augment $self:ident F32Store $memarg:ident) => {
        $self.0.augment_op($memarg.memory, AugmentedOp::F32Store);
    };
    (augment $self:ident F64Store $memarg:ident) => {
        $self.0.augment_op($memarg.memory, AugmentedOp::F64Store);
    };
    (augment $self:ident I32Store8 $memarg:ident) => {
        $self.0.augment_op($memarg.memory, AugmentedOp::I32Store8);
    };
    (augment $self:ident I32Store16 $memarg:ident) => {
        $self.0.augment_op($memarg.memory, AugmentedOp::I32Store16);
    };

    (augment $self:ident MemorySize $mem:ident) => {
        $self.0.augment_op($mem, AugmentedOp::MemorySize);
    };

    // Catch-all which asserts that none of the `$arg` looks like a memory
    // index ty catch any missing instructions from the list above.
    (augment $self:ident $op:ident $($arg:ident)*) => {
        $(
            define_visit!(assert_not_mem $op $arg);
        )*
    };

    (assert_not_mem $op:ident mem) => {panic!(concat!("missed case ", stringify!($op)));};
    (assert_not_mem $op:ident src_mem) => {panic!(concat!("missed case ", stringify!($op)));};
    (assert_not_mem $op:ident dst_mem) => {panic!(concat!("missed case ", stringify!($op)));};
    (assert_not_mem $op:ident memarg) => {panic!(concat!("missed case ", stringify!($op)));};
    (assert_not_mem $op:ident $other:ident) => {};
}

impl<'a> VisitOperator<'a> for CollectMemOps<'_, 'a> {
    type Output = ();

    wasmparser::for_each_visit_operator!(define_visit);
}

impl<'a> VisitSimdOperator<'a> for CollectMemOps<'_, 'a> {
    wasmparser::for_each_visit_simd_operator!(define_visit);
}

impl AugmentedOp {
    fn encode_type(&self, section: &mut TypeSection) {
        use wasm_encoder::ValType::{F32, F64, I32, I64};
        match self {
            // Loads take two arguments: the first is the address being loaded
            // from and the second is the static offset that was listed on the
            // relevant load instruction.
            AugmentedOp::I32Load
            | AugmentedOp::I32Load8U
            | AugmentedOp::I32Load8S
            | AugmentedOp::I32Load16U
            | AugmentedOp::I32Load16S => {
                section.ty().function([I32, I32], [I32]);
            }
            AugmentedOp::I64Load => {
                section.ty().function([I32, I32], [I64]);
            }
            AugmentedOp::F32Load => {
                section.ty().function([I32, I32], [F32]);
            }
            AugmentedOp::F64Load => {
                section.ty().function([I32, I32], [F64]);
            }

            // Stores, like loads, take an additional argument than usual which
            // is the static offset on the store instruction.
            AugmentedOp::I32Store | AugmentedOp::I32Store8 | AugmentedOp::I32Store16 => {
                section.ty().function([I32, I32, I32], []);
            }
            AugmentedOp::I64Store => {
                section.ty().function([I32, I64, I32], []);
            }
            AugmentedOp::F32Store => {
                section.ty().function([I32, F32, I32], []);
            }
            AugmentedOp::F64Store => {
                section.ty().function([I32, F64, I32], []);
            }

            AugmentedOp::MemorySize => {
                section.ty().function([], [I32]);
            }
        }
    }
}

/// Re-encodes a given module to avoid multi-memory features
///
/// This re-encoder implements `wasm_encoder::reencode::Reencode` in order
/// to rewrite a given multi-memory module to run in environments without multi-memory
/// by using imports that bridge all non-0-idx memory access
///
/// We must generally remap functions and memories:
///
/// * Functions must be remapped due to injecting specialized memory-access imports (shifted by # of new imports)
/// * Memories are remapped to ensure there are no non-0-idx memories (i.e. no multi-memory use)
///
struct MultiMemoryCoreReencoder<'a, 'b> {
    augmenter: &'a Augmenter<'b>,
}

impl Reencode for MultiMemoryCoreReencoder<'_, '_> {
    type Error = std::convert::Infallible;

    fn function_index(
        &mut self,
        index: u32,
    ) -> Result<u32, wasm_encoder::reencode::Error<Self::Error>> {
        Ok(self.augmenter.remap_func(index))
    }

    fn memory_index(
        &mut self,
        index: u32,
    ) -> Result<u32, wasm_encoder::reencode::Error<Self::Error>> {
        Ok(self.augmenter.remap_memory(index))
    }
}

impl MultiMemoryCoreReencoder<'_, '_> {
    fn translate(&mut self, func: &mut Function, operator: wasmparser::Operator<'_>) -> Result<()> {
        use wasmparser::Operator::*;

        match operator {
            I32Load { memarg } if memarg.memory > 0 => {
                self.augment(func, AugmentedOp::I32Load, memarg)
            }
            I32Load8U { memarg } if memarg.memory > 0 => {
                self.augment(func, AugmentedOp::I32Load8U, memarg)
            }
            I32Load8S { memarg } if memarg.memory > 0 => {
                self.augment(func, AugmentedOp::I32Load8S, memarg)
            }
            I32Load16U { memarg } if memarg.memory > 0 => {
                self.augment(func, AugmentedOp::I32Load16U, memarg)
            }
            I32Load16S { memarg } if memarg.memory > 0 => {
                self.augment(func, AugmentedOp::I32Load16S, memarg)
            }
            I64Load { memarg } if memarg.memory > 0 => {
                self.augment(func, AugmentedOp::I64Load, memarg)
            }
            F32Load { memarg } if memarg.memory > 0 => {
                self.augment(func, AugmentedOp::F32Load, memarg)
            }
            F64Load { memarg } if memarg.memory > 0 => {
                self.augment(func, AugmentedOp::F64Load, memarg)
            }
            I32Store { memarg } if memarg.memory > 0 => {
                self.augment(func, AugmentedOp::I32Store, memarg)
            }
            I32Store8 { memarg } if memarg.memory > 0 => {
                self.augment(func, AugmentedOp::I32Store8, memarg)
            }
            I32Store16 { memarg } if memarg.memory > 0 => {
                self.augment(func, AugmentedOp::I32Store16, memarg)
            }
            I64Store { memarg } if memarg.memory > 0 => {
                self.augment(func, AugmentedOp::I64Store, memarg)
            }
            F32Store { memarg } if memarg.memory > 0 => {
                self.augment(func, AugmentedOp::F32Store, memarg)
            }
            F64Store { memarg } if memarg.memory > 0 => {
                self.augment(func, AugmentedOp::F64Store, memarg)
            }
            MemorySize { mem } if mem > 0 => {
                let mem = MemoryIndex::from_u32(mem - 1);
                let function = self.augmenter.augments[&(mem, AugmentedOp::MemorySize)];
                func.instruction(&wasm_encoder::Instruction::Call(function));
            }
            operator => {
                func.instruction(&self.instruction(operator)?);
            }
        }
        Ok(())
    }

    fn augment(&self, func: &mut Function, op: AugmentedOp, memarg: wasmparser::MemArg) {
        use wasm_encoder::Instruction::{Call, I32Const};

        let memory = MemoryIndex::from_u32(memarg.memory - 1);
        let function = self.augmenter.augments[&(memory, op)];
        func.instruction(&I32Const(memarg.offset as i32));
        func.instruction(&Call(function));
    }
}