Move wasm data sections out of wasmtime_environ::Module (#3231)

* Reduce indentation in `to_paged`

Use a few early-returns from `match` to avoid lots of extra indentation.

* Move wasm data sections out of `wasmtime_environ::Module`

This is the first step down the road of #3230. The long-term goal is
that `Module` is always `bincode`-decoded, but wasm data segments are a
possibly very-large portion of this residing in modules which we don't
want to shove through bincode. This refactors the internals of wasmtime
to be ok with this data living separately from the `Module` itself,
providing access at necessary locations.

Wasm data segments are now extracted from a wasm module and
concatenated directly. Data sections then describe ranges within this
concatenated list of data, and passive data works the same way. This
implementation does not lend itself to eventually optimizing the case
where passive data is dropped and no longer needed. That's left for a
future PR.

crates/environ/src/module.rs

@@ -1,11 +1,11 @@
 //! Data structures for representing decoded wasm modules.
 
-use crate::{EntityRef, PrimaryMap, Tunables};
+use crate::{EntityRef, ModuleTranslation, PrimaryMap, Tunables, WASM_PAGE_SIZE};
 use indexmap::IndexMap;
 use serde::{Deserialize, Serialize};
 use std::collections::{BTreeMap, BTreeSet};
 use std::convert::TryFrom;
-use std::sync::Arc;
+use std::ops::Range;
 use wasmtime_types::*;
 
 /// Implemenation styles for WebAssembly linear memory.
@@ -106,8 +106,11 @@ pub struct MemoryInitializer {
     pub base: Option<GlobalIndex>,
     /// The offset to add to the base.
     pub offset: u64,
-    /// The data to write into the linear memory.
+    /// The range of the data to write within the linear memory.
-    pub data: Box<[u8]>,
+    ///
+    /// This range indexes into a separately stored data section which will be
+    /// provided with the compiled module's code as well.
+    pub data: Range<u32>,
 }
 
 /// The type of WebAssembly linear memory initialization to use for a module.
@@ -125,6 +128,7 @@ pub enum MemoryInitialization {
     ///
     /// This is the default memory initialization type.
     Segmented(Vec<MemoryInitializer>),
+
     /// Memory initialization is paged.
     ///
     /// To be paged, the following requirements must be met:
@@ -141,111 +145,140 @@ pub enum MemoryInitialization {
     /// any work to point the kernel at the right linear memory page to use.
     Paged {
         /// The map of defined memory index to a list of initialization pages.
-        /// The list of page data is sparse, with None representing a zero page.
+        ///
-        /// Each page of initialization data is WebAssembly page-sized (64 KiB).
+        /// The list of page data is sparse, with each element starting with
-        /// The size of the list will be the maximum page written to by a data segment.
+        /// the offset in memory where it will be placed (specified here, as
-        map: PrimaryMap<DefinedMemoryIndex, Vec<Option<Box<[u8]>>>>,
+        /// a page index, with a `u64`). Each page of initialization data is
+        /// WebAssembly page-sized (64 KiB). Pages whose offset are not
+        /// specified in this array start with 0s in memory. The `Range`
+        /// indices, like those in `MemoryInitializer`, point within a data
+        /// segment that will come as an auxiliary descriptor with other data
+        /// such as the compiled code for the wasm module.
+        map: PrimaryMap<DefinedMemoryIndex, Vec<(u64, Range<u32>)>>,
         /// Whether or not an out-of-bounds data segment was observed.
         /// This is used to fail module instantiation after the pages are initialized.
         out_of_bounds: bool,
     },
 }
 
-impl MemoryInitialization {
+impl ModuleTranslation<'_> {
-    /// Attempts to convert segmented memory initialization into paged initialization for the given module.
+    /// Attempts to convert segmented memory initialization into paged
+    /// initialization for the module that this translation represents.
     ///
-    /// Returns `None` if the initialization cannot be paged or if it is already paged.
+    /// If this module's memory initialization is not compatible with paged
-    pub fn to_paged(&self, module: &Module) -> Option<Self> {
+    /// initialization then this won't change anything. Otherwise if it is
-        const WASM_PAGE_SIZE: usize = crate::WASM_PAGE_SIZE as usize;
+    /// compatible then the `memory_initialization` field will be updated.
+    pub fn try_paged_init(&mut self) {
+        let initializers = match &self.module.memory_initialization {
+            MemoryInitialization::Segmented(list) => list,
+            MemoryInitialization::Paged { .. } => return,
+        };
+        let page_size = u64::from(WASM_PAGE_SIZE);
+        let num_defined_memories =
+            self.module.memory_plans.len() - self.module.num_imported_memories;
+        let mut out_of_bounds = false;
 
-        match self {
+        // Initially all memories start out as all zeros, represented with a
-            Self::Paged { .. } => None,
+        // lack of entries in the `BTreeMap` here. The map indexes byte offset
-            Self::Segmented(initializers) => {
+        // (which is always wasm-page-aligned) to the contents of the page, with
-                let num_defined_memories = module.memory_plans.len() - module.num_imported_memories;
+        // missing entries implicitly as all zeros.
-                let mut out_of_bounds = false;
+        let mut page_contents = PrimaryMap::with_capacity(num_defined_memories);
-                let mut map = PrimaryMap::with_capacity(num_defined_memories);
+        for _ in 0..num_defined_memories {
+            page_contents.push(BTreeMap::new());
+        }
 
-                for _ in 0..num_defined_memories {
+        assert_eq!(initializers.len(), self.data.len());
-                    map.push(Vec::new());
+        for (initializer, data) in initializers.iter().zip(&self.data) {
+            let memory_index = match (
+                self.module.defined_memory_index(initializer.memory_index),
+                initializer.base.is_some(),
+            ) {
+                (None, _) | (_, true) => {
+                    // If the initializer references an imported memory or uses a global base,
+                    // the complete set of segments will need to be processed at module instantiation
+                    return;
                 }
+                (Some(index), false) => index,
+            };
+            if out_of_bounds {
+                continue;
+            }
 
-                for initializer in initializers {
+            // Perform a bounds check on the segment
-                    match (
+            //
-                        module.defined_memory_index(initializer.memory_index),
+            // As this segment is referencing a defined memory without a global
-                        initializer.base.is_some(),
+            // base, the last byte written to by the segment cannot exceed the
-                    ) {
+            // memory's initial minimum size
-                        (None, _) | (_, true) => {
+            let len = u64::try_from(initializer.data.len()).unwrap();
-                            // If the initializer references an imported memory or uses a global base,
+            let end = match initializer.offset.checked_add(len) {
-                            // the complete set of segments will need to be processed at module instantiation
+                Some(end) => end,
-                            return None;
+                None => {
-                        }
+                    out_of_bounds = true;
-                        (Some(index), false) => {
+                    continue;
-                            if out_of_bounds {
-                                continue;
-                            }
-
-                            // Perform a bounds check on the segment
-                            // As this segment is referencing a defined memory without a global base, the last byte
-                            // written to by the segment cannot exceed the memory's initial minimum size
-                            let offset = usize::try_from(initializer.offset).unwrap();
-                            let end = match offset.checked_add(initializer.data.len()) {
-                                Some(end) => end,
-                                None => {
-                                    out_of_bounds = true;
-                                    continue;
-                                }
-                            };
-                            if end
-                                > ((module.memory_plans[initializer.memory_index].memory.minimum
-                                    as usize)
-                                    * WASM_PAGE_SIZE)
-                            {
-                                out_of_bounds = true;
-                                continue;
-                            }
-
-                            let pages = &mut map[index];
-                            let mut page_index = offset / WASM_PAGE_SIZE;
-                            let mut page_offset = offset % WASM_PAGE_SIZE;
-                            let mut data_offset = 0;
-                            let mut data_remaining = initializer.data.len();
-
-                            if data_remaining == 0 {
-                                continue;
-                            }
-
-                            // Copy the initialization data by each WebAssembly-sized page (64 KiB)
-                            loop {
-                                if page_index >= pages.len() {
-                                    pages.resize(page_index + 1, None);
-                                }
-
-                                let page = pages[page_index].get_or_insert_with(|| {
-                                    vec![0; WASM_PAGE_SIZE].into_boxed_slice()
-                                });
-                                let len =
-                                    std::cmp::min(data_remaining, WASM_PAGE_SIZE - page_offset);
-
-                                page[page_offset..page_offset + len].copy_from_slice(
-                                    &initializer.data[data_offset..(data_offset + len)],
-                                );
-
-                                if len == data_remaining {
-                                    break;
-                                }
-
-                                page_index += 1;
-                                page_offset = 0;
-                                data_offset += len;
-                                data_remaining -= len;
-                            }
-                        }
-                    };
                 }
+            };
+            let memory = &self.module.memory_plans[initializer.memory_index].memory;
+            let initial_memory_end = memory.minimum * page_size;
+            if end > initial_memory_end {
+                out_of_bounds = true;
+                continue;
+            }
 
-                Some(Self::Paged { map, out_of_bounds })
+            // Perform the same style of initialization that instantiating the
+            // module performs at this point, except initialize our
+            // `page_contents` map which is indexed by page number and contains
+            // the actual page contents.
+            //
+            // This is done iteratively page-by-page until the entire data
+            // segment has been copied into the page map.
+            let contents = &mut page_contents[memory_index];
+            let mut page_index = initializer.offset / page_size;
+            let mut page_offset = (initializer.offset % page_size) as usize;
+            let mut data = &data[..];
+
+            while !data.is_empty() {
+                // If this page hasn't been seen before, then it starts out as
+                // all zeros.
+                let page = contents
+                    .entry(page_index)
+                    .or_insert_with(|| vec![0; page_size as usize]);
+                let page = &mut page[page_offset..];
+
+                let len = std::cmp::min(data.len(), page.len());
+                page[..len].copy_from_slice(&data[..len]);
+
+                page_index += 1;
+                page_offset = 0;
+                data = &data[len..];
             }
         }
+
+        // If we've gotten this far then we're switching to paged
+        // initialization. The contents of the initial wasm memory are
+        // specified by `page_contents`, so the job now is to transform data
+        // representation of wasm memory back into the representation we use
+        // in a `Module`.
+        //
+        // This is done by clearing `self.data`, the original data segments,
+        // since those are now all represented in `page_contents`. Afterwards
+        // all the pages are subsequently pushed onto `self.data` and the
+        // offsets within `self.data` are recorded in each segment that's part
+        // of `Paged`.
+        self.data.clear();
+        let mut map = PrimaryMap::with_capacity(page_contents.len());
+        let mut offset = 0;
+        for (memory, pages) in page_contents {
+            let mut page_offsets = Vec::with_capacity(pages.len());
+            for (byte_offset, page) in pages {
+                let end = offset + (page.len() as u32);
+                page_offsets.push((byte_offset, offset..end));
+                offset = end;
+                self.data.push(page.into());
+            }
+            let index = map.push(page_offsets);
+            assert_eq!(index, memory);
+        }
+        self.module.memory_initialization = MemoryInitialization::Paged { map, out_of_bounds };
     }
 }
 
@@ -351,12 +384,8 @@ pub struct Module {
     /// The map from passive element index (element segment index space) to index in `passive_elements`.
     pub passive_elements_map: BTreeMap<ElemIndex, usize>,
 
-    /// WebAssembly passive data segments.
-    #[serde(with = "passive_data_serde")]
-    pub passive_data: Vec<Arc<[u8]>>,
-
     /// The map from passive data index (data segment index space) to index in `passive_data`.
-    pub passive_data_map: BTreeMap<DataIndex, usize>,
+    pub passive_data_map: BTreeMap<DataIndex, Range<u32>>,
 
     /// WebAssembly function names.
     pub func_names: BTreeMap<FuncIndex, String>,
@@ -627,47 +656,3 @@ pub struct InstanceSignature {
     /// The name of what's being exported as well as its type signature.
     pub exports: IndexMap<String, EntityType>,
 }
-
-mod passive_data_serde {
-    use super::Arc;
-    use serde::{de::SeqAccess, de::Visitor, ser::SerializeSeq, Deserializer, Serializer};
-    use std::fmt;
-
-    pub(super) fn serialize<S>(data: &Vec<Arc<[u8]>>, ser: S) -> Result<S::Ok, S::Error>
-    where
-        S: Serializer,
-    {
-        let mut seq = ser.serialize_seq(Some(data.len()))?;
-        for v in data {
-            seq.serialize_element(v.as_ref())?;
-        }
-        seq.end()
-    }
-
-    struct PassiveDataVisitor;
-    impl<'de> Visitor<'de> for PassiveDataVisitor {
-        type Value = Vec<Arc<[u8]>>;
-
-        fn expecting(&self, formatter: &mut fmt::Formatter) -> fmt::Result {
-            formatter.write_str("a passive data sequence")
-        }
-
-        fn visit_seq<M>(self, mut access: M) -> Result<Self::Value, M::Error>
-        where
-            M: SeqAccess<'de>,
-        {
-            let mut data = Vec::with_capacity(access.size_hint().unwrap_or(0));
-            while let Some(value) = access.next_element::<Vec<u8>>()? {
-                data.push(value.into());
-            }
-            Ok(data)
-        }
-    }
-
-    pub(super) fn deserialize<'de, D>(de: D) -> Result<Vec<Arc<[u8]>>, D::Error>
-    where
-        D: Deserializer<'de>,
-    {
-        de.deserialize_seq(PassiveDataVisitor)
-    }
-}

crates/environ/src/module_environ.rs

@@ -9,6 +9,7 @@ use crate::{
     WasmFuncType, WasmResult,
 };
 use cranelift_entity::packed_option::ReservedValue;
+use std::borrow::Cow;
 use std::collections::{hash_map::Entry, HashMap};
 use std::convert::{TryFrom, TryInto};
 use std::mem;
@@ -67,6 +68,23 @@ pub struct ModuleTranslation<'data> {
     /// configuration.
     pub has_unparsed_debuginfo: bool,
 
+    /// List of data segments found in this module which should be concatenated
+    /// together for the final compiled artifact.
+    ///
+    /// These data segments, when concatenated, are indexed by the
+    /// `MemoryInitializer` type.
+    pub data: Vec<Cow<'data, [u8]>>,
+
+    /// Total size of all data pushed onto `data` so far.
+    total_data: u32,
+
+    /// List of passive element segments found in this module which will get
+    /// concatenated for the final artifact.
+    pub passive_data: Vec<&'data [u8]>,
+
+    /// Total size of all passive data pushed into `passive_data` so far.
+    total_passive_data: u32,
+
     /// When we're parsing the code section this will be incremented so we know
     /// which function is currently being defined.
     code_index: u32,
@@ -577,14 +595,32 @@ impl<'data> ModuleEnvironment<'data> {
 
                 let cnt = usize::try_from(data.get_count()).unwrap();
                 initializers.reserve_exact(cnt);
+                self.result.data.reserve_exact(cnt);
 
                 for (index, entry) in data.into_iter().enumerate() {
                     let wasmparser::Data { kind, data } = entry?;
+                    let mk_range = |total: &mut u32| -> Result<_, WasmError> {
+                        let range = u32::try_from(data.len())
+                            .ok()
+                            .and_then(|size| {
+                                let start = *total;
+                                let end = start.checked_add(size)?;
+                                Some(start..end)
+                            })
+                            .ok_or_else(|| {
+                                WasmError::Unsupported(format!(
+                                    "more than 4 gigabytes of data in wasm module",
+                                ))
+                            })?;
+                        *total += range.end - range.start;
+                        Ok(range)
+                    };
                     match kind {
                         DataKind::Active {
                             memory_index,
                             init_expr,
                         } => {
+                            let range = mk_range(&mut self.result.total_data)?;
                             let memory_index = MemoryIndex::from_u32(memory_index);
                             let mut init_expr_reader = init_expr.get_binary_reader();
                             let (base, offset) = match init_expr_reader.read_operator()? {
@@ -600,21 +636,23 @@ impl<'data> ModuleEnvironment<'data> {
                                     )));
                                 }
                             };
+
                             initializers.push(MemoryInitializer {
                                 memory_index,
                                 base,
                                 offset,
-                                data: data.into(),
+                                data: range,
                             });
+                            self.result.data.push(data.into());
                         }
                         DataKind::Passive => {
                             let data_index = DataIndex::from_u32(index as u32);
-                            let index = self.result.module.passive_data.len();
+                            let range = mk_range(&mut self.result.total_passive_data)?;
-                            self.result.module.passive_data.push(Arc::from(data));
+                            self.result.passive_data.push(data);
                             self.result
                                 .module
                                 .passive_data_map
-                                .insert(data_index, index);
+                                .insert(data_index, range);
                         }
                     }
                 }

crates/jit/src/instantiate.rs

@@ -51,6 +51,9 @@ pub struct CompilationArtifacts {
     /// ELF image with functions code.
     obj: Box<[u8]>,
 
+    /// All data segments referenced by this module, both active and passive.
+    wasm_data: Box<[u8]>,
+
     /// Descriptions of compiled functions
     funcs: PrimaryMap<DefinedFuncIndex, FunctionInfo>,
 
@@ -92,14 +95,45 @@ impl CompilationArtifacts {
         tunables: &Tunables,
     ) -> CompilationArtifacts {
         let ModuleTranslation {
-            module,
+            mut module,
             debuginfo,
             has_unparsed_debuginfo,
+            data,
+            passive_data,
             ..
         } = translation;
+
+        // Concatenate all the wasm data together, placing both active and
+        // passive data into the same chunk of data. Note that this
+        // implementation doesn't allow for unmapping or somehow releasing
+        // passive data on `data.drop`, and if we want to do that in the future
+        // we'll have to change this to store passive data segments separately
+        // from the main data segments.
+        //
+        // Also note that here we have to update all passive data segments and
+        // their relative indices.
+        let wasm_data_size = data
+            .iter()
+            .map(|s| s.len())
+            .chain(passive_data.iter().map(|s| s.len()))
+            .sum();
+        let mut wasm_data = Vec::with_capacity(wasm_data_size);
+        for data in data.iter() {
+            wasm_data.extend_from_slice(data);
+        }
+        let total_data_len = wasm_data.len();
+        for data in passive_data.iter() {
+            wasm_data.extend_from_slice(data);
+        }
+        for (_, range) in module.passive_data_map.iter_mut() {
+            range.start = range.start.checked_add(total_data_len as u32).unwrap();
+            range.end = range.end.checked_add(total_data_len as u32).unwrap();
+        }
+
         CompilationArtifacts {
             module: Arc::new(module),
             obj: obj.into_boxed_slice(),
+            wasm_data: wasm_data.into(),
             funcs,
             native_debug_info_present: tunables.generate_native_debuginfo,
             debug_info: if tunables.parse_wasm_debuginfo {
@@ -203,6 +237,15 @@ impl CompiledModule {
         &self.artifacts
     }
 
+    /// Returns the concatenated list of all data associated with this wasm
+    /// module.
+    ///
+    /// This is used for initialization of memories and all data ranges stored
+    /// in a `Module` are relative to the slice returned here.
+    pub fn wasm_data(&self) -> &[u8] {
+        &self.artifacts.wasm_data
+    }
+
     /// Return a reference-counting pointer to a module.
     pub fn module(&self) -> &Arc<Module> {
         &self.artifacts.module

crates/runtime/src/instance.rs

@@ -16,9 +16,9 @@ use memoffset::offset_of;
 use more_asserts::assert_lt;
 use std::alloc::Layout;
 use std::any::Any;
-use std::collections::BTreeMap;
 use std::convert::TryFrom;
 use std::hash::Hash;
+use std::ops::Range;
 use std::ptr::NonNull;
 use std::sync::Arc;
 use std::{mem, ptr, slice};
@@ -147,6 +147,12 @@ pub(crate) struct Instance {
     /// If the index is present in the set, the segment has been dropped.
     dropped_data: EntitySet<DataIndex>,
 
+    /// A slice pointing to all data that is referenced by this instance. This
+    /// data is managed externally so this is effectively an unsafe reference,
+    /// and this does not live for the `'static` lifetime so the API boundaries
+    /// here are careful to never hand out static references.
+    wasm_data: &'static [u8],
+
     /// Hosts can store arbitrary per-instance information here.
     ///
     /// Most of the time from Wasmtime this is `Box::new(())`, a noop
@@ -509,22 +515,6 @@ impl Instance {
         self.vmctx_plus_offset(self.offsets.vmctx_anyfuncs_begin())
     }
 
-    fn find_passive_segment<'a, I, D, T>(
-        index: I,
-        index_map: &BTreeMap<I, usize>,
-        data: &'a Vec<D>,
-        dropped: &EntitySet<I>,
-    ) -> &'a [T]
-    where
-        D: AsRef<[T]>,
-        I: EntityRef + Ord,
-    {
-        match index_map.get(&index) {
-            Some(index) if !dropped.contains(I::new(*index)) => data[*index].as_ref(),
-            _ => &[],
-        }
-    }
-
     /// The `table.init` operation: initializes a portion of a table with a
     /// passive element.
     ///
@@ -544,12 +534,13 @@ impl Instance {
         // inform `rustc` that the lifetime of the elements here are
         // disconnected from the lifetime of `self`.
         let module = self.module.clone();
-        let elements = Self::find_passive_segment(
+
-            elem_index,
+        let elements = match module.passive_elements_map.get(&elem_index) {
-            &module.passive_elements_map,
+            Some(index) if !self.dropped_elements.contains(elem_index) => {
-            &module.passive_elements,
+                module.passive_elements[*index].as_ref()
-            &self.dropped_elements,
+            }
-        );
+            _ => &[],
+        };
         self.table_init_segment(table_index, elements, dst, src, len)
     }
 
@@ -601,9 +592,7 @@ impl Instance {
     pub(crate) fn elem_drop(&mut self, elem_index: ElemIndex) {
         // https://webassembly.github.io/reference-types/core/exec/instructions.html#exec-elem-drop
 
-        if let Some(index) = self.module.passive_elements_map.get(&elem_index) {
+        self.dropped_elements.insert(elem_index);
-            self.dropped_elements.insert(ElemIndex::new(*index));
-        }
 
         // Note that we don't check that we actually removed a segment because
         // dropping a non-passive segment is a no-op (not a trap).
@@ -700,23 +689,21 @@ impl Instance {
         src: u32,
         len: u32,
     ) -> Result<(), Trap> {
-        // TODO: this `clone()` shouldn't be necessary but is used for now to
+        let range = match self.module.passive_data_map.get(&data_index).cloned() {
-        // inform `rustc` that the lifetime of the elements here are
+            Some(range) if !self.dropped_data.contains(data_index) => range,
-        // disconnected from the lifetime of `self`.
+            _ => 0..0,
-        let module = self.module.clone();
+        };
-        let data = Self::find_passive_segment(
+        self.memory_init_segment(memory_index, range, dst, src, len)
-            data_index,
+    }
-            &module.passive_data_map,
+
-            &module.passive_data,
+    pub(crate) fn wasm_data(&self, range: Range<u32>) -> &[u8] {
-            &self.dropped_data,
+        &self.wasm_data[range.start as usize..range.end as usize]
-        );
-        self.memory_init_segment(memory_index, &data, dst, src, len)
     }
 
     pub(crate) fn memory_init_segment(
         &mut self,
         memory_index: MemoryIndex,
-        data: &[u8],
+        range: Range<u32>,
         dst: u64,
         src: u32,
         len: u32,
@@ -724,6 +711,7 @@ impl Instance {
         // https://webassembly.github.io/bulk-memory-operations/core/exec/instructions.html#exec-memory-init
 
         let memory = self.get_memory(memory_index);
+        let data = self.wasm_data(range);
         let dst = self.validate_inbounds(memory.current_length, dst, len.into())?;
         let src = self.validate_inbounds(data.len(), src.into(), len.into())?;
         let len = len as usize;
@@ -741,9 +729,7 @@ impl Instance {
 
     /// Drop the given data segment, truncating its length to zero.
     pub(crate) fn data_drop(&mut self, data_index: DataIndex) {
-        if let Some(index) = self.module.passive_data_map.get(&data_index) {
+        self.dropped_data.insert(data_index);
-            self.dropped_data.insert(DataIndex::new(*index));
-        }
 
         // Note that we don't check that we actually removed a segment because
         // dropping a non-passive segment is a no-op (not a trap).

crates/runtime/src/instance/allocator.rs

@@ -61,6 +61,16 @@ pub struct InstanceAllocationRequest<'a> {
     /// We use a number of `PhantomPinned` declarations to indicate this to the
     /// compiler. More info on this in `wasmtime/src/store.rs`
     pub store: Option<*mut dyn Store>,
+
+    /// A list of all wasm data that can be referenced by the module that
+    /// will be allocated. The `Module` given here has active/passive data
+    /// segments that are specified as relative indices into this list of bytes.
+    ///
+    /// Note that this is an unsafe pointer. The pointer is expected to live for
+    /// the entire duration of the instance at this time. It's the
+    /// responsibility of the callee when allocating to ensure that this data
+    /// outlives the instance.
+    pub wasm_data: *const [u8],
 }
 
 /// An link error while instantiating a module.
@@ -337,10 +347,10 @@ fn initialize_memories(
         instance
             .memory_init_segment(
                 init.memory_index,
-                &init.data,
+                init.data.clone(),
                 get_memory_init_start(init, instance)?,
                 0,
-                u32::try_from(init.data.len()).unwrap(),
+                init.data.end - init.data.start,
             )
             .map_err(InstantiationError::Trap)?;
     }
@@ -391,13 +401,11 @@ fn initialize_instance(
                 let slice =
                     unsafe { slice::from_raw_parts_mut(memory.base, memory.current_length) };
 
-                for (page_index, page) in pages.iter().enumerate() {
+                for (page_index, page) in pages {
-                    if let Some(data) = page {
+                    debug_assert_eq!(page.end - page.start, WASM_PAGE_SIZE);
-                        debug_assert_eq!(data.len(), WASM_PAGE_SIZE as usize);
+                    let start = (*page_index * u64::from(WASM_PAGE_SIZE)) as usize;
-                        let start = page_index * WASM_PAGE_SIZE as usize;
+                    let end = start + WASM_PAGE_SIZE as usize;
-                        let end = start + WASM_PAGE_SIZE as usize;
+                    slice[start..end].copy_from_slice(instance.wasm_data(page.clone()));
-                        slice[start..end].copy_from_slice(data);
-                    }
                 }
             }
 
@@ -652,8 +660,9 @@ unsafe impl InstanceAllocator for OnDemandInstanceAllocator {
                 memories,
                 tables,
                 dropped_elements: EntitySet::with_capacity(req.module.passive_elements.len()),
-                dropped_data: EntitySet::with_capacity(req.module.passive_data.len()),
+                dropped_data: EntitySet::with_capacity(req.module.passive_data_map.len()),
                 host_state,
+                wasm_data: &*req.wasm_data,
                 vmctx: VMContext {
                     _marker: marker::PhantomPinned,
                 },

crates/runtime/src/instance/allocator/pooling.rs

@@ -364,6 +364,7 @@ impl InstancePool {
                     dropped_elements: EntitySet::new(),
                     dropped_data: EntitySet::new(),
                     host_state: Box::new(()),
+                    wasm_data: &[],
                     vmctx: VMContext {
                         _marker: marker::PhantomPinned,
                     },
@@ -382,6 +383,7 @@ impl InstancePool {
         instance.module = req.module.clone();
         instance.offsets = VMOffsets::new(HostPtr, instance.module.as_ref());
         instance.host_state = std::mem::replace(&mut req.host_state, Box::new(()));
+        instance.wasm_data = &*req.wasm_data;
 
         let mut limiter = req.store.and_then(|s| (*s).limiter());
         Self::set_instance_memories(
@@ -492,6 +494,7 @@ impl InstancePool {
         // fresh allocation later on.
         instance.module = self.empty_module.clone();
         instance.offsets = VMOffsets::new(HostPtr, &self.empty_module);
+        instance.wasm_data = &[];
 
         self.free_list.lock().unwrap().push(index);
     }
@@ -527,7 +530,6 @@ impl InstancePool {
         }
 
         debug_assert!(instance.dropped_data.is_empty());
-        instance.dropped_data.resize(module.passive_data.len());
 
         Ok(())
     }
@@ -1410,6 +1412,7 @@ mod test {
                             shared_signatures: VMSharedSignatureIndex::default().into(),
                             host_state: Box::new(()),
                             store: None,
+                            wasm_data: &[],
                         },
                     )
                     .expect("allocation should succeed"),
@@ -1432,6 +1435,7 @@ mod test {
                 shared_signatures: VMSharedSignatureIndex::default().into(),
                 host_state: Box::new(()),
                 store: None,
+                wasm_data: &[],
             },
         ) {
             Err(InstantiationError::Limit(3)) => {}

crates/runtime/src/instance/allocator/pooling/uffd.rs

@@ -269,7 +269,9 @@ unsafe fn initialize_wasm_page(
     if let MemoryInitialization::Paged { map, .. } = &instance.module.memory_initialization {
         let pages = &map[memory_index];
 
-        if let Some(Some(data)) = pages.get(page_index) {
+        let pos = pages.binary_search_by_key(&(page_index as u64), |k| k.0);
+        if let Ok(i) = pos {
+            let data = instance.wasm_data(pages[i].1.clone());
             debug_assert_eq!(data.len(), WASM_PAGE_SIZE);
 
             log::trace!(
@@ -530,6 +532,7 @@ mod test {
                                 shared_signatures: VMSharedSignatureIndex::default().into(),
                                 host_state: Box::new(()),
                                 store: None,
+                                wasm_data: &[],
                             },
                         )
                         .expect("instance should allocate"),

crates/wasmtime/src/instance.rs

@@ -736,6 +736,7 @@ impl<'a> Instantiator<'a> {
                         shared_signatures: self.cur.module.signatures().as_module_map().into(),
                         host_state: Box::new(Instance(instance_to_be)),
                         store: Some(store.traitobj),
+                        wasm_data: compiled_module.wasm_data(),
                     })?;
 
             // The instance still has lots of setup, for example

crates/wasmtime/src/module.rs

@@ -379,13 +379,7 @@ impl Module {
             // If configured, attempt to use paged memory initialization
             // instead of the default mode of memory initialization
             if cfg!(all(feature = "uffd", target_os = "linux")) {
-                if let Some(init) = translation
+                translation.try_paged_init();
-                    .module
-                    .memory_initialization
-                    .to_paged(&translation.module)
-                {
-                    translation.module.memory_initialization = init;
-                }
             }
 
             Ok(CompilationArtifacts::new(translation, obj, funcs, tunables))

crates/wasmtime/src/store.rs

@@ -211,6 +211,7 @@ impl<T> Store<T> {
                     imports: Default::default(),
                     module: Arc::new(wasmtime_environ::Module::default()),
                     store: None,
+                    wasm_data: &[],
                 })
                 .expect("failed to allocate default callee")
         };

crates/wasmtime/src/trampoline.rs

@@ -48,6 +48,7 @@ fn create_handle(
                 shared_signatures: shared_signature_id.into(),
                 host_state,
                 store: Some(store.traitobj),
+                wasm_data: &[],
             },
         )?;
 

crates/wasmtime/src/trampoline/func.rs

@@ -132,6 +132,7 @@ pub unsafe fn create_raw_function(
             shared_signatures: sig.into(),
             host_state,
             store: None,
+            wasm_data: &[],
         })?,
     )
 }