dfir_rs/util/
accumulator.rs

1//! Accumulator trait and implementors.
2use std::collections::hash_map::Entry;
3
4/// Generalization of fold, reduce, etc.
5pub trait Accumulator<ValAccum, ValIn> {
6    /// Accumulates a value into an either occupied or vacant table entry.
7    fn accumulate<Key>(&mut self, entry: Entry<'_, Key, ValAccum>, item: ValIn);
8}
9
10/// Fold with an initialization and fold function.
11pub struct Fold<InitFn, FoldFn> {
12    init_fn: InitFn,
13    fold_fn: FoldFn,
14}
15
16impl<InitFn, FoldFn> Fold<InitFn, FoldFn> {
17    /// Create a `Fold` [`Accumulator`] with the given `InitFn` and `FoldFn`.
18    pub fn new<Accum, Item>(init_fn: InitFn, fold_fn: FoldFn) -> Self
19    where
20        Self: Accumulator<Accum, Item>,
21    {
22        Self { init_fn, fold_fn }
23    }
24}
25
26impl<InitFn, FoldFn, Accum, Item> Accumulator<Accum, Item> for Fold<InitFn, FoldFn>
27where
28    InitFn: Fn() -> Accum,
29    FoldFn: Fn(&mut Accum, Item),
30{
31    fn accumulate<Key>(&mut self, entry: Entry<'_, Key, Accum>, item: Item) {
32        let prev_item = entry.or_insert_with(|| (self.init_fn)());
33        let () = (self.fold_fn)(prev_item, item);
34    }
35}
36
37/// Reduce with a reduce function.
38pub struct Reduce<ReduceFn> {
39    reduce_fn: ReduceFn,
40}
41
42impl<ReduceFn> Reduce<ReduceFn> {
43    /// Create a `Reduce` [`Accumulator`] with the given `ReduceFn`.
44    pub fn new<Item>(reduce_fn: ReduceFn) -> Self
45    where
46        Self: Accumulator<Item, Item>,
47    {
48        Self { reduce_fn }
49    }
50}
51
52impl<ReduceFn, Item> Accumulator<Item, Item> for Reduce<ReduceFn>
53where
54    ReduceFn: Fn(&mut Item, Item),
55{
56    fn accumulate<Key>(&mut self, entry: Entry<'_, Key, Item>, item: Item) {
57        match entry {
58            Entry::Vacant(entry) => {
59                entry.insert(item);
60            }
61            Entry::Occupied(mut entry) => {
62                let prev_item = entry.get_mut();
63                let () = (self.reduce_fn)(prev_item, item);
64            }
65        }
66    }
67}
68
69/// Fold but with initialization by converting the first received item.
70pub struct FoldFrom<InitFn, FoldFn> {
71    init_fn: InitFn,
72    fold_fn: FoldFn,
73}
74
75impl<InitFn, FoldFn> FoldFrom<InitFn, FoldFn> {
76    /// Create a `FoldFrom` [`Accumulator`] with the given `InitFn` and `FoldFn`.
77    pub fn new<Accum, Item>(init_fn: InitFn, fold_fn: FoldFn) -> Self
78    where
79        Self: Accumulator<Accum, Item>,
80    {
81        Self { init_fn, fold_fn }
82    }
83}
84
85impl<InitFn, FoldFn, Accum, Item> Accumulator<Accum, Item> for FoldFrom<InitFn, FoldFn>
86where
87    InitFn: Fn(Item) -> Accum,
88    FoldFn: Fn(&mut Accum, Item),
89{
90    fn accumulate<Key>(&mut self, entry: Entry<'_, Key, Accum>, item: Item) {
91        match entry {
92            Entry::Vacant(entry) => {
93                entry.insert((self.init_fn)(item));
94            }
95            Entry::Occupied(mut entry) => {
96                let prev_item = entry.get_mut();
97                let () = (self.fold_fn)(prev_item, item);
98            }
99        }
100    }
101}