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