1use std::{error, fmt, io};
10
11pub const WASM_PAGE_SIZE_IN_BYTES: u64 = 64 * 1024; static CANISTER_STABLE_MEMORY: CanisterStableMemory = CanisterStableMemory {};
15
16pub trait StableMemory {
18 fn stable_size(&self) -> u64;
20
21 fn stable_grow(&self, new_pages: u64) -> Result<u64, StableMemoryError>;
28
29 fn stable_write(&self, offset: u64, buf: &[u8]);
34
35 fn stable_read(&self, offset: u64, buf: &mut [u8]);
37}
38
39#[derive(Debug)]
41pub enum StableMemoryError {
42 OutOfMemory,
44 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#[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
86pub fn stable_size() -> u64 {
88 CANISTER_STABLE_MEMORY.stable_size()
89}
90
91pub fn stable_grow(new_pages: u64) -> Result<u64, StableMemoryError> {
98 CANISTER_STABLE_MEMORY.stable_grow(new_pages)
99}
100
101pub fn stable_write(offset: u64, buf: &[u8]) {
106 CANISTER_STABLE_MEMORY.stable_write(offset, buf);
107}
108
109pub fn stable_read(offset: u64, buf: &mut [u8]) {
111 CANISTER_STABLE_MEMORY.stable_read(offset, buf);
112}
113
114pub 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 unsafe {
129 vec.set_len(size);
130 }
131 vec
132}
133
134#[derive(Debug)]
143pub struct StableIO<M: StableMemory = CanisterStableMemory> {
144 offset: u64,
146
147 capacity: u64,
149
150 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 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 pub fn offset(&self) -> u64 {
173 self.offset
174 }
175
176 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 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 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 &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 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 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)) }
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#[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 #[inline]
293 pub fn with_memory(memory: M, offset: u64) -> Self {
294 Self(StableIO::<M>::with_memory(memory, offset))
295 }
296
297 #[inline]
299 pub fn offset(&self) -> u64 {
300 self.0.offset()
301 }
302
303 #[inline]
305 pub fn grow(&mut self, new_pages: u64) -> Result<(), StableMemoryError> {
306 self.0.grow(new_pages)
307 }
308
309 #[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#[derive(Debug)]
353pub struct BufferedStableWriter<M: StableMemory = CanisterStableMemory> {
354 inner: io::BufWriter<StableWriter<M>>,
355}
356
357impl BufferedStableWriter {
358 pub fn new(buffer_size: usize) -> BufferedStableWriter {
360 BufferedStableWriter::with_writer(buffer_size, StableWriter::default())
361 }
362}
363
364impl<M: StableMemory> BufferedStableWriter<M> {
365 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 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#[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 #[inline]
411 pub fn with_memory(memory: M, offset: u64) -> Self {
412 Self(StableIO::<M>::with_memory(memory, offset))
413 }
414
415 #[inline]
417 pub fn offset(&self) -> u64 {
418 self.0.offset()
419 }
420
421 #[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#[derive(Debug)]
457pub struct BufferedStableReader<M: StableMemory = CanisterStableMemory> {
458 inner: io::BufReader<StableReader<M>>,
459}
460
461impl BufferedStableReader {
462 pub fn new(buffer_size: usize) -> BufferedStableReader {
464 BufferedStableReader::with_reader(buffer_size, StableReader::default())
465 }
466}
467
468impl<M: StableMemory> BufferedStableReader<M> {
469 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 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 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}