//! Wizer: the WebAssembly pre-initializer!
//!
//! See the [`Wizer`] struct for details.

#![deny(missing_docs)]

use anyhow::Context;
use rayon::iter::{IntoParallelIterator, ParallelExtend, ParallelIterator};
use std::collections::{HashMap, HashSet};
use std::convert::TryFrom;
use std::fmt::Display;
use std::path::PathBuf;
#[cfg(feature = "structopt")]
use structopt::StructOpt;

const WASM_PAGE_SIZE: u32 = 65_536;
const NATIVE_PAGE_SIZE: u32 = 4_096;

const DEFAULT_INHERIT_STDIO: bool = true;
const DEFAULT_INHERIT_ENV: bool = false;

const DEFAULT_WASM_MULTI_VALUE: bool = true;
const DEFAULT_WASM_MULTI_MEMORY: bool = true;

/// Wizer: the WebAssembly pre-initializer!
///
/// Don't wait for your Wasm module to initialize itself, pre-initialize it!
/// Wizer instantiates your WebAssembly module, executes its initialization
/// function, and then serializes the instance's initialized state out into a
/// new WebAssembly module. Now you can use this new, pre-initialized
/// WebAssembly module to hit the ground running, without making your users wait
/// for that first-time set up code to complete.
///
/// ## Caveats
///
/// * The initialization function may not call any imported functions. Doing so
///   will trigger a trap and `wizer` will exit.
///
/// * The Wasm module may not import globals, tables, or memories.
///
/// * Reference types are not supported yet. This is tricky because it would
///   allow the Wasm module to mutate tables, and we would need to be able to
///   snapshot the new table state, but funcrefs and externrefs don't have
///   identity and aren't comparable in the Wasm spec, which makes snapshotting
///   difficult.
#[cfg_attr(feature = "structopt", derive(StructOpt))]
#[derive(Clone, Debug)]
pub struct Wizer {
    /// The Wasm export name of the function that should be executed to
    /// initialize the Wasm module.
    #[cfg_attr(
        feature = "structopt",
        structopt(short = "f", long = "init-func", default_value = "wizer.initialize")
    )]
    init_func: String,

    /// Any function renamings to perform.
    ///
    /// A renaming specification `dst=src` renames a function export `src` to
    /// `dst`, overwriting any previous `dst` export.
    ///
    /// Multiple renamings can be specified. It is an error to specify more than
    /// one source to rename to a destination name, or to specify more than one
    /// renaming destination for one source.
    ///
    /// This option can be used, for example, to replace a `_start` entry point
    /// in an initialized module with an alternate entry point.
    #[cfg_attr(
        feature = "structopt",
        structopt(
            short = "r",
            long = "rename-func",
            alias = "func-rename",
            value_name = "dst=src"
        )
    )]
    func_renames: Vec<String>,

    /// Allow WASI imports to be called during initialization.
    ///
    /// This can introduce diverging semantics because the initialization can
    /// observe nondeterminism that might have gone a different way at runtime
    /// than it did at initialization time.
    ///
    /// If your Wasm module uses WASI's `get_random` to add randomness to
    /// something as a security mitigation (e.g. something akin to ASLR or the
    /// way Rust's hash maps incorporate a random nonce) then note that, if the
    /// randomization is added during initialization time and you don't ever
    /// re-randomize at runtime, then that randomization will become per-module
    /// rather than per-instance.
    #[cfg_attr(feature = "structopt", structopt(long = "allow-wasi"))]
    allow_wasi: bool,

    /// When using WASI during initialization, should `stdin`, `stderr`, and
    /// `stdout` be inherited?
    ///
    /// This is true by default.
    #[cfg_attr(
        feature = "structopt",
        structopt(long = "inherit-stdio", value_name = "true|false")
    )]
    inherit_stdio: Option<bool>,

    /// When using WASI during initialization, should environment variables be
    /// inherited?
    ///
    /// This is false by default.
    #[cfg_attr(
        feature = "structopt",
        structopt(long = "inherit-env", value_name = "true|false")
    )]
    inherit_env: Option<bool>,

    /// When using WASI during initialization, which file system directories
    /// should be made available?
    ///
    /// None are available by default.
    #[cfg_attr(
        feature = "structopt",
        structopt(long = "dir", parse(from_os_str), value_name = "directory")
    )]
    dirs: Vec<PathBuf>,

    /// Enable or disable Wasm multi-memory proposal.
    ///
    /// Enabled by default.
    #[cfg_attr(feature = "structopt", structopt(long, value_name = "true|false"))]
    wasm_multi_memory: Option<bool>,

    /// Enable or disable the Wasm multi-value proposal.
    ///
    /// Enabled by default.
    #[cfg_attr(feature = "structopt", structopt(long, value_name = "true|false"))]
    wasm_multi_value: Option<bool>,
}

