inc_complete/storage/
hashmapped.rs

1use dashmap::DashMap;
2
3use crate::{Cell, storage::StorageFor};
4use std::hash::{BuildHasher, Hash};
5
6use super::Computation;
7
8pub struct HashMapStorage<K, Hasher = rustc_hash::FxBuildHasher>
9where
10    K: Computation + Eq + Hash,
11    Hasher: BuildHasher,
12{
13    key_to_cell: DashMap<K, Cell, Hasher>,
14    cell_to_key: DashMap<Cell, (K, Option<K::Output>), Hasher>,
15}
16
17impl<K, H> Default for HashMapStorage<K, H>
18where
19    K: Computation + Eq + Hash,
20    H: Default + BuildHasher + Clone,
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 + Computation,
33    K::Output: Eq + Clone,
34    H: BuildHasher + Clone,
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        // key_to_cell must be written last to avoid data races
42        self.cell_to_key.insert(cell, (key.clone(), None));
43        self.key_to_cell.insert(key, cell);
44    }
45
46    fn try_get_input(&self, cell: Cell) -> Option<K> {
47        let key_ref = self.cell_to_key.get(&cell)?;
48        Some(key_ref.0.clone())
49    }
50
51    fn get_input(&self, cell: Cell) -> K {
52        self.cell_to_key.get(&cell).unwrap().0.clone()
53    }
54
55    fn get_output(&self, cell: Cell) -> Option<K::Output> {
56        self.cell_to_key.get(&cell).unwrap().1.clone()
57    }
58
59    fn update_output(&self, cell: Cell, new_value: K::Output) -> bool {
60        let mut previous_output = self.cell_to_key.get_mut(&cell).unwrap();
61        let changed = K::ASSUME_CHANGED
62            || previous_output
63                .1
64                .as_ref()
65                .is_none_or(|value| *value != new_value);
66        previous_output.1 = Some(new_value);
67        changed
68    }
69
70    fn gc(&mut self, used_cells: &std::collections::HashSet<Cell>) {
71        // Remove cells that are not in the used set
72        self.cell_to_key.retain(|cell, _| used_cells.contains(cell));
73        self.key_to_cell.retain(|_, cell| used_cells.contains(cell));
74    }
75}
76
77impl<K, H> serde::Serialize for HashMapStorage<K, H>
78where
79    K: serde::Serialize + Computation + Eq + Hash + Clone,
80    K::Output: serde::Serialize + Clone,
81    H: BuildHasher + Clone,
82{
83    fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
84    where
85        S: serde::Serializer,
86    {
87        let mut cell_to_key_vec: Vec<(Cell, (K, Option<K::Output>))> =
88            Vec::with_capacity(self.cell_to_key.len());
89
90        for kv in self.cell_to_key.iter() {
91            let cell = *kv.key();
92            let (key, value) = kv.value().clone();
93            cell_to_key_vec.push((cell, (key, value)));
94        }
95
96        cell_to_key_vec.serialize(serializer)
97    }
98}
99
100impl<'de, K, H> serde::Deserialize<'de> for HashMapStorage<K, H>
101where
102    K: serde::Deserialize<'de> + Hash + Eq + Computation + Clone,
103    K::Output: serde::Deserialize<'de>,
104    H: Default + BuildHasher + Clone,
105{
106    fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
107    where
108        D: serde::Deserializer<'de>,
109    {
110        let cell_to_key_vec: Vec<(Cell, (K, Option<K::Output>))> =
111            serde::Deserialize::deserialize(deserializer)?;
112
113        let key_to_cell = DashMap::default();
114        let cell_to_key = DashMap::default();
115
116        for (cell, (key, value)) in cell_to_key_vec {
117            key_to_cell.insert(key.clone(), cell);
118            cell_to_key.insert(cell, (key, value));
119        }
120
121        Ok(HashMapStorage {
122            cell_to_key,
123            key_to_cell,
124        })
125    }
126}
inc_complete/storage/hashmapped.rs

inc_complete/storage/
hashmapped.rs
