ic_stable_structures/
base_vec.rs

1//! The common code for all Vec-like data structures.
2//!
3//! # V1 layout
4//!
5//! ```text
6//! ---------------------------------------- <- Address 0
7//! Magic                  ↕ 3 bytes
8//! ----------------------------------------
9//! Layout version         ↕ 1 byte
10//! ----------------------------------------
11//! Number of entries = L  ↕ 8 bytes
12//! ----------------------------------------
13//! Max entry size         ↕ 4 bytes
14//! ----------------------------------------
15//! Fixed size flag        ↕ 1 byte
16//! ----------------------------------------
17//! Reserved space         ↕ 47 bytes
18//! ---------------------------------------- <- Address 64
19//! E_0                    ↕ SLOT_SIZE bytes
20//! ----------------------------------------
21//! E_1                    ↕ SLOT_SIZE bytes
22//! ----------------------------------------
23//! ...
24//! ----------------------------------------
25//! E_(L-1)                ↕ SLOT_SIZE bytes
26//! ----------------------------------------
27//! Unallocated space
28//! ```
29//!
30//! The `SLOT_SIZE` constant depends on the item type. If the item
31//! type is fixed in size, the `SLOT_SIZE` is equal to the max size.
32//! Otherwise, the `SLOT_SIZE` is the max size plus the number of
33//! bytes required to represent integers up to that max size.
34use crate::storable::{bounds, bytes_to_store_size_bounded};
35use crate::{
36    read_to_vec, read_u32, read_u64, safe_write, write, write_u32, write_u64, Address, GrowFailed,
37    Memory, Storable,
38};
39use std::borrow::{Borrow, Cow};
40use std::cmp::min;
41use std::fmt;
42use std::marker::PhantomData;
43use std::ops::Range;
44
45const LAYOUT_VERSION: u8 = 1;
46/// The offset where the user data begins.
47const DATA_OFFSET: u64 = 64;
48/// The offset where the vector length resides.
49const LEN_OFFSET: u64 = 4;
50
51#[derive(Debug)]
52struct HeaderV1 {
53    magic: [u8; 3],
54    version: u8,
55    len: u64,
56    max_size: u32,
57    is_fixed_size: bool,
58}
59
60#[derive(PartialEq, Eq, Debug)]
61pub enum InitError {
62    /// The memory already contains another data structure.
63    /// Use [Vec::new] to overwrite it.
64    BadMagic { actual: [u8; 3], expected: [u8; 3] },
65    /// The current version of [Vec] does not support the of the
66    /// memory layout.
67    IncompatibleVersion(u8),
68    /// The vector type is not compatible with the current vector
69    /// layout: the type's bounds differ from the original initialization
70    /// parameters.
71    IncompatibleElementType,
72    /// Failed to allocate memory for the vector.
73    OutOfMemory,
74}
75
76impl fmt::Display for InitError {
77    fn fmt(&self, fmt: &mut fmt::Formatter<'_>) -> fmt::Result {
78        match self {
79            Self::BadMagic{ actual, expected } => {
80                write!(fmt, "bad magic number {actual:?}, expected {expected:?}")
81            }
82            Self::IncompatibleVersion(version)
83            => write!(
84                fmt,
85                "unsupported layout version {version}; supported version numbers are 1..={LAYOUT_VERSION}"
86            ),
87            Self::IncompatibleElementType =>
88                write!(fmt, "the bounds (either max_size or is_fixed_size) of the element type do not match the persisted vector attributes"),
89            Self::OutOfMemory => write!(fmt, "failed to allocate memory for vector metadata"),
90        }
91    }
92}
93
94impl std::error::Error for InitError {}
95
96pub struct BaseVec<T: Storable, M: Memory> {
97    memory: M,
98    _marker: PhantomData<T>,
99}
100
101impl<T: Storable, M: Memory> BaseVec<T, M> {
102    /// Creates a new empty vector in the specified memory,
103    /// overwriting any data structures the memory might have
104    /// contained previously.
105    ///
106    /// Complexity: O(1)
107    pub fn new(memory: M, magic: [u8; 3]) -> Result<Self, GrowFailed> {
108        let t_bounds = bounds::<T>();
109
110        let header = HeaderV1 {
111            magic,
112            version: LAYOUT_VERSION,
113            len: 0,
114            max_size: t_bounds.max_size,
115            is_fixed_size: t_bounds.is_fixed_size,
116        };
117        Self::write_header(&header, &memory)?;
118        Ok(Self {
119            memory,
120            _marker: PhantomData,
121        })
122    }
123
124    /// Initializes a vector in the specified memory.
125    ///
126    /// Complexity: O(1)
127    ///
128    /// PRECONDITION: the memory is either empty or contains a valid
129    /// stable vector.
130    pub fn init(memory: M, magic: [u8; 3]) -> Result<Self, InitError> {
131        if memory.size() == 0 {
132            return Self::new(memory, magic).map_err(|_| InitError::OutOfMemory);
133        }
134        let header = Self::read_header(&memory);
135        if header.magic != magic {
136            return Err(InitError::BadMagic {
137                actual: header.magic,
138                expected: magic,
139            });
140        }
141        if header.version != LAYOUT_VERSION {
142            return Err(InitError::IncompatibleVersion(header.version));
143        }
144        let t_bounds = bounds::<T>();
145        if header.max_size != t_bounds.max_size || header.is_fixed_size != t_bounds.is_fixed_size {
146            return Err(InitError::IncompatibleElementType);
147        }
148
149        Ok(Self {
150            memory,
151            _marker: PhantomData,
152        })
153    }
154
155    /// Returns the underlying memory instance.
156    pub fn into_memory(self) -> M {
157        self.memory
158    }
159
160    /// Returns true if the vector is empty.
161    ///
162    /// Complexity: O(1)
163    pub fn is_empty(&self) -> bool {
164        self.len() == 0
165    }
166
167    /// Returns the number of items in the vector.
168    ///
169    /// Complexity: O(1)
170    pub fn len(&self) -> u64 {
171        read_u64(&self.memory, Address::from(LEN_OFFSET))
172    }
173
174    /// Sets the item at the specified index to the specified value.
175    ///
176    /// Complexity: O(max_size(T))
177    ///
178    /// PRECONDITION: index < self.len()
179    pub fn set(&self, index: u64, item: &T) {
180        assert!(index < self.len());
181
182        let offset = DATA_OFFSET + slot_size::<T>() as u64 * index;
183        let bytes = item.to_bytes_checked();
184        let data_offset = self
185            .write_entry_size(offset, bytes.len() as u32)
186            .expect("unreachable: cannot fail to write to pre-allocated area");
187        write(&self.memory, data_offset, bytes.borrow());
188    }
189
190    /// Returns the item at the specified index.
191    ///
192    /// Complexity: O(max_size(T))
193    pub fn get(&self, index: u64) -> Option<T> {
194        if index < self.len() {
195            Some(self.read_entry(index))
196        } else {
197            None
198        }
199    }
200
201    /// Adds a new item at the end of the vector.
202    ///
203    /// Complexity: O(max_size(T))
204    pub fn push(&self, item: &T) -> Result<(), GrowFailed> {
205        let index = self.len();
206        let offset = DATA_OFFSET + slot_size::<T>() as u64 * index;
207        let bytes = item.to_bytes_checked();
208        let data_offset = self.write_entry_size(offset, bytes.len() as u32)?;
209        safe_write(&self.memory, data_offset, bytes.borrow())?;
210        // NB. We update the size only after we ensure that the data
211        // write succeeded.
212        self.set_len(index + 1);
213        Ok(())
214    }
215
216    /// Removes the item at the end of the vector.
217    ///
218    /// Complexity: O(max_size(T))
219    pub fn pop(&self) -> Option<T> {
220        let len = self.len();
221        if len == 0 {
222            return None;
223        }
224        let value = self.read_entry(len - 1);
225        self.set_len(len - 1);
226        Some(value)
227    }
228
229    pub fn iter(&self) -> Iter<'_, T, M> {
230        Iter {
231            vec: self,
232            buf: vec![],
233            range: Range {
234                start: 0,
235                end: self.len(),
236            },
237        }
238    }
239
240    /// Reads the item at the specified index without any bound checks.
241    fn read_entry(&self, index: u64) -> T {
242        let mut data = std::vec::Vec::new();
243        self.read_entry_to(index, &mut data);
244        T::from_bytes(Cow::Owned(data))
245    }
246
247    /// Reads the item at the specified index without any bound checks.
248    fn read_entry_to(&self, index: u64, buf: &mut Vec<u8>) {
249        let offset = DATA_OFFSET + slot_size::<T>() as u64 * index;
250        let (data_offset, data_size) = self.read_entry_size(offset);
251        read_to_vec(&self.memory, data_offset.into(), buf, data_size);
252    }
253
254    /// Sets the vector's length.
255    fn set_len(&self, new_len: u64) {
256        write_u64(&self.memory, Address::from(LEN_OFFSET), new_len);
257    }
258
259    /// Writes the size of the item at the specified offset.
260    fn write_entry_size(&self, offset: u64, size: u32) -> Result<u64, GrowFailed> {
261        let t_bounds = bounds::<T>();
262        debug_assert!(size <= t_bounds.max_size);
263
264        if t_bounds.is_fixed_size {
265            Ok(offset)
266        } else if t_bounds.max_size <= u8::MAX as u32 {
267            safe_write(&self.memory, offset, &[size as u8; 1])?;
268            Ok(offset + 1)
269        } else if t_bounds.max_size <= u16::MAX as u32 {
270            safe_write(&self.memory, offset, &(size as u16).to_le_bytes())?;
271            Ok(offset + 2)
272        } else {
273            safe_write(&self.memory, offset, &size.to_le_bytes())?;
274            Ok(offset + 4)
275        }
276    }
277
278    /// Reads the size of the entry at the specified offset.
279    fn read_entry_size(&self, offset: u64) -> (u64, usize) {
280        let t_bounds = bounds::<T>();
281        if t_bounds.is_fixed_size {
282            (offset, t_bounds.max_size as usize)
283        } else if t_bounds.max_size <= u8::MAX as u32 {
284            let mut size = [0u8; 1];
285            self.memory.read(offset, &mut size);
286            (offset + 1, size[0] as usize)
287        } else if t_bounds.max_size <= u16::MAX as u32 {
288            let mut size = [0u8; 2];
289            self.memory.read(offset, &mut size);
290            (offset + 2, u16::from_le_bytes(size) as usize)
291        } else {
292            let size = read_u32(&self.memory, Address::from(offset));
293            (offset + 4, size as usize)
294        }
295    }
296
297    /// Write the layout header to the memory.
298    fn write_header(header: &HeaderV1, memory: &M) -> Result<(), GrowFailed> {
299        safe_write(memory, 0, &header.magic)?;
300        memory.write(3, &[header.version; 1]);
301        write_u64(memory, Address::from(4), header.len);
302        write_u32(memory, Address::from(12), header.max_size);
303        memory.write(16, &[if header.is_fixed_size { 1u8 } else { 0u8 }; 1]);
304        Ok(())
305    }
306
307    /// Reads the header from the specified memory.
308    ///
309    /// PRECONDITION: memory.size() > 0
310    fn read_header(memory: &M) -> HeaderV1 {
311        debug_assert!(memory.size() > 0);
312
313        let mut magic = [0u8; 3];
314        let mut version = [0u8; 1];
315        let mut is_fixed_size = [0u8; 1];
316        memory.read(0, &mut magic);
317        memory.read(3, &mut version);
318        let len = read_u64(memory, Address::from(LEN_OFFSET));
319        let max_size = read_u32(memory, Address::from(12));
320        memory.read(16, &mut is_fixed_size);
321
322        HeaderV1 {
323            magic,
324            version: version[0],
325            len,
326            max_size,
327            is_fixed_size: is_fixed_size[0] != 0,
328        }
329    }
330
331    pub fn to_vec(&self) -> Vec<T> {
332        self.iter().collect()
333    }
334}
335
336impl<T: Storable + fmt::Debug, M: Memory> fmt::Debug for BaseVec<T, M> {
337    fn fmt(&self, fmt: &mut fmt::Formatter<'_>) -> fmt::Result {
338        self.to_vec().fmt(fmt)
339    }
340}
341
342fn slot_size<T: Storable>() -> u32 {
343    let t_bounds = bounds::<T>();
344    t_bounds.max_size + bytes_to_store_size_bounded(&t_bounds)
345}
346
347pub struct Iter<'a, T, M>
348where
349    T: Storable,
350    M: Memory,
351{
352    vec: &'a BaseVec<T, M>,
353    buf: std::vec::Vec<u8>,
354    range: Range<u64>,
355}
356
357impl<T, M> Iterator for Iter<'_, T, M>
358where
359    T: Storable,
360    M: Memory,
361{
362    type Item = T;
363
364    fn next(&mut self) -> Option<T> {
365        if self.range.is_empty() || self.vec.len() <= self.range.start {
366            return None;
367        }
368
369        self.vec.read_entry_to(self.range.start, &mut self.buf);
370        self.range.start = self.range.start.saturating_add(1);
371        Some(T::from_bytes(Cow::Borrowed(&self.buf)))
372    }
373
374    fn size_hint(&self) -> (usize, Option<usize>) {
375        (
376            min(self.vec.len(), self.range.end).saturating_sub(self.range.start) as usize,
377            None,
378        )
379    }
380
381    fn count(self) -> usize {
382        min(self.vec.len(), self.range.end)
383            .saturating_sub(self.range.start)
384            .try_into()
385            .expect("Cannot express count as usize")
386    }
387
388    fn nth(&mut self, n: usize) -> Option<T> {
389        self.range.start = self.range.start.saturating_add(n as u64);
390        self.next()
391    }
392}
393
394impl<T, M> DoubleEndedIterator for Iter<'_, T, M>
395where
396    T: Storable,
397    M: Memory,
398{
399    fn next_back(&mut self) -> Option<Self::Item> {
400        if self.range.is_empty() || self.vec.len() < self.range.end {
401            return None;
402        }
403
404        self.vec.read_entry_to(self.range.end - 1, &mut self.buf);
405        self.range.end = self.range.end.saturating_sub(1);
406        Some(T::from_bytes(Cow::Borrowed(&self.buf)))
407    }
408
409    fn nth_back(&mut self, n: usize) -> Option<Self::Item> {
410        self.range.end = self.range.end.saturating_sub(n as u64);
411        self.next_back()
412    }
413}