1use std::collections::{BTreeMap, BTreeSet};
2use std::fmt;
3use pretty::{Pretty, DocAllocator, DocBuilder, BoxAllocator};
4
5#[derive(Clone, PartialEq, Eq, PartialOrd, Ord, Hash, Debug)]
7pub struct Ctx<K, V>(BTreeMap<K, V>);
8
9impl<'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
34impl<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
44pub struct CtxIterator<'a, K, V> {
46 iter: std::collections::btree_map::Iter<'a, K, V>,
47}
48
49impl<'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
67impl<'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#[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
94impl<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
235impl<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#[derive(Debug, Clone, PartialEq, Eq, Hash)]
246pub struct Set<V>(BTreeSet<V>);
247
248impl<'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
276pub 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 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 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 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 pub fn modify<F>(&mut self, mut f: F)
415 where
416 V: Clone,
417 F: FnMut(&mut V),
418 {
419 let mut vec: Vec<V> = self.0.iter().cloned().collect();
421
422 for elem in vec.iter_mut() {
424 f(elem);
425 }
426 self.0 = vec.into_iter().collect();
428 }
429
430 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#[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}