Skip to main content

shuttle_engine/runtime/
storage.rs

1use std::any::Any;
2use std::collections::{HashMap, VecDeque};
3
4/// A unique identifier for a storage slot
5#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
6pub struct StorageKey(pub usize, pub usize); // (identifier, type)
7
8/// A map of storage values.
9///
10/// We remember the insertion order into the storage HashMap so that destruction is deterministic.
11/// Values are Option<_> because we need to be able to incrementally destruct them, as it's valid
12/// for TLS destructors to initialize new TLS slots. When a slot is destructed, its key is removed
13/// from `order` and its value is replaced with None.
14#[derive(Debug)]
15pub struct StorageMap {
16    locals: HashMap<StorageKey, Option<Box<dyn Any>>>,
17    order: VecDeque<StorageKey>,
18}
19
20impl StorageMap {
21    pub fn new() -> Self {
22        Self {
23            locals: HashMap::new(),
24            order: VecDeque::new(),
25        }
26    }
27
28    pub fn get<T: 'static>(&self, key: StorageKey) -> Option<Result<&T, AlreadyDestructedError>> {
29        self.locals.get(&key).map(|val| {
30            val.as_ref()
31                .map(|val| {
32                    Ok(val
33                        .downcast_ref::<T>()
34                        .expect("local value must downcast to expected type"))
35                })
36                .unwrap_or(Err(AlreadyDestructedError))
37        })
38    }
39
40    pub fn init<T: 'static>(&mut self, key: StorageKey, value: T) {
41        let result = self.locals.insert(key, Some(Box::new(value)));
42        assert!(result.is_none(), "cannot reinitialize a storage slot");
43        self.order.push_back(key);
44    }
45
46    /// Whether no slot is still initialized.
47    pub fn is_empty(&self) -> bool {
48        self.order.is_empty()
49    }
50
51    /// Return ownership of the next still-initialized storage slot.
52    pub fn pop(&mut self) -> Option<Box<dyn Any>> {
53        let key = self.order.pop_front()?;
54        let value = self
55            .locals
56            .get_mut(&key)
57            .expect("keys in `order` must exist")
58            .take()
59            .expect("keys in `order` must not yet be destructed");
60        Some(value)
61    }
62}
63
64#[derive(Debug)]
65#[non_exhaustive]
66pub struct AlreadyDestructedError;