b3_stable_structures/
lib.rs

1#![doc = include_str!("../README.md")]
2mod base_vec;
3pub mod btreemap;
4pub mod cell;
5pub use cell::{Cell as StableCell, Cell};
6pub mod file_mem;
7#[cfg(target_arch = "wasm32")]
8mod ic0_memory; // Memory API for canisters.
9pub mod log;
10pub use log::{Log as StableLog, Log};
11pub mod memory_manager;
12pub mod min_heap;
13pub mod reader;
14pub mod storable;
15#[cfg(test)]
16mod tests;
17mod types;
18pub mod vec;
19pub use min_heap::{MinHeap, MinHeap as StableMinHeap};
20pub use vec::{Vec as StableVec, Vec};
21pub mod vec_mem;
22pub mod writer;
23pub use btreemap::{BTreeMap, BTreeMap as StableBTreeMap};
24pub use file_mem::FileMemory;
25#[cfg(target_arch = "wasm32")]
26pub use ic0_memory::Ic0StableMemory;
27use std::error;
28use std::fmt::{Display, Formatter};
29pub use storable::{BoundedStorable, Storable};
30use types::Address;
31pub use vec_mem::VectorMemory;
32
33#[cfg(target_arch = "wasm32")]
34pub type DefaultMemoryImpl = Ic0StableMemory;
35
36#[cfg(not(target_arch = "wasm32"))]
37pub type DefaultMemoryImpl = VectorMemory;
38
39const WASM_PAGE_SIZE: u64 = 65536;
40
41/// The maximum number of stable memory pages a canister can address.
42pub const MAX_PAGES: u64 = u64::MAX / WASM_PAGE_SIZE;
43
44pub trait Memory {
45    /// Returns the current size of the stable memory in WebAssembly
46    /// pages. (One WebAssembly page is 64Ki bytes.)
47    fn size(&self) -> u64;
48
49    /// Tries to grow the memory by new_pages many pages containing
50    /// zeroes.  If successful, returns the previous size of the
51    /// memory (in pages).  Otherwise, returns -1.
52    fn grow(&self, pages: u64) -> i64;
53
54    /// Copies the data referred to by offset out of the stable memory
55    /// and replaces the corresponding bytes in dst.
56    fn read(&self, offset: u64, dst: &mut [u8]);
57
58    /// Copies the data referred to by src and replaces the
59    /// corresponding segment starting at offset in the stable memory.
60    fn write(&self, offset: u64, src: &[u8]);
61}
62
63// A helper function that reads a single 32bit integer encoded as
64// little-endian from the specified memory at the specified offset.
65fn read_u32<M: Memory>(m: &M, addr: Address) -> u32 {
66    let mut buf: [u8; 4] = [0; 4];
67    m.read(addr.get(), &mut buf);
68    u32::from_le_bytes(buf)
69}
70
71// A helper function that reads a single 64bit integer encoded as
72// little-endian from the specified memory at the specified offset.
73fn read_u64<M: Memory>(m: &M, addr: Address) -> u64 {
74    let mut buf: [u8; 8] = [0; 8];
75    m.read(addr.get(), &mut buf);
76    u64::from_le_bytes(buf)
77}
78
79// Writes a single 32-bit integer encoded as little-endian.
80fn write_u32<M: Memory>(m: &M, addr: Address, val: u32) {
81    write(m, addr.get(), &val.to_le_bytes());
82}
83
84// Writes a single 64-bit integer encoded as little-endian.
85fn write_u64<M: Memory>(m: &M, addr: Address, val: u64) {
86    write(m, addr.get(), &val.to_le_bytes());
87}
88
89#[derive(Debug, PartialEq, Eq)]
90pub struct GrowFailed {
91    current_size: u64,
92    delta: u64,
93}
94
95impl Display for GrowFailed {
96    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
97        write!(
98            f,
99            "Failed to grow memory: current size={}, delta={}",
100            self.current_size,
101            self.delta
102        )
103    }
104}
105
106impl error::Error for GrowFailed {}
107
108/// Writes the bytes at the specified offset, growing the memory size if needed.
109fn safe_write<M: Memory>(memory: &M, offset: u64, bytes: &[u8]) -> Result<(), GrowFailed> {
110    let last_byte = offset
111        .checked_add(bytes.len() as u64)
112        .expect("Address space overflow");
113
114    let size_pages = memory.size();
115    let size_bytes = size_pages
116        .checked_mul(WASM_PAGE_SIZE)
117        .expect("Address space overflow");
118
119    if size_bytes < last_byte {
120        let diff_bytes = last_byte - size_bytes;
121        let diff_pages = diff_bytes
122            .checked_add(WASM_PAGE_SIZE - 1)
123            .expect("Address space overflow")
124            / WASM_PAGE_SIZE;
125        if memory.grow(diff_pages) == -1 {
126            return Err(GrowFailed {
127                current_size: size_pages,
128                delta: diff_pages,
129            });
130        }
131    }
132    memory.write(offset, bytes);
133    Ok(())
134}
135
136/// Like [safe_write], but panics if the memory.grow fails.
137fn write<M: Memory>(memory: &M, offset: u64, bytes: &[u8]) {
138    if let Err(GrowFailed {
139        current_size,
140        delta,
141    }) = safe_write(memory, offset, bytes)
142    {
143        panic!(
144            "Failed to grow memory from {} pages to {} pages (delta = {} pages).",
145            current_size,
146            current_size + delta,
147            delta
148        );
149    }
150}
151
152// Reads a struct from memory.
153fn read_struct<T, M: Memory>(addr: Address, memory: &M) -> T {
154    let mut t: T = unsafe { core::mem::zeroed() };
155    let t_slice = unsafe {
156        core::slice::from_raw_parts_mut(&mut t as *mut _ as *mut u8, core::mem::size_of::<T>())
157    };
158    memory.read(addr.get(), t_slice);
159    t
160}
161
162// Writes a struct to memory.
163fn write_struct<T, M: Memory>(t: &T, addr: Address, memory: &M) {
164    let slice = unsafe {
165        core::slice::from_raw_parts(t as *const _ as *const u8, core::mem::size_of::<T>())
166    };
167
168    write(memory, addr.get(), slice)
169}
170
171/// RestrictedMemory creates a limited view of another memory.  This
172/// allows one to divide the main memory into non-intersecting ranges
173/// and use different layouts in each region.
174#[derive(Clone)]
175pub struct RestrictedMemory<M: Memory> {
176    page_range: core::ops::Range<u64>,
177    memory: M,
178}
179
180impl<M: Memory> RestrictedMemory<M> {
181    pub fn new(memory: M, page_range: core::ops::Range<u64>) -> Self {
182        assert!(page_range.end <= MAX_PAGES);
183        Self { memory, page_range }
184    }
185}
186
187impl<M: Memory> Memory for RestrictedMemory<M> {
188    fn size(&self) -> u64 {
189        let base_size = self.memory.size();
190        if base_size < self.page_range.start {
191            0
192        } else if base_size > self.page_range.end {
193            self.page_range.end - self.page_range.start
194        } else {
195            base_size - self.page_range.start
196        }
197    }
198
199    fn grow(&self, delta: u64) -> i64 {
200        let base_size = self.memory.size();
201        if base_size < self.page_range.start {
202            self.memory
203                .grow(self.page_range.start - base_size + delta)
204                .min(0)
205        } else if base_size >= self.page_range.end {
206            if delta == 0 {
207                (self.page_range.end - self.page_range.start) as i64
208            } else {
209                -1
210            }
211        } else {
212            let pages_left = self.page_range.end - base_size;
213            if pages_left < delta {
214                -1
215            } else {
216                let r = self.memory.grow(delta);
217                if r < 0 {
218                    r
219                } else {
220                    r - self.page_range.start as i64
221                }
222            }
223        }
224    }
225
226    fn read(&self, offset: u64, dst: &mut [u8]) {
227        self.memory
228            .read(self.page_range.start * WASM_PAGE_SIZE + offset, dst)
229    }
230
231    fn write(&self, offset: u64, src: &[u8]) {
232        self.memory
233            .write(self.page_range.start * WASM_PAGE_SIZE + offset, src)
234    }
235}