fn translate_val_type(ty: wasmparser::Type) -> wasm_encoder::ValType {
    use wasm_encoder::ValType;
    use wasmparser::Type::*;
    match ty {
        I32 => ValType::I32,
        I64 => ValType::I64,
        F32 => ValType::F32,
        F64 => ValType::F64,
        V128 | FuncRef | ExternRef | ExnRef => panic!("not supported"),
        Func | EmptyBlockType => unreachable!(),
    }
}

fn translate_global_type(ty: wasmparser::GlobalType) -> wasm_encoder::GlobalType {
    wasm_encoder::GlobalType {
        val_type: translate_val_type(ty.content_type),
        mutable: ty.mutable,
    }
}

struct FuncRenames {
    /// For a given export name that we encounter in the original module, a map
    /// to a new name, if any, to emit in the output module.
    rename_src_to_dst: HashMap<String, String>,
    /// A set of export names that we ignore in the original module (because
    /// they are overwritten by renamings).
    rename_dsts: HashSet<String>,
}

impl FuncRenames {
    fn parse(renames: &Vec<String>) -> anyhow::Result<FuncRenames> {
        let mut ret = FuncRenames {
            rename_src_to_dst: HashMap::new(),
            rename_dsts: HashSet::new(),
        };
        if renames.is_empty() {
            return Ok(ret);
        }

        for rename_spec in renames {
            let equal = rename_spec
                .trim()
                .find('=')
                .ok_or_else(|| anyhow::anyhow!("Invalid function rename part: {}", rename_spec))?;
            // TODO: use .split_off() when the API is stabilized.
            let dst = rename_spec[..equal].to_owned();
            let src = rename_spec[equal + 1..].to_owned();
            if ret.rename_dsts.contains(&dst) {
                anyhow::bail!("Duplicated function rename dst {}", dst);
            }
            if ret.rename_src_to_dst.contains_key(&src) {
                anyhow::bail!("Duplicated function rename src {}", src);
            }
            ret.rename_dsts.insert(dst.clone());
            ret.rename_src_to_dst.insert(src, dst);
        }

        Ok(ret)
    }
}

impl Wizer {
    /// Construct a new `Wizer` builder.
    pub fn new() -> Self {
        Wizer {
            init_func: "wizer.initialize".into(),
            func_renames: vec![],
            allow_wasi: false,
            inherit_stdio: None,
            inherit_env: None,
            dirs: vec![],
            wasm_multi_memory: None,
            wasm_multi_value: None,
        }
    }

    /// The export name of the initializer function.
    ///
    /// Defaults to `"wizer.initialize"`.
    pub fn init_func(&mut self, init_func: impl Into<String>) -> &mut Self {
        self.init_func = init_func.into();
        self
    }

    /// Add a function rename to perform.
    pub fn func_rename(&mut self, new_name: impl Display, old_name: impl Display) -> &mut Self {
        self.func_renames.push(format!("{}={}", new_name, old_name));
        self
    }

    /// Allow WASI imports to be called during initialization?
    ///
    /// This can introduce diverging semantics because the initialization can
    /// observe nondeterminism that might have gone a different way at runtime
    /// than it did at initialization time.
    ///
    /// If your Wasm module uses WASI's `get_random` to add randomness to
    /// something as a security mitigation (e.g. something akin to ASLR or the
    /// way Rust's hash maps incorporate a random nonce) then note that, if the
    /// randomization is added during initialization time and you don't ever
    /// re-randomize at runtime, then that randomization will become per-module
    /// rather than per-instance.
    ///
    /// Defaults to `false`.
    pub fn allow_wasi(&mut self, allow: bool) -> &mut Self {
        self.allow_wasi = allow;
        self
    }

    /// When using WASI during initialization, should `stdin`, `stdout`, and
    /// `stderr` be inherited?
    ///
    /// Defaults to `true`.
    pub fn inherit_stdio(&mut self, inherit: bool) -> &mut Self {
        self.inherit_stdio = Some(inherit);
        self
    }

