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

inc_complete/storage/
indexmapped.rs
