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    fn gc(&mut self, used_cells: &std::collections::HashSet<Cell>) {
73        let guard = Guard::new();
74        for cell in self.cell_to_key.iter(&guard) {
75            if !used_cells.contains(cell.0) {
76                self.cell_to_key.remove(cell.0);
77            }
78        }
79        for cell in self.key_to_cell.iter(&guard) {
80            if !used_cells.contains(cell.1) {
81                self.key_to_cell.remove(cell.0);
82            }
83        }
84    }
85}
86
87impl<K> serde::Serialize for TreeIndexStorage<K>
88where
89    K: serde::Serialize + OutputType + Eq + Clone + 'static,
90    K::Output: serde::Serialize + Clone + 'static,
91{
92    fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
93    where
94        S: serde::Serializer,
95    {
96        let mut cell_to_key_vec: Vec<(Cell, (K, Option<K::Output>))> =
97            Vec::with_capacity(self.cell_to_key.len());
98
99        let guard = Guard::new();
100        for (cell, (key, value)) in self.cell_to_key.iter(&guard) {
101            cell_to_key_vec.push((*cell, (key.clone(), value.clone())));
102        }
103
104        cell_to_key_vec.serialize(serializer)
105    }
106}
107
108impl<'de, K> serde::Deserialize<'de> for TreeIndexStorage<K>
109where
110    K: serde::Deserialize<'de> + Ord + OutputType + Clone + 'static,
111    K::Output: serde::Deserialize<'de> + Clone,
112{
113    fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
114    where
115        D: serde::Deserializer<'de>,
116    {
117        let cell_to_key_vec: Vec<(Cell, (K, Option<K::Output>))> =
118            serde::Deserialize::deserialize(deserializer)?;
119
120        let key_to_cell = TreeIndex::new();
121        let cell_to_key = TreeIndex::new();
122
123        for (cell, (key, value)) in cell_to_key_vec {
124            key_to_cell.insert(key.clone(), cell).ok();
125            cell_to_key.insert(cell, (key, value)).ok();
126        }
127
128        Ok(TreeIndexStorage {
129            cell_to_key,
130            key_to_cell,
131        })
132    }
133}
inc_complete/storage/indexmapped.rs

inc_complete/storage/
indexmapped.rs