    /// When using WASI during initialization, should the environment variables
    /// be inherited?
    ///
    /// Defaults to `false`.
    pub fn inherit_env(&mut self, inherit: bool) -> &mut Self {
        self.inherit_env = Some(inherit);
        self
    }

    /// When using WASI during initialization, which file system directories
    /// should be made available?
    ///
    /// None are available by default.
    pub fn dir(&mut self, directory: impl Into<PathBuf>) -> &mut Self {
        self.dirs.push(directory.into());
        self
    }

    /// Enable or disable the Wasm multi-memory proposal.
    ///
    /// Defaults to `true`.
    pub fn wasm_multi_memory(&mut self, enable: bool) -> &mut Self {
        self.wasm_multi_memory = Some(enable);
        self
    }

    /// Enable or disable the Wasm multi-value proposal.
    ///
    /// Defaults to `true`.
    pub fn wasm_multi_value(&mut self, enable: bool) -> &mut Self {
        self.wasm_multi_value = Some(enable);
        self
    }

    /// Initialize the given Wasm, snapshot it, and return the serialized
    /// snapshot as a new, pre-initialized Wasm module.
    pub fn run(&self, wasm: &[u8]) -> anyhow::Result<Vec<u8>> {
        // Parse rename spec.
        let renames = FuncRenames::parse(&self.func_renames)?;

        // Make sure we're given valid Wasm from the get go.
        self.wasm_validate(wasm)?;

        let wasm = self.prepare_input_wasm(wasm)?;
        debug_assert!(
            self.wasm_validate(&wasm).is_ok(),
            "if the Wasm was originally valid, then our preparation step shouldn't invalidate it"
        );

        let config = self.wasmtime_config()?;
        let engine = wasmtime::Engine::new(&config)?;
        let store = wasmtime::Store::new(&engine);
        let module = wasmtime::Module::new(store.engine(), &wasm)?;
        self.validate_init_func(&module)?;

        let instance = self.initialize(&store, &module)?;
        let snapshot = self.snapshot(&instance);
        let initialized_wasm = self.rewrite(&wasm, &snapshot, &renames);

        Ok(initialized_wasm)
    }

    // NB: keep this in sync with the wasmparser features.
    fn wasmtime_config(&self) -> anyhow::Result<wasmtime::Config> {
        let mut config = wasmtime::Config::new();

        // Enable Wasmtime's code cache. This makes it so that repeated
        // wizenings of the same Wasm module (e.g. with different WASI inputs)
        // doesn't require re-compiling the Wasm to native code every time.
        config.cache_config_load_default()?;

        // Proposals we support.
        config.wasm_multi_memory(self.wasm_multi_memory.unwrap_or(DEFAULT_WASM_MULTI_MEMORY));
        config.wasm_multi_value(self.wasm_multi_value.unwrap_or(DEFAULT_WASM_MULTI_VALUE));

        // Proposoals that we should add support for.
        config.wasm_reference_types(false);
        config.wasm_module_linking(false);
        config.wasm_simd(false);
        config.wasm_threads(false);
        config.wasm_bulk_memory(false);

        Ok(config)
    }

    // NB: keep this in sync with the Wasmtime config.
    fn wasm_features(&self) -> wasmparser::WasmFeatures {
        wasmparser::WasmFeatures {
            // Proposals that we support.
            multi_memory: self.wasm_multi_memory.unwrap_or(DEFAULT_WASM_MULTI_MEMORY),
            multi_value: self.wasm_multi_value.unwrap_or(DEFAULT_WASM_MULTI_VALUE),

            // Proposals that we should add support for.
            reference_types: false,
            module_linking: false,
            simd: false,
            threads: false,
            tail_call: false,
            bulk_memory: false,
            memory64: false,
            exceptions: false,

            // We will never want to enable this.
            deterministic_only: false,
        }
    }

    fn wasm_validate(&self, wasm: &[u8]) -> anyhow::Result<()> {
        log::debug!("Validating input Wasm");
        let mut validator = wasmparser::Validator::new();
        validator.wasm_features(self.wasm_features());
        validator.validate_all(wasm)?;
        Ok(())
    }

