inc_complete/storage/
indexmapped.rs

1use scc::{TreeIndex, ebr::Guard};
2
3use crate::{Cell, storage::StorageFor};
4
5use super::OutputType;
6
7/// TreeIndexStorage is backed internally by a `scc::TreeIndex`,
8/// a type optimized for read-heavy workloads. See [scc's documentation](https://docs.rs/scc/2.3.4/scc/#treeindex)
9/// for performance details.
10pub struct TreeIndexStorage<K: OutputType> {
11    key_to_cell: TreeIndex<K, Cell>,
12    cell_to_key: TreeIndex<Cell, (K, Option<K::Output>)>,
13}
14
15impl<K: OutputType> Default for TreeIndexStorage<K> {
16    fn default() -> Self {
17        Self {
18            key_to_cell: Default::default(),
19            cell_to_key: Default::default(),
20        }
21    }
22}
23
24impl<K> StorageFor<K> for TreeIndexStorage<K>
25where
26    K: Clone + Ord + OutputType + 'static,
27    K::Output: Clone + Eq,
28{
29    fn get_cell_for_computation(&self, key: &K) -> Option<Cell> {
30        self.key_to_cell.peek_with(key, |_, v| *v)
31    }
32
33    fn insert_new_cell(&self, cell: Cell, key: K) {
34        self.key_to_cell.insert(key.clone(), cell).ok();
35        self.cell_to_key.insert(cell, (key, None)).ok();
36    }
37
38    fn get_input(&self, cell: Cell) -> K {
39        self.cell_to_key
40            .peek_with(&cell, |_, (k, _)| k.clone())
41            .unwrap()
42    }
43
44    fn get_output(&self, cell: Cell) -> Option<K::Output> {
45        self.cell_to_key
46            .peek_with(&cell, |_, (_, v)| v.clone())
47            .unwrap()
48    }
49
50    fn update_output(&self, cell: Cell, new_value: K::Output) -> bool {
51        let changed = K::ASSUME_CHANGED
52            || self
53                .cell_to_key
54                .peek_with(&cell, |_, old_value| {
55                    old_value.1.as_ref().is_none_or(|value| *value != new_value)
56                })
57                .unwrap();
58
59        // TreeIndex is read-optimized so this write will be slow
60        if changed {
61            let key = self
62                .cell_to_key
63                .peek_with(&cell, |_, (k, _)| k.clone())
64                .unwrap();
65            self.cell_to_key.remove(&cell);
66            self.cell_to_key.insert(cell, (key, Some(new_value))).ok();
67        }
68
69        changed
70    }
71}
72
73impl<K> serde::Serialize for TreeIndexStorage<K>
74where
75    K: serde::Serialize + OutputType + Eq + Clone + 'static,
76    K::Output: serde::Serialize + Clone + 'static,
77{
78    fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
79    where
80        S: serde::Serializer,
81    {
82        let mut cell_to_key_vec: Vec<(Cell, (K, Option<K::Output>))> =
83            Vec::with_capacity(self.cell_to_key.len());
84
85        let guard = Guard::new();
86        for (cell, (key, value)) in self.cell_to_key.iter(&guard) {
87            cell_to_key_vec.push((*cell, (key.clone(), value.clone())));
88        }
89
90        cell_to_key_vec.serialize(serializer)
91    }
92}
93
94impl<'de, K> serde::Deserialize<'de> for TreeIndexStorage<K>
95where
96    K: serde::Deserialize<'de> + Ord + OutputType + Clone + 'static,
97    K::Output: serde::Deserialize<'de> + Clone,
98{
99    fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
100    where
101        D: serde::Deserializer<'de>,
102    {
103        let cell_to_key_vec: Vec<(Cell, (K, Option<K::Output>))> =
104            serde::Deserialize::deserialize(deserializer)?;
105
106        let key_to_cell = TreeIndex::new();
107        let cell_to_key = TreeIndex::new();
108
109        for (cell, (key, value)) in cell_to_key_vec {
110            key_to_cell.insert(key.clone(), cell).ok();
111            cell_to_key.insert(cell, (key, value)).ok();
112        }
113
114        Ok(TreeIndexStorage {
115            cell_to_key,
116            key_to_cell,
117        })
118    }
119}
inc_complete/storage/indexmapped.rs

inc_complete/storage/
indexmapped.rs
