inc_complete/storage/
hashmapped.rs

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