    /// Rewrite the input Wasm with our own custom exports for all globals, and
    /// memories. This way we can reflect on their values later on in the
    /// snapshot phase.
    ///
    /// TODO: will have to also export tables once we support reference types.
    fn prepare_input_wasm(&self, full_wasm: &[u8]) -> anyhow::Result<Vec<u8>> {
        log::debug!("Preparing input Wasm");

        let mut wasm = full_wasm;
        let mut parser = wasmparser::Parser::new(0);
        let mut module = wasm_encoder::Module::new();

        // Count how many globals and memories we see in this module, so that we
        // can export them all.
        let mut memory_count = 0;
        let mut global_count = 0;

        loop {
            let (payload, consumed) =
                match parser.parse(wasm, true).context("failed to parse Wasm")? {
                    wasmparser::Chunk::NeedMoreData(_) => anyhow::bail!("invalid Wasm module"),
                    wasmparser::Chunk::Parsed { payload, consumed } => (payload, consumed),
                };
            wasm = &wasm[consumed..];

            use wasmparser::Payload::*;
            use wasmparser::SectionReader;
            match payload {
                Version { .. } => continue,
                TypeSection(types) => {
                    module.section(&wasm_encoder::RawSection {
                        id: wasm_encoder::SectionId::Type as u8,
                        data: &full_wasm[types.range().start..types.range().end],
                    });
                }
                ImportSection(imports) => {
                    module.section(&wasm_encoder::RawSection {
                        id: wasm_encoder::SectionId::Import as u8,
                        data: &full_wasm[imports.range().start..imports.range().end],
                    });
                }
                AliasSection(_)
                | InstanceSection(_)
                | ModuleSectionStart { .. }
                | ModuleSectionEntry { .. } => {
                    anyhow::bail!("module linking is not supported yet")
                }
                FunctionSection(funcs) => {
                    module.section(&wasm_encoder::RawSection {
                        id: wasm_encoder::SectionId::Function as u8,
                        data: &full_wasm[funcs.range().start..funcs.range().end],
                    });
                }
                TableSection(tables) => {
                    module.section(&wasm_encoder::RawSection {
                        id: wasm_encoder::SectionId::Table as u8,
                        data: &full_wasm[tables.range().start..tables.range().end],
                    });
                }
                MemorySection(mems) => {
                    memory_count += mems.get_count();
                    module.section(&wasm_encoder::RawSection {
                        id: wasm_encoder::SectionId::Memory as u8,
                        data: &full_wasm[mems.range().start..mems.range().end],
                    });
                }
                GlobalSection(globals) => {
                    global_count += globals.get_count();
                    module.section(&wasm_encoder::RawSection {
                        id: wasm_encoder::SectionId::Global as u8,
                        data: &full_wasm[globals.range().start..globals.range().end],
                    });
                }
                ExportSection(mut exports) => {
                    let count = exports.get_count();
                    let mut exports_encoder = wasm_encoder::ExportSection::new();
                    for _ in 0..count {
                        let export = exports.read()?;
                        exports_encoder.export(
                            export.field,
                            match export.kind {
                                wasmparser::ExternalKind::Function => {
                                    wasm_encoder::Export::Function(export.index)
                                }
                                wasmparser::ExternalKind::Table => {
                                    wasm_encoder::Export::Table(export.index)
                                }
                                wasmparser::ExternalKind::Memory => {
                                    wasm_encoder::Export::Memory(export.index)
                                }
                                wasmparser::ExternalKind::Global => {
                                    wasm_encoder::Export::Global(export.index)
                                }
                                wasmparser::ExternalKind::Type
                                | wasmparser::ExternalKind::Module
                                | wasmparser::ExternalKind::Instance => {
                                    anyhow::bail!("module linking is not supported yet");
                                }
                                wasmparser::ExternalKind::Event => {
                                    anyhow::bail!("exceptions are not supported yet")
                                }
                            },
                        );
                    }
                    // Export all of the globals and memories under known names
                    // so we can manipulate them later.
                    for i in 0..global_count {
                        let name = format!("__wizer_global_{}", i);
                        exports_encoder.export(&name, wasm_encoder::Export::Global(i));
                    }
                    for i in 0..memory_count {
                        let name = format!("__wizer_memory_{}", i);
                        exports_encoder.export(&name, wasm_encoder::Export::Memory(i));
                    }
                    module.section(&exports_encoder);
                }
                StartSection { func, range: _ } => {
                    module.section(&wasm_encoder::StartSection {
                        function_index: func,
                    });
                }
                ElementSection(elems) => {
                    module.section(&wasm_encoder::RawSection {
                        id: wasm_encoder::SectionId::Element as u8,
                        data: &full_wasm[elems.range().start..elems.range().end],
                    });
                }
                DataCountSection { .. } => anyhow::bail!("bulk memory is not supported yet"),
                DataSection(data) => {
                    module.section(&wasm_encoder::RawSection {
                        id: wasm_encoder::SectionId::Data as u8,
                        data: &full_wasm[data.range().start..data.range().end],
                    });
                }
                CustomSection {
                    name,
                    data,
                    data_offset: _,
                } => {
                    module.section(&wasm_encoder::CustomSection { name, data });
                }
                CodeSectionStart {
                    range,
                    count: _,
                    size: _,
                } => {
                    let data = &full_wasm[range.start..range.end];
                    module.section(&wasm_encoder::RawSection {
                        id: wasm_encoder::SectionId::Code as u8,
                        data,
                    });
                }
                CodeSectionEntry(_) => continue,
                UnknownSection { .. } => anyhow::bail!("unknown section"),
                EventSection(_) => anyhow::bail!("exceptions are not supported yet"),
                End => return Ok(module.finish()),
            }
        }
    }

