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` (or `last`) 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    /// The keys of the slots that are destructed after all of those in `order`, so that the
19    /// destructors of those can still use them.
20    last: VecDeque<StorageKey>,
21}
22
23impl StorageMap {
24    pub fn new() -> Self {
25        Self {
26            locals: HashMap::new(),
27            order: VecDeque::new(),
28            last: VecDeque::new(),
29        }
30    }
31
32    pub fn get<T: 'static>(&self, key: StorageKey) -> Option<Result<&T, AlreadyDestructedError>> {
33        self.locals.get(&key).map(|val| {
34            val.as_ref()
35                .map(|val| {
36                    Ok(val
37                        .downcast_ref::<T>()
38                        .expect("local value must downcast to expected type"))
39                })
40                .unwrap_or(Err(AlreadyDestructedError))
41        })
42    }
43
44    pub fn init<T: 'static>(&mut self, key: StorageKey, value: T) {
45        self.insert(key, value);
46        self.order.push_back(key);
47    }
48
49    /// Like `init`, for a slot that is destructed after every slot initialized with `init`.
50    pub fn init_destructed_last<T: 'static>(&mut self, key: StorageKey, value: T) {
51        self.insert(key, value);
52        self.last.push_back(key);
53    }
54
55    fn insert<T: 'static>(&mut self, key: StorageKey, value: T) {
56        let result = self.locals.insert(key, Some(Box::new(value)));
57        assert!(result.is_none(), "cannot reinitialize a storage slot");
58    }
59
60    /// Whether no slot is still initialized.
61    pub fn is_empty(&self) -> bool {
62        self.order.is_empty() && self.last.is_empty()
63    }
64
65    /// Return ownership of the next still-initialized storage slot: the oldest of those that are
66    /// not destructed last, or, once there are none of those, the oldest of the others.
67    pub fn pop(&mut self) -> Option<Box<dyn Any>> {
68        let key = self.order.pop_front().or_else(|| self.last.pop_front())?;
69        let value = self
70            .locals
71            .get_mut(&key)
72            .expect("keys in `order` must exist")
73            .take()
74            .expect("keys in `order` must not yet be destructed");
75        Some(value)
76    }
77}
78
79#[derive(Debug)]
80#[non_exhaustive]
81pub struct AlreadyDestructedError;
shuttle_engine/runtime/storage.rs

shuttle_engine/runtime/
storage.rs
