inc_complete/storage/
hashmapped.rs

1use scc::HashMap;
2
3use crate::{Cell, storage::StorageFor};
4use std::hash::{BuildHasher, Hash};
5
6use super::OutputType;
7
8pub struct HashMapStorage<K, Hasher = rustc_hash::FxBuildHasher> where
9    K: OutputType + Eq + Hash,
10    Hasher: BuildHasher,
11{
12    key_to_cell: HashMap<K, Cell, Hasher>,
13    cell_to_key: HashMap<Cell, (K, Option<K::Output>), Hasher>,
14}
15
16impl<K, H> Default for HashMapStorage<K, H> where
17    K: OutputType + Eq + Hash,
18    H: Default + BuildHasher,
19{
20    fn default() -> Self {
21        Self {
22            key_to_cell: Default::default(),
23            cell_to_key: Default::default(),
24        }
25    }
26}
27
28impl<K, H> StorageFor<K> for HashMapStorage<K, H>
29where
30    K: Clone + Eq + Hash + OutputType,
31    K::Output: Eq + Clone,
32    H: BuildHasher,
33{
34    fn get_cell_for_computation(&self, key: &K) -> Option<Cell> {
35        self.key_to_cell.get(key).map(|value| *value)
36    }
37
38    fn insert_new_cell(&self, cell: Cell, key: K) {
39        self.key_to_cell.insert(key.clone(), cell).ok();
40        self.cell_to_key.insert(cell, (key, None)).ok();
41    }
42
43    fn get_input(&self, cell: Cell) -> K {
44        self.cell_to_key.get(&cell).unwrap().0.clone()
45    }
46
47    fn get_output(&self, cell: Cell) -> Option<K::Output> {
48        self.cell_to_key.get(&cell).unwrap().1.clone()
49    }
50
51    fn update_output(&self, cell: Cell, new_value: K::Output) -> bool {
52        let mut previous_output = self.cell_to_key.get(&cell).unwrap();
53        let changed = previous_output
54            .1
55            .as_ref()
56            .is_none_or(|value| *value != new_value);
57        previous_output.1 = Some(new_value);
58        changed
59    }
60}
61
62impl<K, H> serde::Serialize for HashMapStorage<K, H>
63where
64    K: serde::Serialize + OutputType + Eq + Hash + Clone,
65    K::Output: serde::Serialize + Clone,
66    H: BuildHasher,
67{
68    fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
69    where
70        S: serde::Serializer,
71    {
72        let mut cell_to_key_vec: Vec<(Cell, (K, Option<K::Output>))> = Vec::with_capacity(self.cell_to_key.len());
73
74        self.cell_to_key.scan(|cell, (key, value)| {
75            cell_to_key_vec.push((*cell, (key.clone(), value.clone())));
76        });
77
78        cell_to_key_vec.serialize(serializer)
79    }
80}
81
82impl<'de, K, H> serde::Deserialize<'de> for HashMapStorage<K, H>
83where
84    K: serde::Deserialize<'de> + Hash + Eq + OutputType + Clone,
85    K::Output: serde::Deserialize<'de>,
86    H: Default + BuildHasher,
87{
88    fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
89    where
90        D: serde::Deserializer<'de>,
91    {
92        let cell_to_key_vec: Vec<(Cell, (K, Option<K::Output>))> = serde::Deserialize::deserialize(deserializer)?;
93
94        let key_to_cell = HashMap::default();
95        let cell_to_key = HashMap::default();
96
97        for (cell, (key, value)) in cell_to_key_vec {
98            key_to_cell.insert(key.clone(), cell).ok();
99            cell_to_key.insert(cell, (key, value)).ok();
100        }
101
102        Ok(HashMapStorage { cell_to_key, key_to_cell })
103    }
104}