ic_cdk/
stable.rs

1//! APIs to manage stable memory.
2//!
3//! You can check the [Internet Computer Specification](https://internetcomputer.org/docs/current/references/ic-interface-spec/#system-api-stable-memory)
4//! for a in-depth explanation of stable memory.
5// mod canister;
6// #[cfg(test)]
7// mod tests;
8
9use std::{error, fmt, io};
10
11/// WASM page size in bytes.
12pub const WASM_PAGE_SIZE_IN_BYTES: u64 = 64 * 1024; // 64KB
13
14static CANISTER_STABLE_MEMORY: CanisterStableMemory = CanisterStableMemory {};
15
16/// A trait defining the stable memory API which each canister running on the IC can make use of
17pub trait StableMemory {
18    /// Gets current size of the stable memory (in WASM pages).
19    fn stable_size(&self) -> u64;
20
21    /// Attempts to grow the stable memory by `new_pages` (added pages).
22    ///
23    /// Returns an error if it wasn't possible. Otherwise, returns the previous
24    /// size that was reserved.
25    ///
26    /// *Note*: Pages are 64KiB in WASM.
27    fn stable_grow(&self, new_pages: u64) -> Result<u64, StableMemoryError>;
28
29    /// Writes data to the stable memory location specified by an offset.
30    ///
31    /// Warning - this will panic if `offset + buf.len()` exceeds the current size of stable memory.
32    /// Use `stable_grow` to request more stable memory if needed.
33    fn stable_write(&self, offset: u64, buf: &[u8]);
34
35    /// Reads data from the stable memory location specified by an offset.
36    fn stable_read(&self, offset: u64, buf: &mut [u8]);
37}
38
39/// A possible error value when dealing with stable memory.
40#[derive(Debug)]
41pub enum StableMemoryError {
42    /// No more stable memory could be allocated.
43    OutOfMemory,
44    /// Attempted to read more stable memory than had been allocated.
45    OutOfBounds,
46}
47
48impl fmt::Display for StableMemoryError {
49    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
50        match self {
51            Self::OutOfMemory => f.write_str("Out of memory"),
52            Self::OutOfBounds => f.write_str("Read exceeds allocated memory"),
53        }
54    }
55}
56
57impl error::Error for StableMemoryError {}
58
59/// A standard implementation of [`StableMemory`].
60///
61/// Useful for creating [`StableWriter`] and [`StableReader`].
62#[derive(Default, Debug, Copy, Clone)]
63pub struct CanisterStableMemory {}
64
65impl StableMemory for CanisterStableMemory {
66    fn stable_size(&self) -> u64 {
67        // SAFETY: ic0.stable64_size is always safe to call.
68        unsafe { ic0::stable64_size() }
69    }
70
71    fn stable_grow(&self, new_pages: u64) -> Result<u64, StableMemoryError> {
72        // SAFETY: ic0.stable64_grow is always safe to call.
73        unsafe {
74            match ic0::stable64_grow(new_pages) {
75                u64::MAX => Err(StableMemoryError::OutOfMemory),
76                x => Ok(x),
77            }
78        }
79    }
80
81    fn stable_write(&self, offset: u64, buf: &[u8]) {
82        // SAFETY: `buf`, being &[u8], is a readable sequence of bytes, and therefore valid to pass to ic0.stable64_write.
83        unsafe {
84            ic0::stable64_write(offset, buf.as_ptr() as u64, buf.len() as u64);
85        }
86    }
87
88    fn stable_read(&self, offset: u64, buf: &mut [u8]) {
89        // SAFETY: `buf`, being &mut [u8], is a writable sequence of bytes, and therefore valid to pass to ic0.stable64_read.
90        unsafe {
91            ic0::stable64_read(buf.as_ptr() as u64, offset, buf.len() as u64);
92        }
93    }
94}
95
96/// Gets current size of the stable memory (in WASM pages).
97pub fn stable_size() -> u64 {
98    CANISTER_STABLE_MEMORY.stable_size()
99}
100
101/// Attempts to grow the stable memory by `new_pages` (added pages).
102///
103/// Returns an error if it wasn't possible. Otherwise, returns the previous
104/// size that was reserved.
105///
106/// *Note*: Pages are 64KiB in WASM.
107pub fn stable_grow(new_pages: u64) -> Result<u64, StableMemoryError> {
108    CANISTER_STABLE_MEMORY.stable_grow(new_pages)
109}
110
111/// Writes data to the stable memory location specified by an offset.
112///
113/// Warning - this will panic if `offset + buf.len()` exceeds the current size of stable memory.
114/// Use `stable_grow` to request more stable memory if needed.
115pub fn stable_write(offset: u64, buf: &[u8]) {
116    CANISTER_STABLE_MEMORY.stable_write(offset, buf)
117}
118
119/// Reads data from the stable memory location specified by an offset.
120pub fn stable_read(offset: u64, buf: &mut [u8]) {
121    CANISTER_STABLE_MEMORY.stable_read(offset, buf)
122}
123
124/// Returns a copy of the stable memory.
125///
126/// This will map the whole memory (even if not all of it has been written to).
127///
128/// # Panics
129///
130/// When the bytes of the stable memory cannot fit into a `Vec` which constrained by the usize.
131pub fn stable_bytes() -> Vec<u8> {
132    let size = (stable_size() << 16)
133        .try_into()
134        .expect("overflow: stable memory too large to read in one go");
135    let mut vec = Vec::with_capacity(size);
136    // SAFETY:
137    // `vec`, being mutable and allocated to `size` bytes, is safe to pass to ic0.stable_read with no offset.
138    // ic0.stable_read writes to all of `vec[0..size]`, so `set_len` is safe to call with the new size.
139    unsafe {
140        ic0::stable64_read(vec.as_ptr() as u64, 0, size as u64);
141        vec.set_len(size);
142    }
143    vec
144}
145
146/// Performs generic IO (read, write, and seek) on stable memory.
147///
148/// Warning: When using write functionality, this will overwrite any existing
149/// data in stable memory as it writes, so ensure you set the `offset` value
150/// accordingly if you wish to preserve existing data.
151///
152/// Will attempt to grow the memory as it writes,
153/// and keep offsets and total capacity.
154#[derive(Debug)]
155pub struct StableIO<M: StableMemory = CanisterStableMemory> {
156    /// The offset of the next write.
157    offset: u64,
158
159    /// The capacity, in pages.
160    capacity: u64,
161
162    /// The stable memory to write data to.
163    memory: M,
164}
165
166impl Default for StableIO {
167    fn default() -> Self {
168        Self::with_memory(CanisterStableMemory::default(), 0)
169    }
170}
171
172impl<M: StableMemory> StableIO<M> {
173    /// Creates a new `StableIO` which writes to the selected memory
174    pub fn with_memory(memory: M, offset: u64) -> Self {
175        let capacity = memory.stable_size();
176        Self {
177            offset,
178            capacity,
179            memory,
180        }
181    }
182
183    /// Returns the offset of the writer
184    pub fn offset(&self) -> u64 {
185        self.offset
186    }
187
188    /// Attempts to grow the memory by adding new pages.
189    pub fn grow(&mut self, new_pages: u64) -> Result<(), StableMemoryError> {
190        let old_page_count = self.memory.stable_grow(new_pages)?;
191        self.capacity = old_page_count + new_pages;
192        Ok(())
193    }
194
195    /// Writes a byte slice to the buffer.
196    ///
197    /// # Errors
198    ///
199    /// When it cannot grow the memory to accommodate the new data.
200    pub fn write(&mut self, buf: &[u8]) -> Result<usize, StableMemoryError> {
201        let required_capacity_bytes = self.offset + buf.len() as u64;
202        let required_capacity_pages = required_capacity_bytes.div_ceil(WASM_PAGE_SIZE_IN_BYTES);
203        let current_pages = self.capacity;
204        let additional_pages_required = required_capacity_pages.saturating_sub(current_pages);
205
206        if additional_pages_required > 0 {
207            self.grow(additional_pages_required)?;
208        }
209
210        self.memory.stable_write(self.offset, buf);
211        self.offset += buf.len() as u64;
212        Ok(buf.len())
213    }
214
215    /// Reads data from the stable memory location specified by an offset.
216    ///
217    /// # Errors
218    ///
219    /// The stable memory size is cached on creation of the StableReader.
220    /// Therefore, in following scenario, it will get an `OutOfBounds` error:
221    /// 1. Create a StableReader
222    /// 2. Write some data to the stable memory which causes it grow
223    /// 3. call `read()` to read the newly written bytes
224    pub fn read(&mut self, buf: &mut [u8]) -> Result<usize, StableMemoryError> {
225        let capacity_bytes = self.capacity * WASM_PAGE_SIZE_IN_BYTES;
226        let read_buf = if buf.len() as u64 + self.offset > capacity_bytes {
227            if self.offset < capacity_bytes {
228                // When usize=u32:
229                //   (capacity_bytes - self.offset) < buf.len() <= u32::MAX == usize::MAX.
230                // So the cast below won't panic.
231                &mut buf[..(capacity_bytes - self.offset).try_into().unwrap()]
232            } else {
233                return Err(StableMemoryError::OutOfBounds);
234            }
235        } else {
236            buf
237        };
238        self.memory.stable_read(self.offset, read_buf);
239        self.offset += read_buf.len() as u64;
240        Ok(read_buf.len())
241    }
242
243    // Helper used to implement io::Seek
244    fn seek(&mut self, offset: io::SeekFrom) -> io::Result<u64> {
245        self.offset = match offset {
246            io::SeekFrom::Start(offset) => offset,
247            io::SeekFrom::End(offset) => {
248                ((self.capacity * WASM_PAGE_SIZE_IN_BYTES) as i64 + offset) as u64
249            }
250            io::SeekFrom::Current(offset) => (self.offset as i64 + offset) as u64,
251        };
252
253        Ok(self.offset)
254    }
255}
256
257impl<M: StableMemory> io::Write for StableIO<M> {
258    fn write(&mut self, buf: &[u8]) -> Result<usize, io::Error> {
259        self.write(buf)
260            .map_err(|e| io::Error::new(io::ErrorKind::OutOfMemory, e))
261    }
262
263    fn flush(&mut self) -> Result<(), io::Error> {
264        // Noop.
265        Ok(())
266    }
267}
268
269impl<M: StableMemory> io::Read for StableIO<M> {
270    fn read(&mut self, buf: &mut [u8]) -> Result<usize, io::Error> {
271        Self::read(self, buf).or(Ok(0)) // Read defines EOF to be success
272    }
273}
274
275impl<M: StableMemory> io::Seek for StableIO<M> {
276    fn seek(&mut self, offset: io::SeekFrom) -> io::Result<u64> {
277        self.seek(offset)
278    }
279}
280
281// impl_stable_io!(u32);
282// impl_stable_io!(u64);
283
284/// A writer to the stable memory.
285///
286/// Warning: This will overwrite any existing data in stable memory as it writes, so ensure you set
287/// the `offset` value accordingly if you wish to preserve existing data.
288///
289/// Will attempt to grow the memory as it writes,
290/// and keep offsets and total capacity.
291#[derive(Debug)]
292pub struct StableWriter<M: StableMemory = CanisterStableMemory>(StableIO<M>);
293
294#[allow(clippy::derivable_impls)]
295impl Default for StableWriter {
296    #[inline]
297    fn default() -> Self {
298        Self(StableIO::default())
299    }
300}
301
302impl<M: StableMemory> StableWriter<M> {
303    /// Creates a new `StableWriter` which writes to the selected memory
304    #[inline]
305    pub fn with_memory(memory: M, offset: u64) -> Self {
306        Self(StableIO::<M>::with_memory(memory, offset))
307    }
308
309    /// Returns the offset of the writer
310    #[inline]
311    pub fn offset(&self) -> u64 {
312        self.0.offset()
313    }
314
315    /// Attempts to grow the memory by adding new pages.
316    #[inline]
317    pub fn grow(&mut self, new_pages: u64) -> Result<(), StableMemoryError> {
318        self.0.grow(new_pages)
319    }
320
321    /// Writes a byte slice to the buffer.
322    ///
323    /// The only condition where this will
324    /// error out is if it cannot grow the memory.
325    #[inline]
326    pub fn write(&mut self, buf: &[u8]) -> Result<usize, StableMemoryError> {
327        self.0.write(buf)
328    }
329}
330
331impl<M: StableMemory> io::Write for StableWriter<M> {
332    #[inline]
333    fn write(&mut self, buf: &[u8]) -> Result<usize, io::Error> {
334        io::Write::write(&mut self.0, buf)
335    }
336
337    #[inline]
338    fn flush(&mut self) -> Result<(), io::Error> {
339        io::Write::flush(&mut self.0)
340    }
341}
342
343impl<M: StableMemory> io::Seek for StableWriter<M> {
344    #[inline]
345    fn seek(&mut self, pos: io::SeekFrom) -> io::Result<u64> {
346        io::Seek::seek(&mut self.0, pos)
347    }
348}
349
350impl<M: StableMemory> From<StableIO<M>> for StableWriter<M> {
351    fn from(io: StableIO<M>) -> Self {
352        Self(io)
353    }
354}
355
356/// A writer to the stable memory which first writes the bytes to an in memory buffer and flushes
357/// the buffer to stable memory each time it becomes full.
358///
359/// Warning: This will overwrite any existing data in stable memory as it writes, so ensure you set
360/// the `offset` value accordingly if you wish to preserve existing data.
361///
362/// Note: Each call to grow or write to stable memory is a relatively expensive operation, so pick a
363/// buffer size large enough to avoid excessive calls to stable memory.
364#[derive(Debug)]
365pub struct BufferedStableWriter<M: StableMemory = CanisterStableMemory> {
366    inner: io::BufWriter<StableWriter<M>>,
367}
368
369impl BufferedStableWriter {
370    /// Creates a new `BufferedStableWriter`
371    pub fn new(buffer_size: usize) -> BufferedStableWriter {
372        BufferedStableWriter::with_writer(buffer_size, StableWriter::default())
373    }
374}
375
376impl<M: StableMemory> BufferedStableWriter<M> {
377    /// Creates a new `BufferedStableWriter` which writes to the selected memory
378    pub fn with_writer(buffer_size: usize, writer: StableWriter<M>) -> BufferedStableWriter<M> {
379        BufferedStableWriter {
380            inner: io::BufWriter::with_capacity(buffer_size, writer),
381        }
382    }
383
384    /// Returns the offset of the writer
385    pub fn offset(&self) -> u64 {
386        self.inner.get_ref().offset()
387    }
388}
389
390impl<M: StableMemory> io::Write for BufferedStableWriter<M> {
391    fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
392        self.inner.write(buf)
393    }
394
395    fn flush(&mut self) -> io::Result<()> {
396        self.inner.flush()
397    }
398}
399
400impl<M: StableMemory> io::Seek for BufferedStableWriter<M> {
401    #[inline]
402    fn seek(&mut self, pos: io::SeekFrom) -> io::Result<u64> {
403        io::Seek::seek(&mut self.inner, pos)
404    }
405}
406
407// A reader to the stable memory.
408///
409/// Keeps an offset and reads off stable memory consecutively.
410#[derive(Debug)]
411pub struct StableReader<M: StableMemory = CanisterStableMemory>(StableIO<M>);
412
413#[allow(clippy::derivable_impls)]
414impl Default for StableReader {
415    fn default() -> Self {
416        Self(StableIO::default())
417    }
418}
419
420impl<M: StableMemory> StableReader<M> {
421    /// Creates a new `StableReader` which reads from the selected memory
422    #[inline]
423    pub fn with_memory(memory: M, offset: u64) -> Self {
424        Self(StableIO::<M>::with_memory(memory, offset))
425    }
426
427    /// Returns the offset of the reader
428    #[inline]
429    pub fn offset(&self) -> u64 {
430        self.0.offset()
431    }
432
433    /// Reads data from the stable memory location specified by an offset.
434    ///
435    /// Note:
436    /// The stable memory size is cached on creation of the StableReader.
437    /// Therefore, in following scenario, it will get an `OutOfBounds` error:
438    /// 1. Create a StableReader
439    /// 2. Write some data to the stable memory which causes it grow
440    /// 3. call `read()` to read the newly written bytes
441    #[inline]
442    pub fn read(&mut self, buf: &mut [u8]) -> Result<usize, StableMemoryError> {
443        self.0.read(buf)
444    }
445}
446
447impl<M: StableMemory> io::Read for StableReader<M> {
448    #[inline]
449    fn read(&mut self, buf: &mut [u8]) -> Result<usize, io::Error> {
450        io::Read::read(&mut self.0, buf)
451    }
452}
453
454impl<M: StableMemory> io::Seek for StableReader<M> {
455    #[inline]
456    fn seek(&mut self, pos: io::SeekFrom) -> io::Result<u64> {
457        io::Seek::seek(&mut self.0, pos)
458    }
459}
460
461impl<M: StableMemory> From<StableIO<M>> for StableReader<M> {
462    fn from(io: StableIO<M>) -> Self {
463        Self(io)
464    }
465}
466
467/// A reader to the stable memory which reads bytes a chunk at a time as each chunk is required.
468#[derive(Debug)]
469pub struct BufferedStableReader<M: StableMemory = CanisterStableMemory> {
470    inner: io::BufReader<StableReader<M>>,
471}
472
473impl BufferedStableReader {
474    /// Creates a new `BufferedStableReader`
475    pub fn new(buffer_size: usize) -> BufferedStableReader {
476        BufferedStableReader::with_reader(buffer_size, StableReader::default())
477    }
478}
479
480impl<M: StableMemory> BufferedStableReader<M> {
481    /// Creates a new `BufferedStableReader` which reads from the selected memory
482    pub fn with_reader(buffer_size: usize, reader: StableReader<M>) -> BufferedStableReader<M> {
483        BufferedStableReader {
484            inner: io::BufReader::with_capacity(buffer_size, reader),
485        }
486    }
487
488    /// Returns the offset of the reader
489    pub fn offset(&self) -> u64 {
490        self.inner.get_ref().offset()
491    }
492}
493
494impl<M: StableMemory> io::Read for BufferedStableReader<M> {
495    fn read(&mut self, buf: &mut [u8]) -> io::Result<usize> {
496        self.inner.read(buf)
497    }
498}
499
500impl<M: StableMemory> io::Seek for BufferedStableReader<M> {
501    #[inline]
502    fn seek(&mut self, pos: io::SeekFrom) -> io::Result<u64> {
503        io::Seek::seek(&mut self.inner, pos)
504    }
505}
506
507#[cfg(test)]
508mod tests {
509    use super::*;
510    use std::rc::Rc;
511    use std::sync::Mutex;
512
513    #[derive(Default)]
514    pub struct TestStableMemory {
515        memory: Rc<Mutex<Vec<u8>>>,
516    }
517
518    impl TestStableMemory {
519        pub fn new(memory: Rc<Mutex<Vec<u8>>>) -> TestStableMemory {
520            let bytes_len = memory.lock().unwrap().len();
521            if bytes_len > 0 {
522                let pages_required = pages_required(bytes_len);
523                let bytes_required = pages_required * WASM_PAGE_SIZE_IN_BYTES;
524                memory
525                    .lock()
526                    .unwrap()
527                    .resize(bytes_required.try_into().unwrap(), 0);
528            }
529
530            TestStableMemory { memory }
531        }
532    }
533
534    impl StableMemory for TestStableMemory {
535        fn stable_size(&self) -> u64 {
536            let bytes_len = self.memory.lock().unwrap().len();
537            pages_required(bytes_len)
538        }
539
540        fn stable_grow(&self, new_pages: u64) -> Result<u64, StableMemoryError> {
541            let new_bytes = new_pages * WASM_PAGE_SIZE_IN_BYTES;
542
543            let mut vec = self.memory.lock().unwrap();
544            let previous_len = vec.len() as u64;
545            let new_len = vec.len() as u64 + new_bytes;
546            vec.resize(new_len.try_into().unwrap(), 0);
547            Ok(previous_len / WASM_PAGE_SIZE_IN_BYTES)
548        }
549
550        fn stable_write(&self, offset: u64, buf: &[u8]) {
551            let offset = offset as usize;
552
553            let mut vec = self.memory.lock().unwrap();
554            if offset + buf.len() > vec.len() {
555                panic!("stable memory out of bounds");
556            }
557            vec[offset..(offset + buf.len())].clone_from_slice(buf);
558        }
559
560        fn stable_read(&self, offset: u64, buf: &mut [u8]) {
561            let offset = offset as usize;
562
563            let vec = self.memory.lock().unwrap();
564            let count_to_copy = buf.len();
565
566            buf[..count_to_copy].copy_from_slice(&vec[offset..offset + count_to_copy]);
567        }
568    }
569
570    fn pages_required(bytes_len: usize) -> u64 {
571        let page_size = WASM_PAGE_SIZE_IN_BYTES;
572        (bytes_len as u64).div_ceil(page_size)
573    }
574
575    mod stable_writer_tests {
576        use super::*;
577        use rstest::rstest;
578        use std::io::{Seek, Write};
579
580        #[rstest]
581        #[case(None)]
582        #[case(Some(1))]
583        #[case(Some(10))]
584        #[case(Some(100))]
585        #[case(Some(1000))]
586        fn write_single_slice(#[case] buffer_size: Option<usize>) {
587            let memory = Rc::new(Mutex::new(Vec::new()));
588            let mut writer = build_writer(TestStableMemory::new(memory.clone()), buffer_size);
589
590            let bytes = vec![1; 100];
591
592            writer.write_all(&bytes).unwrap();
593            writer.flush().unwrap();
594
595            let result = &*memory.lock().unwrap();
596
597            assert_eq!(bytes, result[..bytes.len()]);
598        }
599
600        #[rstest]
601        #[case(None)]
602        #[case(Some(1))]
603        #[case(Some(10))]
604        #[case(Some(100))]
605        #[case(Some(1000))]
606        fn write_many_slices(#[case] buffer_size: Option<usize>) {
607            let memory = Rc::new(Mutex::new(Vec::new()));
608            let mut writer = build_writer(TestStableMemory::new(memory.clone()), buffer_size);
609
610            for i in 1..100 {
611                let bytes = vec![i as u8; i];
612                writer.write_all(&bytes).unwrap();
613            }
614            writer.flush().unwrap();
615
616            let result = &*memory.lock().unwrap();
617
618            let mut offset = 0;
619            for i in 1..100 {
620                let bytes = &result[offset..offset + i];
621                assert_eq!(bytes, vec![i as u8; i]);
622                offset += i;
623            }
624        }
625
626        #[rstest]
627        #[case(None)]
628        #[case(Some(1))]
629        #[case(Some(10))]
630        #[case(Some(100))]
631        #[case(Some(1000))]
632        fn ensure_only_requests_min_number_of_pages_required(#[case] buffer_size: Option<usize>) {
633            let memory = Rc::new(Mutex::new(Vec::new()));
634            let mut writer = build_writer(TestStableMemory::new(memory.clone()), buffer_size);
635
636            let mut total_bytes = 0;
637            for i in 1..10000 {
638                let bytes = vec![i as u8; i];
639                writer.write_all(&bytes).unwrap();
640                total_bytes += i;
641            }
642            writer.flush().unwrap();
643
644            let capacity_pages = TestStableMemory::new(memory).stable_size();
645            let min_pages_required = (total_bytes as u64).div_ceil(WASM_PAGE_SIZE_IN_BYTES);
646
647            assert_eq!(capacity_pages, min_pages_required);
648        }
649
650        #[test]
651        fn check_offset() {
652            const WRITE_SIZE: usize = 1025;
653
654            let memory = Rc::new(Mutex::new(Vec::new()));
655            let mut writer = StableWriter::with_memory(TestStableMemory::new(memory.clone()), 0);
656            assert_eq!(writer.offset(), 0);
657            assert_eq!(writer.write(&vec![0; WRITE_SIZE]).unwrap(), WRITE_SIZE);
658            assert_eq!(writer.offset(), WRITE_SIZE as u64);
659
660            let mut writer = BufferedStableWriter::with_writer(
661                WRITE_SIZE - 1,
662                StableWriter::with_memory(TestStableMemory::new(memory), 0),
663            );
664            assert_eq!(writer.offset(), 0);
665            assert_eq!(writer.write(&vec![0; WRITE_SIZE]).unwrap(), WRITE_SIZE);
666            assert_eq!(writer.offset(), WRITE_SIZE as u64);
667        }
668
669        #[test]
670        fn test_seek() {
671            let memory = Rc::new(Mutex::new(Vec::new()));
672            let mut writer = StableWriter::with_memory(TestStableMemory::new(memory.clone()), 0);
673            writer
674                .seek(std::io::SeekFrom::Start(WASM_PAGE_SIZE_IN_BYTES))
675                .unwrap();
676            assert_eq!(writer.stream_position().unwrap(), WASM_PAGE_SIZE_IN_BYTES);
677            assert_eq!(writer.write(&[1_u8]).unwrap(), 1);
678            assert_eq!(
679                writer.seek(std::io::SeekFrom::End(0)).unwrap(),
680                WASM_PAGE_SIZE_IN_BYTES * 2
681            );
682            let capacity_pages = TestStableMemory::new(memory).stable_size();
683            assert_eq!(capacity_pages, 2);
684        }
685
686        fn build_writer(memory: TestStableMemory, buffer_size: Option<usize>) -> Box<dyn Write> {
687            let writer = StableWriter::with_memory(memory, 0);
688            if let Some(buffer_size) = buffer_size {
689                Box::new(BufferedStableWriter::with_writer(buffer_size, writer))
690            } else {
691                Box::new(writer)
692            }
693        }
694    }
695
696    mod stable_reader_tests {
697        use super::*;
698        use rstest::rstest;
699        use std::io::{Read, Seek};
700
701        #[rstest]
702        #[case(None)]
703        #[case(Some(1))]
704        #[case(Some(10))]
705        #[case(Some(100))]
706        #[case(Some(1000))]
707        fn reads_all_bytes(#[case] buffer_size: Option<usize>) {
708            let input = vec![1; 10_000];
709            let memory = Rc::new(Mutex::new(input.clone()));
710            let mut reader = build_reader(TestStableMemory::new(memory), buffer_size);
711
712            let mut output = Vec::new();
713            reader.read_to_end(&mut output).unwrap();
714
715            assert_eq!(input, output[..input.len()]);
716        }
717
718        #[test]
719        fn check_offset() {
720            const READ_SIZE: usize = 1025;
721
722            let memory = Rc::new(Mutex::new(vec![1; READ_SIZE]));
723            let mut reader = StableReader::with_memory(TestStableMemory::new(memory.clone()), 0);
724            assert_eq!(reader.offset(), 0);
725            let mut bytes = vec![0; READ_SIZE];
726            assert_eq!(reader.read(&mut bytes).unwrap(), READ_SIZE);
727            assert_eq!(reader.offset(), READ_SIZE as u64);
728
729            let mut reader = BufferedStableReader::with_reader(
730                READ_SIZE - 1,
731                StableReader::with_memory(TestStableMemory::new(memory), 0),
732            );
733            assert_eq!(reader.offset(), 0);
734            let mut bytes = vec![0; READ_SIZE];
735            assert_eq!(reader.read(&mut bytes).unwrap(), READ_SIZE);
736            assert_eq!(reader.offset(), READ_SIZE as u64);
737        }
738
739        #[test]
740        fn test_seek() {
741            const SIZE: usize = 1025;
742            let memory = Rc::new(Mutex::new((0..SIZE).map(|v| v as u8).collect::<Vec<u8>>()));
743            let mut reader = StableReader::with_memory(TestStableMemory::new(memory), 0);
744            let mut bytes = vec![0_u8; 1];
745
746            const OFFSET: usize = 200;
747            reader
748                .seek(std::io::SeekFrom::Start(OFFSET as u64))
749                .unwrap();
750            assert_eq!(reader.stream_position().unwrap() as usize, OFFSET);
751            assert_eq!(reader.read(&mut bytes).unwrap(), 1);
752            assert_eq!(&bytes, &[OFFSET as u8]);
753            assert_eq!(
754                reader.seek(std::io::SeekFrom::End(0)).unwrap(),
755                WASM_PAGE_SIZE_IN_BYTES
756            );
757            reader
758                .seek(std::io::SeekFrom::Start(WASM_PAGE_SIZE_IN_BYTES * 2))
759                .unwrap();
760            // out of bounds so should fail
761            assert!(reader.read(&mut bytes).is_err());
762        }
763
764        fn build_reader(memory: TestStableMemory, buffer_size: Option<usize>) -> Box<dyn Read> {
765            let reader = StableReader::with_memory(memory, 0);
766            if let Some(buffer_size) = buffer_size {
767                Box::new(BufferedStableReader::with_reader(buffer_size, reader))
768            } else {
769                Box::new(reader)
770            }
771        }
772    }
773}