tinywasm/store/memory/
paged.rs

1use alloc::boxed::Box;
2use alloc::vec::Vec;
3use core::cmp::min;
4use core::hint::cold_path;
5
6use super::{LinearMemory, checked_effective_addr};
7
8/// A sparse chunked linear memory.
9///
10/// This backend stores memory in fixed-size chunks, which makes growth cheaper because it avoids
11/// resizing and copying one large contiguous buffer.
12///
13/// The tradeoff is that reads and writes do a bit more bookkeeping and may need to cross chunk
14/// boundaries, so they are usually slightly slower than [`super::VecMemory`].
15///
16/// In particular, [`LinearMemory::read`] and [`LinearMemory::write`] return at most the bytes up to
17/// the end of the current chunk. Higher-level exact helpers loop over these short operations when
18/// they need a full range.
19pub struct PagedMemory {
20    len: usize,
21    chunk_size: usize,
22    chunk_shift: u32,
23    chunk_mask: usize,
24    chunks: Vec<Option<Box<[u8]>>>,
25}
26
27impl PagedMemory {
28    /// Tries to create a new sparse memory with `len` addressable bytes and the given `chunk_size`.
29    ///
30    /// Prefer this backend when grow behavior matters more than absolute read and write speed.
31    pub fn try_new(len: usize, chunk_size: usize) -> Result<Self, crate::Trap> {
32        assert!(chunk_size.is_power_of_two(), "chunk_size must be a power of two");
33
34        let mut memory = Self {
35            len: 0,
36            chunk_size,
37            chunk_shift: chunk_size.trailing_zeros(),
38            chunk_mask: chunk_size - 1,
39            chunks: Vec::new(),
40        };
41        memory.grow_to(len)?;
42        Ok(memory)
43    }
44
45    #[inline(always)]
46    fn allocate_chunk(&self) -> Result<Box<[u8]>, crate::Trap> {
47        let mut chunk = Vec::new();
48        match chunk.try_reserve_exact(self.chunk_size) {
49            Ok(()) => {}
50            Err(_) => {
51                cold_path();
52                return Err(crate::Trap::OutOfMemory);
53            }
54        }
55        chunk.resize(self.chunk_size, 0);
56        Ok(chunk.into_boxed_slice())
57    }
58
59    #[inline(always)]
60    fn chunk_mut(&mut self, chunk_idx: usize) -> Result<&mut [u8], crate::Trap> {
61        if self.chunks[chunk_idx].is_none() {
62            self.chunks[chunk_idx] = Some(self.allocate_chunk()?);
63        }
64
65        Ok(self.chunks[chunk_idx].as_deref_mut().unwrap_or_else(|| unreachable!()))
66    }
67
68    #[inline(always)]
69    fn chunk_slice(&self, chunk_idx: usize) -> Option<&[u8]> {
70        self.chunks[chunk_idx].as_deref()
71    }
72
73    #[inline(always)]
74    fn checked_end(&self, addr: usize, len: usize) -> Option<usize> {
75        let end = addr.checked_add(len)?;
76        if end > self.len {
77            return None;
78        }
79        Some(end)
80    }
81
82    #[inline(always)]
83    fn copy_within_single_chunk(&mut self, dst: usize, src: usize, len: usize) -> bool {
84        if len == 0 {
85            return true;
86        }
87
88        if self.checked_end(src, len).is_none() || self.checked_end(dst, len).is_none() {
89            return false;
90        }
91
92        let src_chunk_idx = src >> self.chunk_shift;
93        let dst_chunk_idx = dst >> self.chunk_shift;
94        if src_chunk_idx != dst_chunk_idx {
95            return false;
96        }
97
98        let src_offset = src & self.chunk_mask;
99        let dst_offset = dst & self.chunk_mask;
100        if src_offset + len > self.chunk_size || dst_offset + len > self.chunk_size {
101            return false;
102        }
103
104        if let Some(Some(chunk)) = self.chunks.get_mut(src_chunk_idx) {
105            chunk.copy_within(src_offset..src_offset + len, dst_offset);
106        }
107
108        true
109    }
110}
111
112#[cfg(feature = "debug")]
113impl core::fmt::Debug for PagedMemory {
114    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
115        let allocated_chunks = self.chunks.iter().filter(|chunk| chunk.is_some()).count();
116        f.debug_struct("PagedMemory")
117            .field("len", &self.len)
118            .field("chunk_size", &self.chunk_size)
119            .field("allocated_chunks", &allocated_chunks)
120            .finish()
121    }
122}
123
124impl LinearMemory for PagedMemory {
125    #[inline(always)]
126    fn len(&self) -> usize {
127        self.len
128    }
129
130    #[inline(always)]
131    fn grow_to(&mut self, new_len: usize) -> Result<(), crate::Trap> {
132        if new_len < self.len {
133            return Err(crate::Trap::MemoryOutOfBounds { offset: new_len, len: 0, max: self.len });
134        }
135
136        let new_chunk_count = if new_len == 0 { 0 } else { new_len.div_ceil(self.chunk_size) };
137        if new_chunk_count > self.chunks.len() {
138            match self.chunks.try_reserve_exact(new_chunk_count - self.chunks.len()) {
139                Ok(()) => {}
140                Err(_) => {
141                    cold_path();
142                    return Err(crate::Trap::OutOfMemory);
143                }
144            }
145            self.chunks.resize_with(new_chunk_count, || None);
146        } else {
147            self.chunks.truncate(new_chunk_count);
148        }
149
150        self.len = new_len;
151        Ok(())
152    }
153
154    #[inline(always)]
155    fn read(&self, addr: usize, dst: &mut [u8]) -> usize {
156        if addr >= self.len || dst.is_empty() {
157            return 0;
158        }
159
160        let chunk_idx = addr >> self.chunk_shift;
161        let chunk_offset = addr & self.chunk_mask;
162        let chunk_end = min((chunk_idx + 1) << self.chunk_shift, self.len);
163        let read_len = min(chunk_end - addr, dst.len());
164        if let Some(chunk) = self.chunk_slice(chunk_idx) {
165            dst[..read_len].copy_from_slice(&chunk[chunk_offset..chunk_offset + read_len]);
166        } else {
167            dst[..read_len].fill(0);
168        }
169
170        read_len
171    }
172
173    #[inline(always)]
174    fn write(&mut self, addr: usize, src: &[u8]) -> usize {
175        if addr >= self.len || src.is_empty() {
176            return 0;
177        }
178
179        let chunk_idx = addr >> self.chunk_shift;
180        let chunk_offset = addr & self.chunk_mask;
181        let write_len = min(min(self.chunk_size - chunk_offset, self.len - addr), src.len());
182
183        let Ok(chunk) = self.chunk_mut(chunk_idx) else {
184            return 0;
185        };
186        chunk[chunk_offset..chunk_offset + write_len].copy_from_slice(&src[..write_len]);
187        write_len
188    }
189
190    #[inline(always)]
191    fn write_all(&mut self, addr: usize, src: &[u8]) -> Option<()> {
192        let end = self.checked_end(addr, src.len())?;
193        let mut pos = addr;
194        let mut src_offset = 0;
195
196        while pos < end {
197            let chunk_idx = pos >> self.chunk_shift;
198            let chunk_offset = pos & self.chunk_mask;
199            let copy_len = min(self.chunk_size - chunk_offset, end - pos);
200
201            let chunk = self.chunk_mut(chunk_idx).ok()?;
202            chunk[chunk_offset..chunk_offset + copy_len].copy_from_slice(&src[src_offset..src_offset + copy_len]);
203
204            pos += copy_len;
205            src_offset += copy_len;
206        }
207
208        Some(())
209    }
210
211    #[inline(always)]
212    fn fill(&mut self, addr: usize, len: usize, val: u8) -> Option<()> {
213        let end = self.checked_end(addr, len)?;
214        let mut pos = addr;
215
216        while pos < end {
217            let chunk_idx = pos >> self.chunk_shift;
218            let chunk_offset = pos & self.chunk_mask;
219            let chunk_start = chunk_idx << self.chunk_shift;
220            let chunk_full_len = min(self.chunk_size, self.len - chunk_start);
221            let chunk_end = min(chunk_start + self.chunk_size, end);
222            let fill_len = chunk_end - pos;
223
224            if val == 0 {
225                if chunk_offset == 0 && fill_len == chunk_full_len {
226                    self.chunks[chunk_idx] = None;
227                } else if let Some(Some(chunk)) = self.chunks.get_mut(chunk_idx) {
228                    chunk[chunk_offset..chunk_offset + fill_len].fill(0);
229                }
230            } else {
231                self.chunk_mut(chunk_idx).ok()?[chunk_offset..chunk_offset + fill_len].fill(val);
232            }
233
234            pos = chunk_end;
235        }
236
237        Some(())
238    }
239
240    #[inline(always)]
241    fn copy_within(&mut self, dst: usize, src: usize, len: usize) -> Option<()> {
242        self.checked_end(src, len)?;
243        self.checked_end(dst, len)?;
244
245        if len == 0 || dst == src {
246            return Some(());
247        }
248
249        if self.copy_within_single_chunk(dst, src, len) {
250            return Some(());
251        }
252
253        let mut buf = [0u8; 256];
254
255        if dst < src || dst >= src + len {
256            let mut copied = 0;
257            while copied < len {
258                let chunk_len = min(buf.len(), len - copied);
259                self.read_exact(src + copied, &mut buf[..chunk_len])?;
260                self.write_all(dst + copied, &buf[..chunk_len])?;
261                copied += chunk_len;
262            }
263        } else {
264            let mut remaining = len;
265            while remaining > 0 {
266                let chunk_len = min(buf.len(), remaining);
267                let chunk_start = remaining - chunk_len;
268                self.read_exact(src + chunk_start, &mut buf[..chunk_len])?;
269                self.write_all(dst + chunk_start, &buf[..chunk_len])?;
270                remaining = chunk_start;
271            }
272        }
273
274        Some(())
275    }
276
277    #[inline(always)]
278    fn read_8(&self, base: u64, offset: u64) -> core::result::Result<u8, crate::Trap> {
279        let addr = checked_effective_addr::<1>(self.len, base, offset)?;
280        let chunk_idx = addr >> self.chunk_shift;
281        let chunk_offset = addr & self.chunk_mask;
282        Ok(self.chunk_slice(chunk_idx).map_or(0, |chunk| chunk[chunk_offset]))
283    }
284
285    #[inline(always)]
286    fn read_16(&self, base: u64, offset: u64) -> core::result::Result<[u8; 2], crate::Trap> {
287        let addr = checked_effective_addr::<2>(self.len, base, offset)?;
288        let chunk_idx = addr >> self.chunk_shift;
289        let chunk_offset = addr & self.chunk_mask;
290        if chunk_offset + 2 <= self.chunk_size {
291            return Ok(match self.chunk_slice(chunk_idx) {
292                Some(chunk) => chunk[chunk_offset..chunk_offset + 2].try_into().unwrap_or_else(|_| unreachable!()),
293                None => [0; 2],
294            });
295        }
296
297        let mut bytes = [0; 2];
298        self.read_exact(addr, &mut bytes).unwrap();
299        Ok(bytes)
300    }
301
302    #[inline(always)]
303    fn read_32(&self, base: u64, offset: u64) -> core::result::Result<[u8; 4], crate::Trap> {
304        let addr = checked_effective_addr::<4>(self.len, base, offset)?;
305        let chunk_idx = addr >> self.chunk_shift;
306        let chunk_offset = addr & self.chunk_mask;
307        if chunk_offset + 4 <= self.chunk_size {
308            return Ok(match self.chunk_slice(chunk_idx) {
309                Some(chunk) => chunk[chunk_offset..chunk_offset + 4].try_into().unwrap_or_else(|_| unreachable!()),
310                None => [0; 4],
311            });
312        }
313
314        let mut bytes = [0; 4];
315        self.read_exact(addr, &mut bytes).unwrap();
316        Ok(bytes)
317    }
318
319    #[inline(always)]
320    fn read_64(&self, base: u64, offset: u64) -> core::result::Result<[u8; 8], crate::Trap> {
321        let addr = checked_effective_addr::<8>(self.len, base, offset)?;
322        let chunk_idx = addr >> self.chunk_shift;
323        let chunk_offset = addr & self.chunk_mask;
324        if chunk_offset + 8 <= self.chunk_size {
325            return Ok(match self.chunk_slice(chunk_idx) {
326                Some(chunk) => chunk[chunk_offset..chunk_offset + 8].try_into().unwrap_or_else(|_| unreachable!()),
327                None => [0; 8],
328            });
329        }
330
331        let mut bytes = [0; 8];
332        self.read_exact(addr, &mut bytes).unwrap();
333        Ok(bytes)
334    }
335
336    #[inline(always)]
337    fn read_128(&self, base: u64, offset: u64) -> core::result::Result<[u8; 16], crate::Trap> {
338        let addr = checked_effective_addr::<16>(self.len, base, offset)?;
339        let chunk_idx = addr >> self.chunk_shift;
340        let chunk_offset = addr & self.chunk_mask;
341        if chunk_offset + 16 <= self.chunk_size {
342            return Ok(match self.chunk_slice(chunk_idx) {
343                Some(chunk) => chunk[chunk_offset..chunk_offset + 16].try_into().unwrap_or_else(|_| unreachable!()),
344                None => [0; 16],
345            });
346        }
347
348        let mut bytes = [0; 16];
349        self.read_exact(addr, &mut bytes).unwrap();
350        Ok(bytes)
351    }
352
353    #[inline(always)]
354    fn write_8(&mut self, base: u64, offset: u64, byte: u8) -> core::result::Result<(), crate::Trap> {
355        let addr = checked_effective_addr::<1>(self.len, base, offset)?;
356        let chunk_idx = addr >> self.chunk_shift;
357        let chunk_offset = addr & self.chunk_mask;
358        self.chunk_mut(chunk_idx)?[chunk_offset] = byte;
359        Ok(())
360    }
361
362    #[inline(always)]
363    fn write_16(&mut self, base: u64, offset: u64, bytes: [u8; 2]) -> core::result::Result<(), crate::Trap> {
364        let addr = checked_effective_addr::<2>(self.len, base, offset)?;
365        let chunk_idx = addr >> self.chunk_shift;
366        let chunk_offset = addr & self.chunk_mask;
367        if chunk_offset + 2 <= self.chunk_size {
368            self.chunk_mut(chunk_idx)?[chunk_offset..chunk_offset + 2].copy_from_slice(&bytes);
369        } else {
370            if self.write_all(addr, &bytes).is_none() {
371                cold_path();
372                return Err(crate::Trap::OutOfMemory);
373            }
374        }
375        Ok(())
376    }
377
378    #[inline(always)]
379    fn write_32(&mut self, base: u64, offset: u64, bytes: [u8; 4]) -> core::result::Result<(), crate::Trap> {
380        let addr = checked_effective_addr::<4>(self.len, base, offset)?;
381        let chunk_idx = addr >> self.chunk_shift;
382        let chunk_offset = addr & self.chunk_mask;
383        if chunk_offset + 4 <= self.chunk_size {
384            self.chunk_mut(chunk_idx)?[chunk_offset..chunk_offset + 4].copy_from_slice(&bytes);
385        } else {
386            if self.write_all(addr, &bytes).is_none() {
387                cold_path();
388                return Err(crate::Trap::OutOfMemory);
389            }
390        }
391        Ok(())
392    }
393
394    #[inline(always)]
395    fn write_64(&mut self, base: u64, offset: u64, bytes: [u8; 8]) -> core::result::Result<(), crate::Trap> {
396        let addr = checked_effective_addr::<8>(self.len, base, offset)?;
397        let chunk_idx = addr >> self.chunk_shift;
398        let chunk_offset = addr & self.chunk_mask;
399        if chunk_offset + 8 <= self.chunk_size {
400            self.chunk_mut(chunk_idx)?[chunk_offset..chunk_offset + 8].copy_from_slice(&bytes);
401        } else {
402            if self.write_all(addr, &bytes).is_none() {
403                cold_path();
404                return Err(crate::Trap::OutOfMemory);
405            }
406        }
407        Ok(())
408    }
409
410    #[inline(always)]
411    fn write_128(&mut self, base: u64, offset: u64, bytes: [u8; 16]) -> core::result::Result<(), crate::Trap> {
412        let addr = checked_effective_addr::<16>(self.len, base, offset)?;
413        let chunk_idx = addr >> self.chunk_shift;
414        let chunk_offset = addr & self.chunk_mask;
415        if chunk_offset + 16 <= self.chunk_size {
416            self.chunk_mut(chunk_idx)?[chunk_offset..chunk_offset + 16].copy_from_slice(&bytes);
417        } else {
418            if self.write_all(addr, &bytes).is_none() {
419                cold_path();
420                return Err(crate::Trap::OutOfMemory);
421            }
422        }
423        Ok(())
424    }
425}
426
427#[cfg(test)]
428mod tests {
429    use super::{LinearMemory, PagedMemory};
430
431    #[test]
432    fn paged_memory_reads_zeroes_from_sparse_chunks() {
433        let memory = PagedMemory::try_new(16, 4).expect("test memory should be constructible");
434        let mut dst = [1; 6];
435        assert_eq!(memory.read(5, &mut dst), 3);
436        assert_eq!(&dst[..3], &[0; 3]);
437        assert_eq!(&dst[3..], &[1; 3]);
438    }
439
440    #[test]
441    fn paged_memory_store_and_load_crosses_chunk_boundaries() {
442        let mut memory = PagedMemory::try_new(16, 4).expect("test memory should be constructible");
443        memory.write_all(3, &[1, 2, 3, 4, 5, 6]).unwrap();
444
445        let mut dst = [0; 6];
446        memory.read_exact(3, &mut dst).unwrap();
447        assert_eq!(dst, [1, 2, 3, 4, 5, 6]);
448    }
449
450    #[test]
451    fn paged_memory_copy_within_handles_overlap() {
452        let mut memory = PagedMemory::try_new(16, 4).expect("test memory should be constructible");
453        memory.write_all(0, &[1, 2, 3, 4, 5, 6]).unwrap();
454        memory.copy_within(2, 0, 6).unwrap();
455
456        let mut dst = [0; 8];
457        memory.read_exact(0, &mut dst).unwrap();
458        assert_eq!(dst, [1, 2, 1, 2, 3, 4, 5, 6]);
459    }
460
461    #[test]
462    fn paged_memory_write_stops_at_chunk_boundary() {
463        let mut memory = PagedMemory::try_new(16, 4).expect("test memory should be constructible");
464        assert_eq!(memory.write(3, &[1, 2, 3, 4]), 1);
465
466        let mut dst = [0; 4];
467        memory.read_exact(3, &mut dst).unwrap();
468        assert_eq!(dst, [1, 0, 0, 0]);
469    }
470}
tinywasm/store/memory/paged.rs

tinywasm/store/memory/
paged.rs