    /// Check that the module exports an initialization function, and that the
    /// function has the correct type.
    fn validate_init_func(&self, module: &wasmtime::Module) -> anyhow::Result<()> {
        log::debug!("Validating the exported initialization function");
        match module.get_export(&self.init_func) {
            Some(wasmtime::ExternType::Func(func_ty)) => {
                if func_ty.params().len() != 0 || func_ty.results().len() != 0 {
                    anyhow::bail!(
                        "the Wasm module's `{}` function export does not have type `[] -> []`",
                        &self.init_func
                    );
                }
            }
            Some(_) => anyhow::bail!(
                "the Wasm module's `{}` export is not a function",
                &self.init_func
            ),
            None => anyhow::bail!(
                "the Wasm module does not have a `{}` export",
                &self.init_func
            ),
        }
        Ok(())
    }

    /// Create dummy imports for instantiating the module.
    fn dummy_imports(
        &self,
        store: &wasmtime::Store,
        module: &wasmtime::Module,
        linker: &mut wasmtime::Linker,
    ) -> anyhow::Result<()> {
        log::debug!("Creating dummy imports");

        for imp in module.imports() {
            if linker.get_one_by_name(imp.module(), imp.name()).is_ok() {
                // Already defined, must be part of WASI.
                continue;
            }

            match imp.ty() {
                wasmtime::ExternType::Func(func_ty) => {
                    let trap = wasmtime::Trap::new(format!(
                        "cannot call imports within the initialization function; attempted \
                         to call `'{}' '{}'`",
                        imp.module(),
                        imp.name().unwrap()
                    ));
                    linker.define(
                        imp.module(),
                        imp.name().unwrap(),
                        wasmtime::Func::new(
                            store,
                            func_ty,
                            move |_caller: wasmtime::Caller, _params, _results| Err(trap.clone()),
                        ),
                    )?;
                }
                wasmtime::ExternType::Global(_global_ty) => {
                    // The Wasm module could use `global.get` to read the
                    // imported value and branch on that or use `global.set` to
                    // update it if it's mutable. We can't create a trapping
                    // dummy value, like we can for functions, and we can't
                    // define a "global segment" to update any imported values
                    // (although I suppose we could inject a `start` function).
                    anyhow::bail!("cannot initialize Wasm modules that import globals")
                }
                wasmtime::ExternType::Table(_table_ty) => {
                    // TODO: we could import a dummy table full of trapping
                    // functions *if* the reference types proposal is not
                    // enabled, so there is no way the Wasm module can
                    // manipulate the table. This would allow initializing such
                    // modules, as long as the initialization didn't do any
                    // `call_indirect`s.
                    anyhow::bail!("cannot initialize Wasm modules that import tables")
                }
                wasmtime::ExternType::Memory(_memory_ty) => {
                    // The Wasm module could read the memory and branch on its
                    // contents, and since we can't create a dummy memory that
                    // matches the "real" import, nor can we create a trapping
                    // dummy version like we can for functions, we can't support
                    // imported memories.
                    anyhow::bail!("cannot initialize Wasm modules that import memories")
                }
                _ => anyhow::bail!("module linking is not supported yet"),
            };
        }

        Ok(())
    }

