Skip to main content

formality_core/
collections.rs

1//! We use btree maps and sets, and we rely on the ordering
2//! properties that they ensure.
3
4use std::collections::{BTreeMap, BTreeSet};
5
6use crate::cast::{Downcast, DowncastFrom, DowncastTo, Upcast, UpcastFrom};
7
8pub type Map<K, V> = BTreeMap<K, V>;
9pub type Set<E> = BTreeSet<E>;
10
11#[macro_export]
12macro_rules! set {
13    () => {
14        $crate::Set::new()
15    };
16
17    ($($t:tt)*) => {
18        $crate::seq![$($t)*].into_iter().collect::<$crate::Set<_>>()
19    };
20}
21
22/// Replacement for the `vec` macro that supports `..foo` notation to flatten
23/// vectors.
24#[macro_export]
25macro_rules! seq {
26    (@elem($iter:ident) .. $e:expr) => {
27        let $iter = $iter.chain($e);
28    };
29
30    (@elem($iter:ident) .. $e:expr, $($t:tt)*) => {
31        let $iter = $iter.chain($e);
32        $crate::seq!(@elem($iter) $($t)*);
33    };
34
35    (@elem($iter:ident) $e:expr, $($t:tt)*) => {
36        let $iter = $iter.chain(std::iter::once($e));
37        $crate::seq!(@elem($iter) $($t)*);
38    };
39
40    (@elem($iter:ident)) => {
41    };
42
43    (@elem($iter:ident) $($t:tt)*) => {
44        let $iter = $iter.chain(std::iter::once($($t)*));
45    };
46
47    () => {
48        vec![]
49    };
50
51    ($($t:tt)*) => {
52        {
53            let iter = std::iter::empty();
54            $crate::seq!(@elem(iter) $($t)*);
55            iter.collect::<Vec<_>>()
56        }
57    };
58
59}
60
61pub trait SetExt<T>: Sized {
62    fn split_first(self) -> Option<(T, Self)>;
63
64    fn union_with(self, other: Self) -> Self;
65
66    fn plus(self, other: T) -> Self;
67}
68
69impl<T: Ord + Clone> SetExt<T> for Set<T> {
70    fn split_first(self) -> Option<(T, Set<T>)> {
71        let mut iter = self.into_iter();
72        iter.next().map(|e| (e, iter.collect()))
73    }
74
75    fn union_with(mut self, other: Self) -> Self {
76        for item in other {
77            self.insert(item);
78        }
79        self
80    }
81
82    fn plus(mut self, other: T) -> Self {
83        self.insert(other);
84        self
85    }
86}
87
88impl<T> DowncastTo<()> for Set<T> {
89    fn downcast_to(&self) -> Option<()> {
90        if self.is_empty() {
91            Some(())
92        } else {
93            None
94        }
95    }
96}
97
98impl<T> DowncastTo<()> for Vec<T> {
99    fn downcast_to(&self) -> Option<()> {
100        if self.is_empty() {
101            Some(())
102        } else {
103            None
104        }
105    }
106}
107
108macro_rules! tuple_upcast {
109    ($($name:ident),*) => {
110        #[allow(non_snake_case)]
111        impl<$($name,)* T> UpcastFrom<($($name,)*)> for Set<T>
112        where
113            $($name: Upcast<Set<T>>,)*
114            T: Ord + Clone,
115        {
116            fn upcast_from(($($name,)*): ($($name,)*)) -> Self {
117                let c = None.into_iter();
118                $(
119                    let $name: Set<T> = $name.upcast();
120                    let c = c.chain($name);
121                )*
122                c.collect()
123            }
124        }
125    }
126}
127
128tuple_upcast!(A, B);
129tuple_upcast!(A, B, C);
130tuple_upcast!(A, B, C, D);
131
132impl<A, T> DowncastTo<(A, Set<T>)> for Set<T>
133where
134    A: DowncastFrom<T>,
135    T: Ord + Clone,
136{
137    fn downcast_to(&self) -> Option<(A, Set<T>)> {
138        if self.is_empty() {
139            None
140        } else {
141            let r = self.clone();
142            let (a, bs) = r.split_first().unwrap();
143            let a: A = a.downcast()?;
144            Some((a, bs))
145        }
146    }
147}
148
149impl<A, T> DowncastTo<(A, Vec<T>)> for Vec<T>
150where
151    A: DowncastFrom<T>,
152    T: Ord + Clone,
153{
154    fn downcast_to(&self) -> Option<(A, Vec<T>)> {
155        if self.is_empty() {
156            None
157        } else {
158            let r = self.clone();
159            let (a, bs) = r.split_first().unwrap();
160            let a: A = a.downcast()?;
161            Some((a, bs.to_vec()))
162        }
163    }
164}
165
166pub trait Deduplicate {
167    fn deduplicate(self) -> Self;
168}
169
170impl<T> Deduplicate for Vec<T>
171where
172    T: Ord + Clone,
173{
174    /// Remove duplicates of `v`
175    fn deduplicate(mut self) -> Vec<T> {
176        let mut s = Set::default();
177        self.retain(|e| s.insert(e.clone()));
178        self
179    }
180}