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    /// Return ownership of the next still-initialized storage slot.
47    pub fn pop(&mut self) -> Option<Box<dyn Any>> {
48        let key = self.order.pop_front()?;
49        let value = self
50            .locals
51            .get_mut(&key)
52            .expect("keys in `order` must exist")
53            .take()
54            .expect("keys in `order` must not yet be destructed");
55        Some(value)
56    }
57}
58
59#[derive(Debug)]
60#[non_exhaustive]
61pub struct AlreadyDestructedError;