inc_complete/storage/
btreemapped.rs

1use std::collections::BTreeMap;
2
3use crate::{Cell, storage::StorageFor};
4
5use super::OutputType;
6
7#[derive(Clone)]
8pub struct BTreeMapStorage<K: Ord + OutputType> {
9    key_to_cell: BTreeMap<K, Cell>,
10    cell_to_key: BTreeMap<Cell, (K, Option<K::Output>)>,
11}
12
13impl<K: Ord + OutputType> Default for BTreeMapStorage<K> {
14    fn default() -> Self {
15        Self {
16            key_to_cell: Default::default(),
17            cell_to_key: Default::default(),
18        }
19    }
20}
21
22impl<K> StorageFor<K> for BTreeMapStorage<K>
23where
24    K: Clone + Ord + OutputType,
25    K::Output: Eq,
26{
27    fn get_cell_for_computation(&self, key: &K) -> Option<Cell> {
28        self.key_to_cell.get(key).copied()
29    }
30
31    fn insert_new_cell(&mut self, cell: Cell, key: K) {
32        self.key_to_cell.insert(key.clone(), cell);
33        self.cell_to_key.insert(cell, (key, None));
34    }
35
36    fn get_input(&self, cell: Cell) -> &K {
37        &self.cell_to_key[&cell].0
38    }
39
40    fn get_output(&self, cell: Cell) -> Option<&K::Output> {
41        self.cell_to_key[&cell].1.as_ref()
42    }
43
44    fn update_output(&mut self, cell: Cell, new_value: K::Output) -> bool {
45        let previous_output = &mut self.cell_to_key.get_mut(&cell).unwrap().1;
46        let changed = previous_output
47            .as_ref()
48            .is_none_or(|value| *value != new_value);
49        *previous_output = Some(new_value);
50        changed
51    }
52}
53
54impl<K> serde::Serialize for BTreeMapStorage<K>
55where
56    K: serde::Serialize + Ord + OutputType,
57    K::Output: serde::Serialize,
58{
59    fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
60    where
61        S: serde::Serializer,
62    {
63        self.cell_to_key.serialize(serializer)
64    }
65}
66
67impl<'de, K> serde::Deserialize<'de> for BTreeMapStorage<K>
68where
69    K: serde::Deserialize<'de> + Ord + OutputType + Clone,
70    K::Output: serde::Deserialize<'de>,
71{
72    fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
73    where
74        D: serde::Deserializer<'de>,
75    {
76        let cell_to_key: BTreeMap<Cell, (K, Option<K::Output>)> =
77            serde::Deserialize::deserialize(deserializer)?;
78        let key_to_cell = cell_to_key
79            .iter()
80            .map(|(cell, (key, _))| (key.clone(), *cell))
81            .collect();
82        Ok(BTreeMapStorage {
83            cell_to_key,
84            key_to_cell,
85        })
86    }
87}