    /// Instantiate the module and call its initialization function.
    fn initialize(
        &self,
        store: &wasmtime::Store,
        module: &wasmtime::Module,
    ) -> anyhow::Result<wasmtime::Instance> {
        log::debug!("Calling the initialization function");

        let mut linker = wasmtime::Linker::new(store);
        if self.allow_wasi {
            let mut ctx = wasi_cap_std_sync::WasiCtxBuilder::new();
            if self.inherit_stdio.unwrap_or(DEFAULT_INHERIT_STDIO) {
                ctx = ctx.inherit_stdio();
            }
            if self.inherit_env.unwrap_or(DEFAULT_INHERIT_ENV) {
                ctx = ctx.inherit_env()?;
            }
            for dir in &self.dirs {
                log::debug!("Preopening directory: {}", dir.display());
                let preopened = unsafe {
                    cap_std::fs::Dir::open_ambient_dir(dir)
                        .with_context(|| format!("failed to open directory: {}", dir.display()))?
                };
                ctx = ctx.preopened_dir(preopened, dir)?;
            }
            let ctx = ctx.build()?;
            wasmtime_wasi::Wasi::new(store, ctx).add_to_linker(&mut linker)?;
        }
        self.dummy_imports(&store, &module, &mut linker)?;
        let instance = linker.instantiate(module)?;

        let init_func = instance
            .get_typed_func::<(), ()>(&self.init_func)
            .expect("checked by `validate_init_func`");
        init_func
            .call(())
            .with_context(|| format!("the `{}` function trapped", self.init_func))?;

        Ok(instance)
    }

