cairo_lang_utils/
collection_arithmetics.rs

1#[cfg(test)]
2#[path = "collection_arithmetics_test.rs"]
3mod test;
4
5use core::hash::{BuildHasher, Hash};
6use core::ops::{Add, Sub};
7
8use crate::ordered_hash_map::{Entry, OrderedHashMap};
9
10/// A trait for types which have a zero value.
11///
12/// Functions may assume the following:
13/// * `x = x + zero() = zero() + x`
14pub trait HasZero {
15    /// Returns the zero value for the type.
16    fn zero() -> Self;
17}
18impl HasZero for i32 {
19    fn zero() -> Self {
20        0
21    }
22}
23impl HasZero for i64 {
24    fn zero() -> Self {
25        0
26    }
27}
28
29/// Returns a map which contains the sum of the values from the given two maps, for each key.
30///
31/// If the key is missing from one of them, it is treated as zero.
32pub fn add_maps<
33    Key: Hash + Eq,
34    Value: HasZero + Add<Output = Value> + Clone + Eq,
35    Rhs: IntoIterator<Item = (Key, Value)>,
36    BH: BuildHasher,
37>(
38    lhs: OrderedHashMap<Key, Value, BH>,
39    rhs: Rhs,
40) -> OrderedHashMap<Key, Value, BH> {
41    merge_maps(lhs, rhs, |a, b| a + b)
42}
43
44/// Returns a map which contains the difference of the values from the given two maps, for each key.
45///
46/// If the key is missing from one of them, it is treated as zero.
47pub fn sub_maps<
48    Key: Hash + Eq,
49    Value: HasZero + Sub<Output = Value> + Clone + Eq,
50    Rhs: IntoIterator<Item = (Key, Value)>,
51    BH: BuildHasher,
52>(
53    lhs: OrderedHashMap<Key, Value, BH>,
54    rhs: Rhs,
55) -> OrderedHashMap<Key, Value, BH> {
56    merge_maps(lhs, rhs, |a, b| a - b)
57}
58
59/// Returns a map which contains the combination by using `action` of the values from the given two
60/// maps, for each key.
61///
62/// If the key is missing from one of them, it is treated as zero.
63fn merge_maps<
64    Key: Hash + Eq,
65    Value: HasZero + Clone + Eq,
66    Rhs: IntoIterator<Item = (Key, Value)>,
67    Action: Fn(Value, Value) -> Value,
68    BH: BuildHasher,
69>(
70    lhs: OrderedHashMap<Key, Value, BH>,
71    rhs: Rhs,
72    action: Action,
73) -> OrderedHashMap<Key, Value, BH> {
74    let mut res = lhs;
75    for (key, rhs_val) in rhs {
76        match res.entry(key) {
77            Entry::Occupied(mut e) => {
78                let new_val = action(e.get().clone(), rhs_val);
79                if new_val == Value::zero() {
80                    e.swap_remove();
81                } else {
82                    e.insert(new_val);
83                }
84            }
85            Entry::Vacant(e) => {
86                if rhs_val != Value::zero() {
87                    e.insert(action(Value::zero(), rhs_val));
88                }
89            }
90        }
91    }
92    res
93}
cairo_lang_utils/collection_arithmetics.rs

cairo_lang_utils/
collection_arithmetics.rs
