dyadic_rationals/
context.rs

1use std::collections::{BTreeMap, BTreeSet};
2use std::fmt;
3use pretty::{Pretty, DocAllocator, DocBuilder, BoxAllocator};
4
5/// General BTreeMap context
6#[derive(Clone, PartialEq, Eq, PartialOrd, Ord, Hash, Debug)]
7pub struct Ctx<K, V>(BTreeMap<K, V>);
8
9/// Pretty printer instance for Ctx
10impl<'a, D, A, K, V> Pretty<'a, D, A> for Ctx<K, V>
11where
12    K: Clone + Pretty<'a, D, A>,
13    V: Clone + Pretty<'a, D, A>,
14    D: DocAllocator<'a, A>,
15    D::Doc: Clone,
16    A: 'a + Clone,
17{
18    fn pretty(self, allocator: &'a D) -> DocBuilder<'a, D, A> {
19        allocator.concat([
20            allocator.text("{"),
21            allocator.intersperse(
22                self.0.iter().map(|(k, v)| {
23                    k.clone().pretty(allocator)
24                        .append(allocator.text(" -> "))
25                        .append(v.clone().pretty(allocator))
26                }),
27                ", ",
28            ),
29            allocator.text("}"),
30        ])
31    }
32}
33
34/// IntoIterator instance for Ctx
35impl<K, V> IntoIterator for Ctx<K, V> {
36    type Item = (K, V);
37    type IntoIter = std::collections::btree_map::IntoIter<K, V>;
38
39    fn into_iter(self) -> Self::IntoIter {
40        self.0.into_iter()
41    }
42}
43
44/// Iterator for borrowing key-value pairs
45pub struct CtxIterator<'a, K, V> {
46    iter: std::collections::btree_map::Iter<'a, K, V>,
47}
48
49/// Iterator instance for Ctx
50impl<'a, K, V> Iterator for CtxIterator<'a, K, V> {
51    type Item = (&'a K, &'a V);
52
53    fn next(&mut self) -> Option<Self::Item> {
54        self.iter.next()
55    }
56}
57
58impl<K, V> FromIterator<(K, V)> for Ctx<K, V>
59where
60    K: Ord
61{
62    fn from_iter<I: IntoIterator<Item = (K, V)>>(iter: I) -> Self {
63        Ctx(BTreeMap::from_iter(iter))
64    }
65}
66
67/// Display instance for Ctx calls the pretty printer
68impl<'a, K, V> fmt::Display for Ctx<K, V>
69where
70    K: Pretty<'a, BoxAllocator, ()> + Clone,
71    V: Pretty<'a, BoxAllocator, ()> + Clone
72{
73    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
74        <Ctx<_, _> as Pretty<'_, BoxAllocator, ()>>::pretty(self.clone(), &BoxAllocator)
75            .1
76            .render_fmt(80, f)
77    }
78}
79
80#[cfg(test)] use arbitrary::{Arbitrary, Unstructured};
81/// Arbitrary instance for Ctx
82#[cfg(test)]
83impl<'a, K: Arbitrary<'a> + Ord> Arbitrary<'a> for Ctx<K, u8> {
84    fn arbitrary(u: &mut Unstructured<'a>) -> arbitrary::Result<Self, arbitrary::Error> {
85        let n: u8 = u.int_in_range(0..=9)?;
86        let mut vars = BTreeMap::new();
87        for _ in 0..n {
88            vars.insert(K::arbitrary(u)?, u.int_in_range(0..=12)?);
89        }
90        Ok(Ctx(vars))
91    }
92}
93
94/// Special and wrapper methods for Ctx
95impl<K: Ord, V> Ctx<K, V> {
96    pub fn new() -> Self {
97        Ctx(BTreeMap::new())
98    }
99
100    pub fn singleton(k: K, v: V) -> Self {
101        Ctx(BTreeMap::from([(k, v)]))
102    }
103
104    pub fn insert_with<FF>(&mut self, k: K, v1: V, f: &FF)
105    where
106        V:Clone,
107        FF: Fn(V, V) -> V
108    {
109        self.0.entry(k).and_modify(|v2| *v2 = f(v1.clone(), v2.clone())).or_insert(v1);
110    }
111
112    pub fn insert(&mut self, k: K, v: V)
113    where V: Clone
114    {
115        self.insert_with(k, v, &|_, v| v);
116    }
117
118    pub fn append_with<It, FF>(&mut self, it: It, f: &FF) -> &mut Self
119    where
120        V: Clone,
121        It: Iterator<Item = (K, V)>,
122        FF: Fn(V, V) -> V,
123    {
124        for (k, v) in it {
125            self.insert_with(k, v, f);
126        }
127        self
128    }
129    pub fn append<It>(&mut self, it: It) -> &mut Self
130    where
131        V: Clone,
132        It: Iterator<Item = (K, V)>,
133    {
134        self.append_with(it, &|_, v| v)
135    }
136
137    pub fn union_with<FF>(&self, other: Self, f: &FF) -> Self
138    where
139        K: Clone,
140        V: Clone,
141        FF: Fn(V, V) -> V
142    {
143        let mut c = self.clone();
144        c.append_with(other.into_iter(), f);
145        c
146    }
147
148    pub fn union(&self, other: Self) -> Self
149    where K: Clone, V: Clone
150    {
151        self.union_with(other, &|_, v| v)
152    }
153
154    pub fn intersection_with<VV, FF>(&self, other: Self, f: &FF) -> Ctx<K, VV>
155    where
156        K: Clone,
157        V: Clone,
158        VV: Clone,
159        FF: Fn(V, V) -> VV
160    {
161        let mut diff = Ctx::new();
162        for (k, v1) in self.0.clone().into_iter() {
163            if let Some(v2) = other.0.get(&k) {
164                diff.insert(k, f(v1.clone(), v2.clone()));
165            }
166        }
167        diff
168    }
169    pub fn is_empty(&self) -> bool {
170        self.0.is_empty()
171    }
172    pub fn get(&self, k: &K) -> Option<&V> {
173        self.0.get(k)
174    }
175
176    pub fn get_mut(&mut self, k: &K) -> Option<&mut V> {
177        self.0.get_mut(k)
178    }
179
180    pub fn remove(&mut self, k: &K) -> Option<V>
181    where
182        V: Clone,
183    {
184        self.0.remove(k)
185    }
186
187    pub fn keys(&self) -> Set<&K> {
188        Set::from(self.0.keys().collect::<Vec<_>>())
189    }
190    pub fn contains(&self, k: &K) -> bool {
191        self.0.contains_key(k)
192    }
193    pub fn iter(&self) -> CtxIterator<K, V> {
194        CtxIterator {
195            iter: self.0.iter(),
196        }
197    }
198    pub fn iter_mut(&mut self) -> impl Iterator<Item = (&K, &mut V)> {
199        self.0.iter_mut()
200    }
201    pub fn modify<F>(&mut self, mut f: F)
202    where
203        F: FnMut(&K, &mut V)
204    {
205        for (k, v) in self.iter_mut() {
206            f(k, v);
207        }
208    }
209    pub fn retain(&mut self, f: impl Fn(&K, &mut V) -> bool) {
210        self.0.retain(f);
211    }
212
213    pub fn extract_if(&mut self, f: impl Fn(&K, &V) -> bool) -> Ctx<K, V>
214    where
215        K: Clone,
216        V: Clone,
217    {
218        let mut c = Ctx::new();
219        for (k, v) in self.0.clone().into_iter() {
220            if f(&k, &v) {
221                c.insert(k, v);
222            }
223        }
224        self.retain(|k, _| !c.contains(k));
225        c
226    }
227}
228
229impl<K: Ord, V> Default for Ctx<K, V> {
230    fn default() -> Self {
231        Ctx(BTreeMap::new())
232    }
233}
234
235/// From instance
236impl<X, Y, K: From<X> + Ord, V: From<Y>, const N: usize> From<[(X, Y); N]> for Ctx<K, V> {
237    fn from(v: [(X, Y); N]) -> Self {
238        Ctx(BTreeMap::from_iter(v.into_iter().map(|(k, v)| (K::from(k), V::from(v)))))
239    }
240}
241
242///////////////////////////////////////////////////////////////////////////////////
243// A set of values with Pretty and Display traits and other useful methods
244///////////////////////////////////////////////////////////////////////////////////
245#[derive(Debug, Clone, PartialEq, Eq, Hash)]
246pub struct Set<V>(BTreeSet<V>);
247
248/// Pretty printer instance
249impl<'a, D, A, V> Pretty<'a, D, A> for Set<V>
250where
251    D: DocAllocator<'a, A>,
252    D::Doc: Clone,
253    A: 'a + Clone,
254    V: Pretty<'a, D, A> + Clone
255{
256    fn pretty(self, allocator: &'a D) -> DocBuilder<'a, D, A> {
257        allocator.concat([
258            allocator.text("{"),
259            allocator.intersperse(
260                self.0.into_iter()
261                    .map(|k| k.pretty(allocator)), ", "),
262            allocator.text("}"),
263        ])
264    }
265}
266
267impl<V> IntoIterator for Set<V> {
268    type Item = V;
269    type IntoIter = std::collections::btree_set::IntoIter<V>;
270
271    fn into_iter(self) -> Self::IntoIter {
272        self.0.into_iter()
273    }
274}
275
276// Iterator for borrowing key-value pairs
277pub struct SetIterator<'a, V> {
278    iter: std::collections::btree_set::Iter<'a, V>,
279}
280
281impl<'a, V> Iterator for SetIterator<'a, V> {
282    type Item = &'a V;
283
284    fn next(&mut self) -> Option<Self::Item> {
285        self.iter.next()
286    }
287}
288
289impl<V> FromIterator<V> for Set<V>
290where
291    V: Ord,
292{
293    fn from_iter<I: IntoIterator<Item = V>>(iter: I) -> Self {
294        Set(BTreeSet::from_iter(iter))
295    }
296}
297
298impl<V> Default for Set<V> {
299    fn default() -> Self {
300        Set(BTreeSet::new())
301    }
302}
303
304impl<'a, V> fmt::Display for Set<V>
305where
306    V: Pretty<'a, BoxAllocator, ()> + Clone
307{
308    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
309        <Set<_> as Pretty<'_, BoxAllocator, ()>>::pretty(self.clone(), &BoxAllocator)
310                .1
311                .render_fmt(80, f)
312    }
313}
314
315impl<V: Ord, const N: usize> From<[V; N]> for Set<V> {
316    fn from(v: [V; N]) -> Self {
317        Set(BTreeSet::from_iter(v.into_iter()))
318    }
319}
320
321impl<V: Ord> From<Vec<V>> for Set<V> {
322    fn from(v: Vec<V>) -> Self {
323        Set(BTreeSet::from_iter(v.into_iter()))
324    }
325}
326
327impl<V: Ord> Set<V> {
328    pub fn new() -> Self where V: Ord {
329        Set(BTreeSet::new())
330    }
331    pub fn singleton(v: V) -> Self {
332        Set(BTreeSet::from([v]))
333    }
334    pub fn insert_with(&mut self, k: V, f: impl Fn(V) -> V) {
335        if self.0.remove(&k) {
336            self.0.insert(f(k));
337        } else {
338            self.0.insert(k);
339        }
340    }
341
342    pub fn insert(&mut self, k: V) {
343        self.0.insert(k);
344    }
345
346    pub fn append_with<FF, It>(&mut self, it: It, f: &FF) -> &mut Self
347    where
348        FF: Fn(V) -> V,
349        It: Iterator<Item = V>,
350    {
351        for k in it {
352            self.insert_with(k, f);
353        }
354        self
355    }
356
357    pub fn pop_first(&mut self) -> Option<V> {
358        self.0.pop_first()
359    }
360
361    pub fn append<It>(&mut self, it: It) -> &mut Self
362    where
363        It: Iterator<Item = V>,
364    {
365        for k in it {
366            self.0.insert(k);
367        }
368        self
369    }
370    pub fn first(&self) -> Option<&V> {
371        self.0.iter().next()
372    }
373    pub fn contains(&self, k: &V) -> bool {
374        self.0.contains(k)
375    }
376    pub fn is_empty(&self) -> bool {
377        self.0.is_empty()
378    }
379    /// Union of two sets with conflict hanlder [f]
380    pub fn union_with<FF>(&self, other: Set<V>, f: &FF) -> Self
381    where
382        V: Clone,
383        FF: Fn(V) -> V
384    {
385        let mut c = self.clone();
386        c.append_with(other.into_iter(), f);
387        c
388    }
389
390    /// Union of two sets
391    pub fn union(&self, other: Set<V>) -> Self
392    where
393        V: Clone
394    {
395        let mut c = self.clone();
396        c.append(other.into_iter());
397        c
398    }
399    /// Intersection of two sets
400    pub fn intersection(&self, other: Set<V>) -> Self
401    where
402        V: Clone
403    {
404        Set(self.0.intersection(&other.0).cloned().collect())
405    }
406
407    pub fn iter(&self) -> SetIterator<V> {
408        SetIterator {
409            iter: self.0.iter(),
410        }
411    }
412
413    // Method to iterate over a mutable Vec, modify elements, and return a new Set<T>
414    pub fn modify<F>(&mut self, mut f: F)
415    where
416        V: Clone,
417        F: FnMut(&mut V),
418    {
419        // Convert BTreeSet<V> to Vec<V>
420        let mut vec: Vec<V> = self.0.iter().cloned().collect();
421
422        // Apply modification function to each element
423        for elem in vec.iter_mut() {
424            f(elem);
425        }
426        // Collect back into Set<V>
427        self.0 = vec.into_iter().collect();
428    }
429
430    // Extract elements from the map that satisfy a predicate
431    pub fn extract_if(&mut self, f: impl Fn(&V) -> bool) -> Set<V>
432    where
433        V: Clone
434    {
435        self.0.extract_if(f).collect()
436    }
437
438    pub fn retain(&mut self, f: impl Fn(&V) -> bool) {
439        self.0.retain(f);
440    }
441}
442
443/// Arbitrary instance for Set
444#[cfg(test)]
445impl<'a, K: Arbitrary<'a> + Ord> Arbitrary<'a> for Set<K> {
446    fn arbitrary(u: &mut Unstructured<'a>) -> arbitrary::Result<Self, arbitrary::Error> {
447        let n: u8 = u.int_in_range(0..=9)?;
448        let mut vars = BTreeSet::new();
449        for _ in 0..n {
450            vars.insert(K::arbitrary(u)?);
451        }
452        Ok(Set(vars))
453    }
454}