    /// Snapshot the given instance's globals, memories, and tables.
    ///
    /// TODO: when we support reference types, we will have to snapshot tables.
    fn snapshot<'a>(&self, instance: &'a wasmtime::Instance) -> Snapshot<'a> {
        // Get the initialized values of all globals.
        log::debug!("Snapshotting global values");
        let mut globals = vec![];
        let mut global_index = 0;
        loop {
            let name = format!("__wizer_global_{}", global_index);
            match instance.get_global(&name) {
                None => break,
                Some(global) => {
                    globals.push(global.get());
                    global_index += 1;
                }
            }
        }

        // Find and record non-zero regions of memory (in parallel).
        //
        // TODO: This could be really slow for large memories. Instead, we
        // should bring our own memories, protect the pages, and keep a table
        // with a dirty bit for each page, so we can just snapshot the pages
        // that actually got changed to non-zero values.
        log::debug!("Snapshotting memories");
        let mut memory_mins = vec![];
        let mut data_segments = vec![];
        let mut memory_index = 0;
        loop {
            let name = format!("__wizer_memory_{}", memory_index);
            match instance.get_memory(&name) {
                None => break,
                Some(memory) => {
                    memory_mins.push(memory.size());

                    let num_wasm_pages = memory.size();
                    let num_native_pages = num_wasm_pages * (WASM_PAGE_SIZE / NATIVE_PAGE_SIZE);

                    let memory: &'a [u8] = unsafe {
                        // Safe because no one else has a (potentially mutable)
                        // view to this memory and we know the memory will live
                        // as long as the instance is alive.
                        std::slice::from_raw_parts(memory.data_ptr(), memory.data_size())
                    };

                    // Consider each "native" page of the memory. (Scare quotes
                    // because we have no guarantee that anyone isn't using huge
                    // page sizes or something). Process each page in
                    // parallel. If any byte has changed, add the whole page as
                    // a data segment. This means that the resulting Wasm module
                    // should instantiate faster, since there are fewer segments
                    // to bounds check on instantiation. Engines could even
                    // theoretically recognize that each of these segments is
                    // page sized and aligned, and use lazy copy-on-write
                    // initialization of each instance's memory.
                    data_segments.par_extend((0..num_native_pages).into_par_iter().filter_map(
                        |i| {
                            let start = i * NATIVE_PAGE_SIZE;
                            let end = ((i + 1) * NATIVE_PAGE_SIZE) as usize;
                            let page = &memory[start as usize..end];
                            for byte in page {
                                if *byte != 0 {
                                    return Some(DataSegment {
                                        memory_index,
                                        offset: start as u32,
                                        data: page,
                                    });
                                }
                            }
                            None
                        },
                    ));

                    memory_index += 1;
                }
            }
        }

        // Sort data segments to enforce determinism in the face of the
        // parallelism above.
        data_segments.sort_by_key(|s| (s.memory_index, s.offset));

        // Merge any contiguous pages, so that the engine can initialize them
        // all at once (ideally with a single copy-on-write `mmap`) rather than
        // initializing each data segment individually.
        for i in (1..data_segments.len()).rev() {
            let a = &data_segments[i - 1];
            let b = &data_segments[i];

            // Only merge segments for the same memory.
            if a.memory_index != b.memory_index {
                continue;
            }

            // Only merge segments if they are contiguous.
            if a.offset + u32::try_from(a.data.len()).unwrap() != b.offset {
                continue;
            }

            // Okay, merge them together into `a` (so that the next iteration
            // can merge it with its predecessor) and then remove `b`!
            data_segments[i - 1].data = unsafe {
                debug_assert_eq!(
                    a.data
                        .as_ptr()
                        .offset(isize::try_from(a.data.len()).unwrap()),
                    b.data.as_ptr()
                );
                std::slice::from_raw_parts(a.data.as_ptr(), a.data.len() + b.data.len())
            };
            data_segments.remove(i);
        }

        Snapshot {
            globals,
            memory_mins,
            data_segments,
        }
    }

    fn rewrite(&self, full_wasm: &[u8], snapshot: &Snapshot, renames: &FuncRenames) -> Vec<u8> {
        log::debug!("Rewriting input Wasm to pre-initialized state");

        let mut wasm = full_wasm;
        let mut parser = wasmparser::Parser::new(0);
        let mut module = wasm_encoder::Module::new();

        // Encode the initialized data segments from the snapshot rather than
        // the original, uninitialized data segments.
        let mut added_data = false;
        let mut add_data_section = |module: &mut wasm_encoder::Module| {
            if added_data || snapshot.data_segments.is_empty() {
                return;
            }
            let mut data_section = wasm_encoder::DataSection::new();
            for DataSegment {
                memory_index,
                offset,
                data,
            } in &snapshot.data_segments
            {
                data_section.active(
                    *memory_index,
                    wasm_encoder::Instruction::I32Const(*offset as i32),
                    data.iter().copied(),
                );
            }
            module.section(&data_section);
            added_data = true;
        };

        loop {
            let (payload, consumed) = match parser.parse(wasm, true).unwrap() {
                wasmparser::Chunk::NeedMoreData(_) => unreachable!(),
                wasmparser::Chunk::Parsed { payload, consumed } => (payload, consumed),
            };
            wasm = &wasm[consumed..];

            use wasmparser::Payload::*;
            use wasmparser::SectionReader;
            match payload {
                Version { .. } => continue,
                TypeSection(types) => {
                    module.section(&wasm_encoder::RawSection {
                        id: wasm_encoder::SectionId::Type as u8,
                        data: &full_wasm[types.range().start..types.range().end],
                    });
                }
                ImportSection(imports) => {
                    module.section(&wasm_encoder::RawSection {
                        id: wasm_encoder::SectionId::Import as u8,
                        data: &full_wasm[imports.range().start..imports.range().end],
                    });
                }
                AliasSection(_) | InstanceSection(_) => unreachable!(),
                FunctionSection(funcs) => {
                    module.section(&wasm_encoder::RawSection {
                        id: wasm_encoder::SectionId::Function as u8,
                        data: &full_wasm[funcs.range().start..funcs.range().end],
                    });
                }
                TableSection(tables) => {
                    module.section(&wasm_encoder::RawSection {
                        id: wasm_encoder::SectionId::Table as u8,
                        data: &full_wasm[tables.range().start..tables.range().end],
                    });
                }
                MemorySection(mut mems) => {
                    // Set the minimum size of each memory to the snapshot's
                    // initialized size for that memory.
                    let mut memory_encoder = wasm_encoder::MemorySection::new();
                    for i in 0..mems.get_count() {
                        let memory = mems.read().unwrap();
                        match memory {
                            wasmparser::MemoryType::M32 { limits, shared: _ } => {
                                memory_encoder.memory(wasm_encoder::MemoryType {
                                    limits: wasm_encoder::Limits {
                                        min: snapshot.memory_mins[i as usize],
                                        max: limits.maximum,
                                    },
                                });
                            }
                            _ => unreachable!(),
                        }
                    }
                    module.section(&memory_encoder);
                }
                GlobalSection(mut globals) => {
                    // Encode the initialized values from the snapshot, rather
                    // than the original values.
                    let mut globals_encoder = wasm_encoder::GlobalSection::new();
                    for i in 0..globals.get_count() {
                        let global = globals.read().unwrap();
                        globals_encoder.global(
                            translate_global_type(global.ty),
                            match snapshot.globals[i as usize] {
                                wasmtime::Val::I32(x) => wasm_encoder::Instruction::I32Const(x),
                                wasmtime::Val::I64(x) => wasm_encoder::Instruction::I64Const(x),
                                wasmtime::Val::F32(x) => {
                                    wasm_encoder::Instruction::F32Const(f32::from_bits(x))
                                }
                                wasmtime::Val::F64(x) => {
                                    wasm_encoder::Instruction::F64Const(f64::from_bits(x))
                                }
                                _ => unreachable!(),
                            },
                        );
                    }
                    module.section(&globals_encoder);
                }
                ExportSection(mut exports) => {
                    // Remove the `__wizer_*` exports we added during the
                    // preparation phase, as well as the initialization
                    // function's export. Removing the latter will enable
                    // further Wasm optimizations (notably GC'ing unused
                    // functions) via `wasm-opt` and similar tools.
                    //
                    // We also perform renames at this stage. We have
                    // precomputed the list of old dsts to remove, and the
                    // mapping from srcs to dsts, so we can do this in one pass.
                    let count = exports.get_count();
                    let mut exports_encoder = wasm_encoder::ExportSection::new();
                    for _ in 0..count {
                        let export = exports.read().unwrap();
                        if export.field.starts_with("__wizer_") || export.field == self.init_func {
                            continue;
                        }
                        if !renames.rename_src_to_dst.contains_key(export.field)
                            && renames.rename_dsts.contains(export.field)
                        {
                            // A rename overwrites this export, and it is not renamed to another
                            // export, so skip it.
                            continue;
                        }
                        let field = match renames.rename_src_to_dst.get(export.field) {
                            None => export.field,
                            // A rename spec renames the export to this new
                            // name, so use it.
                            Some(dst) => dst,
                        };
                        exports_encoder.export(
                            field,
                            match export.kind {
                                wasmparser::ExternalKind::Function => {
                                    wasm_encoder::Export::Function(export.index)
                                }
                                wasmparser::ExternalKind::Table => {
                                    wasm_encoder::Export::Table(export.index)
                                }
                                wasmparser::ExternalKind::Memory => {
                                    wasm_encoder::Export::Memory(export.index)
                                }
                                wasmparser::ExternalKind::Global => {
                                    wasm_encoder::Export::Global(export.index)
                                }
                                wasmparser::ExternalKind::Type
                                | wasmparser::ExternalKind::Module
                                | wasmparser::ExternalKind::Instance
                                | wasmparser::ExternalKind::Event => unreachable!(),
                            },
                        );
                    }
                    module.section(&exports_encoder);
                }
                StartSection { .. } => {
                    // Skip the `start` function -- it's already been run!
                    continue;
                }
                ElementSection(elems) => {
                    module.section(&wasm_encoder::RawSection {
                        id: wasm_encoder::SectionId::Element as u8,
                        data: &full_wasm[elems.range().start..elems.range().end],
                    });
                }
                DataCountSection { .. } => unreachable!(),
                DataSection(_) => {
                    // TODO: supporting bulk memory will require copying over
                    // any active or declared segments.
                    add_data_section(&mut module);
                }
                CustomSection {
                    name,
                    data,
                    data_offset: _,
                } => {
                    // Some tools expect the name custom section to come last,
                    // even though custom sections are allowed in any order.
                    if name == "name" {
                        add_data_section(&mut module);
                    }

                    module.section(&wasm_encoder::CustomSection { name, data });
                }
                CodeSectionStart {
                    range,
                    count: _,
                    size: _,
                } => {
                    let data = &full_wasm[range.start..range.end];
                    module.section(&wasm_encoder::RawSection {
                        id: wasm_encoder::SectionId::Code as u8,
                        data,
                    });
                }
                CodeSectionEntry(_) => continue,
                ModuleSectionStart { .. }
                | ModuleSectionEntry { .. }
                | UnknownSection { .. }
                | EventSection(_) => unreachable!(),
                End => {
                    add_data_section(&mut module);
                    return module.finish();
                }
            }
        }
    }
}

/// A "snapshot" of non-default Wasm state after an instance has been
/// initialized.
struct Snapshot<'a> {
    /// Maps global index to its initialized value.
    globals: Vec<wasmtime::Val>,

    /// A new minimum size for each memory (in units of pages).
    memory_mins: Vec<u32>,

    /// Segments of non-zero memory.
    data_segments: Vec<DataSegment<'a>>,
}

struct DataSegment<'a> {
    /// The index of this data segment's memory.
    memory_index: u32,
    /// The offset within the memory that `data` should be copied to.
    offset: u32,
    /// This segment's (non-zero) data.
    data: &'a [u8],
}