inc_complete/storage/
indexmapped.rs

1use scc::{TreeIndex, ebr::Guard};
2
3use crate::{Cell, storage::StorageFor};
4
5use super::Computation;
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: Computation> {
11    key_to_cell: TreeIndex<K, Cell>,
12    cell_to_key: TreeIndex<Cell, (K, Option<K::Output>)>,
13}
14
15impl<K: Computation> 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 + Computation + '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 try_get_input(&self, cell: Cell) -> Option<K> {
39        self.cell_to_key.peek_with(&cell, |_, (k, _)| k.clone())
40    }
41
42    fn get_output(&self, cell: Cell) -> Option<K::Output> {
43        self.cell_to_key
44            .peek_with(&cell, |_, (_, v)| v.clone())
45            .unwrap()
46    }
47
48    fn update_output(&self, cell: Cell, new_value: K::Output) -> bool {
49        let changed = K::ASSUME_CHANGED
50            || self
51                .cell_to_key
52                .peek_with(&cell, |_, old_value| {
53                    old_value.1.as_ref().is_none_or(|value| *value != new_value)
54                })
55                .unwrap();
56
57        // TreeIndex is read-optimized so this write will be slow
58        if changed {
59            let key = self
60                .cell_to_key
61                .peek_with(&cell, |_, (k, _)| k.clone())
62                .unwrap();
63            self.cell_to_key.remove(&cell);
64            self.cell_to_key.insert(cell, (key, Some(new_value))).ok();
65        }
66
67        changed
68    }
69
70    fn gc(&mut self, used_cells: &std::collections::HashSet<Cell>) {
71        let guard = Guard::new();
72        for cell in self.cell_to_key.iter(&guard) {
73            if !used_cells.contains(cell.0) {
74                self.cell_to_key.remove(cell.0);
75            }
76        }
77        for cell in self.key_to_cell.iter(&guard) {
78            if !used_cells.contains(cell.1) {
79                self.key_to_cell.remove(cell.0);
80            }
81        }
82    }
83}
84
85impl<K> serde::Serialize for TreeIndexStorage<K>
86where
87    K: serde::Serialize + Computation + Eq + Clone + 'static,
88    K::Output: serde::Serialize + Clone + 'static,
89{
90    fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
91    where
92        S: serde::Serializer,
93    {
94        let mut cell_to_key_vec: Vec<(Cell, (K, Option<K::Output>))> =
95            Vec::with_capacity(self.cell_to_key.len());
96
97        let guard = Guard::new();
98        for (cell, (key, value)) in self.cell_to_key.iter(&guard) {
99            cell_to_key_vec.push((*cell, (key.clone(), value.clone())));
100        }
101
102        cell_to_key_vec.serialize(serializer)
103    }
104}
105
106impl<'de, K> serde::Deserialize<'de> for TreeIndexStorage<K>
107where
108    K: serde::Deserialize<'de> + Ord + Computation + Clone + 'static,
109    K::Output: serde::Deserialize<'de> + Clone,
110{
111    fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
112    where
113        D: serde::Deserializer<'de>,
114    {
115        let cell_to_key_vec: Vec<(Cell, (K, Option<K::Output>))> =
116            serde::Deserialize::deserialize(deserializer)?;
117
118        let key_to_cell = TreeIndex::new();
119        let cell_to_key = TreeIndex::new();
120
121        for (cell, (key, value)) in cell_to_key_vec {
122            key_to_cell.insert(key.clone(), cell).ok();
123            cell_to_key.insert(cell, (key, value)).ok();
124        }
125
126        Ok(TreeIndexStorage {
127            cell_to_key,
128            key_to_cell,
129        })
130    }
131}
inc_complete/storage/indexmapped.rs

inc_complete/storage/
indexmapped.rs
