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

inc_complete/storage/
indexmapped.rs
