Skip to main content

inc_complete/
accumulate.rs

1use std::{collections::BTreeSet, marker::PhantomData};
2
3use dashmap::DashMap;
4use serde::{Deserialize, Serialize};
5
6use crate::{Cell, Computation, Run, Storage, StorageFor};
7
8/// An accumulator is a collection which can accumulate a given cell key associated with multiple
9/// values of a given type. `Accumulator<MyItem>` is an example of such a type.
10///
11/// This is most often a hashmap or similar map
12pub trait Accumulate<Item> {
13    /// Push an item to the context of the given cell
14    fn accumulate(&self, cell: Cell, item: Item);
15
16    /// Retrieve all items associated with the given cell.
17    /// Note that this should only include the exact cell given, not any
18    /// values accumulated from dependencies.
19    fn get_accumulated<Items>(&self, cell: Cell) -> Items
20    where
21        Items: FromIterator<Item>;
22}
23
24pub struct Accumulator<Item> {
25    map: DashMap<Cell, Vec<Item>>,
26}
27
28impl<Item> Default for Accumulator<Item> {
29    fn default() -> Self {
30        Self {
31            map: Default::default(),
32        }
33    }
34}
35
36impl<Item> Accumulator<Item> {
37    pub fn clear(&self, cell: Cell) {
38        self.map.remove(&cell);
39    }
40}
41
42impl<Item: Clone> Accumulate<Item> for Accumulator<Item> {
43    fn accumulate(&self, cell: Cell, item: Item) {
44        self.map.entry(cell).or_default().push(item);
45    }
46
47    fn get_accumulated<Items>(&self, cell: Cell) -> Items
48    where
49        Items: FromIterator<Item>,
50    {
51        if let Some(items) = self.map.get(&cell) {
52            FromIterator::from_iter(items.iter().cloned())
53        } else {
54            FromIterator::from_iter(std::iter::empty())
55        }
56    }
57}
58
59#[derive(Serialize, Deserialize, Debug, Copy, Clone, PartialEq, Eq, Hash, PartialOrd, Ord)]
60#[serde(transparent)]
61pub struct Accumulated<Item> {
62    cell: Cell,
63    _item: std::marker::PhantomData<Item>,
64}
65
66impl<Item> Accumulated<Item> {
67    pub(crate) fn new(cell: Cell) -> Self {
68        Self {
69            cell,
70            _item: PhantomData,
71        }
72    }
73}
74
75impl<Item: 'static> Computation for Accumulated<Item> {
76    type Output = BTreeSet<Item>;
77    const IS_INPUT: bool = false;
78    const ASSUME_CHANGED: bool = false;
79
80    fn computation_id() -> u32 {
81        // TODO: Find another way to keep this as a unique ID or at least assert it is different
82        // than every other id.
83        // Arbitrary semi-random value meant to not be easily accidentally used for ids in user code
84        0x54325243
85    }
86}
87
88impl<S, Item> Run<S> for Accumulated<Item>
89where
90    Item: 'static + Ord,
91    S: Storage + StorageFor<Accumulated<Item>> + Accumulate<Item>,
92{
93    fn run(&self, db: &crate::DbHandle<S>) -> Self::Output {
94        db.get_accumulated_with_cell(self.cell)
95    }
96}
97
98impl<Item: Serialize + Clone> Serialize for Accumulator<Item> {
99    fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
100    where
101        S: serde::Serializer,
102    {
103        let vec: Vec<(Cell, Vec<Item>)> = self
104            .map
105            .iter()
106            .map(|entry| (*entry.key(), entry.value().clone()))
107            .collect();
108
109        vec.serialize(serializer)
110    }
111}
112
113impl<'de, Item: Deserialize<'de>> Deserialize<'de> for Accumulator<Item> {
114    fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
115    where
116        D: serde::Deserializer<'de>,
117    {
118        let vec: Vec<(Cell, Vec<Item>)> = Deserialize::deserialize(deserializer)?;
119        let map = vec.into_iter().collect();
120        Ok(Accumulator { map })
121    }
122}