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