abstutil/
collections.rs

1use std::cmp::Ord;
2use std::collections::{BTreeMap, BTreeSet};
3use std::marker::PhantomData;
4
5use anyhow::Result;
6
7use itertools::Itertools;
8use serde::{Deserialize, Serialize};
9
10// TODO Ideally derive Serialize and Deserialize, but I can't seem to express the lifetimes
11// correctly.
12#[derive(PartialEq, Clone)]
13pub struct MultiMap<K, V>
14where
15    K: Ord + PartialEq + Clone,
16    V: Ord + PartialEq + Clone,
17{
18    map: BTreeMap<K, BTreeSet<V>>,
19    empty: BTreeSet<V>,
20}
21
22impl<K, V> MultiMap<K, V>
23where
24    K: Ord + PartialEq + Clone,
25    V: Ord + PartialEq + Clone,
26{
27    pub fn new() -> MultiMap<K, V> {
28        MultiMap {
29            map: BTreeMap::new(),
30            empty: BTreeSet::new(),
31        }
32    }
33
34    pub fn insert(&mut self, key: K, value: V) {
35        self.map
36            .entry(key)
37            .or_insert_with(BTreeSet::new)
38            .insert(value);
39    }
40
41    pub fn remove(&mut self, key: K, value: V) {
42        if !self.map.contains_key(&key) {
43            return;
44        }
45        self.map.get_mut(&key).unwrap().remove(&value);
46        if self.map[&key].is_empty() {
47            self.map.remove(&key);
48        }
49    }
50
51    pub fn get(&self, key: K) -> &BTreeSet<V> {
52        self.map.get(&key).unwrap_or(&self.empty)
53    }
54
55    pub fn set(&mut self, key: K, values: BTreeSet<V>) {
56        self.map.insert(key, values);
57    }
58
59    pub fn len(&self) -> usize {
60        self.map.len()
61    }
62
63    pub fn is_empty(&self) -> bool {
64        self.map.is_empty()
65    }
66
67    #[allow(clippy::should_implement_trait)]
68    pub fn borrow(&self) -> &BTreeMap<K, BTreeSet<V>> {
69        &self.map
70    }
71
72    pub fn consume(self) -> BTreeMap<K, BTreeSet<V>> {
73        self.map
74    }
75}
76impl<K, V> Default for MultiMap<K, V>
77where
78    K: Ord + PartialEq + Clone,
79    V: Ord + PartialEq + Clone,
80{
81    fn default() -> MultiMap<K, V> {
82        MultiMap::new()
83    }
84}
85
86/// A counter per key
87// Be careful with PartialEq -- some entries may have an explicit 0, others not
88#[derive(Serialize, Deserialize, Clone, PartialEq)]
89pub struct Counter<T: Ord + PartialEq + Clone> {
90    map: BTreeMap<T, usize>,
91    sum: usize,
92}
93
94impl<T: Ord + PartialEq + Clone> Default for Counter<T> {
95    fn default() -> Counter<T> {
96        Counter::new()
97    }
98}
99
100impl<T: Ord + PartialEq + Clone> Counter<T> {
101    pub fn new() -> Counter<T> {
102        Counter {
103            map: BTreeMap::new(),
104            sum: 0,
105        }
106    }
107
108    pub fn add(&mut self, val: T, amount: usize) -> usize {
109        let entry = self.map.entry(val).or_insert(0);
110        *entry += amount;
111        self.sum += amount;
112        *entry
113    }
114    pub fn inc(&mut self, val: T) -> usize {
115        self.add(val, 1)
116    }
117
118    /// If the key is missing, returns 0
119    pub fn get(&self, val: T) -> usize {
120        self.map.get(&val).cloned().unwrap_or(0)
121    }
122
123    /// Values with the same count are grouped together
124    pub fn sorted_asc(&self) -> Vec<Vec<T>> {
125        let mut list = self.map.iter().collect::<Vec<_>>();
126        list.sort_by_key(|(_, cnt)| *cnt);
127        list.into_iter()
128            .group_by(|(_, cnt)| *cnt)
129            .into_iter()
130            .map(|(_, group)| group.into_iter().map(|(val, _)| val.clone()).collect())
131            .collect()
132    }
133
134    pub fn highest_n(&self, n: usize) -> Vec<(T, usize)> {
135        let mut list: Vec<(T, usize)> = self
136            .map
137            .iter()
138            .map(|(key, cnt)| (key.clone(), *cnt))
139            .collect();
140        list.sort_by_key(|(_, cnt)| *cnt);
141        list.reverse();
142        list.truncate(n);
143        list
144    }
145
146    /// If two keys share the maximum, returns one of them arbitrarily (and deterministically)
147    pub fn max_key(&self) -> T {
148        self.map.iter().max_by_key(|pair| pair.1).unwrap().0.clone()
149    }
150
151    pub fn max(&self) -> usize {
152        self.map.values().max().cloned().unwrap_or(0)
153    }
154    pub fn sum(&self) -> usize {
155        self.sum
156    }
157
158    pub fn compare(mut self, mut other: Counter<T>) -> Vec<(T, usize, usize)> {
159        for key in self.map.keys() {
160            other.map.entry(key.clone()).or_insert(0);
161        }
162        for key in other.map.keys() {
163            self.map.entry(key.clone()).or_insert(0);
164        }
165        self.map
166            .into_iter()
167            .map(|(k, cnt)| (k.clone(), cnt, other.map[&k]))
168            .collect()
169    }
170
171    #[allow(clippy::should_implement_trait)]
172    pub fn borrow(&self) -> &BTreeMap<T, usize> {
173        &self.map
174    }
175    pub fn consume(self) -> BTreeMap<T, usize> {
176        self.map
177    }
178
179    pub fn is_empty(&self) -> bool {
180        self.map.is_empty()
181    }
182
183    pub fn extend(&mut self, other: Counter<T>) {
184        self.map.extend(other.map);
185        self.sum += other.sum;
186    }
187
188    /// Remove all entries that aren't in the specified set of keys
189    pub fn subset(&mut self, keys: &BTreeSet<T>) {
190        let mut sum = 0;
191        self.map.retain(|k, v| {
192            if keys.contains(k) {
193                true
194            } else {
195                sum += *v;
196                false
197            }
198        });
199        self.sum -= sum;
200    }
201}
202
203pub fn wraparound_get<T>(vec: &[T], idx: isize) -> &T {
204    let len = vec.len() as isize;
205    let idx = idx % len;
206    let idx = if idx >= 0 { idx } else { idx + len };
207    &vec[idx as usize]
208}
209
210pub fn contains_duplicates<T: Ord>(vec: &[T]) -> bool {
211    let mut set = BTreeSet::new();
212    for item in vec {
213        if set.contains(item) {
214            return true;
215        }
216        set.insert(item);
217    }
218    false
219}
220
221/// Use when your key is just PartialEq, not Ord or Hash.
222pub struct VecMap<K, V> {
223    inner: Vec<(K, V)>,
224}
225
226impl<K: Clone + PartialEq, V> VecMap<K, V> {
227    pub fn new() -> VecMap<K, V> {
228        VecMap { inner: Vec::new() }
229    }
230
231    pub fn consume(self) -> Vec<(K, V)> {
232        self.inner
233    }
234
235    pub fn mut_or_insert<F: Fn() -> V>(&mut self, key: K, ctor: F) -> &mut V {
236        if let Some(idx) = self.inner.iter().position(|(k, _)| key == *k) {
237            return &mut self.inner[idx].1;
238        }
239        self.inner.push((key, ctor()));
240        &mut self.inner.last_mut().unwrap().1
241    }
242
243    /// Doesn't dedupe
244    pub fn push(&mut self, key: K, value: V) {
245        self.inner.push((key, value));
246    }
247
248    pub fn get(&self, key: &K) -> Option<&V> {
249        for (k, v) in &self.inner {
250            if k == key {
251                return Some(v);
252            }
253        }
254        None
255    }
256
257    pub fn len(&self) -> usize {
258        self.inner.len()
259    }
260
261    pub fn is_empty(&self) -> bool {
262        self.inner.is_empty()
263    }
264
265    pub fn clear(&mut self) {
266        self.inner.clear();
267    }
268}
269
270impl<K: Clone + PartialEq, V> Default for VecMap<K, V> {
271    fn default() -> Self {
272        VecMap::new()
273    }
274}
275
276/// Convenience functions around a string->string map
277#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
278pub struct Tags(BTreeMap<String, String>);
279
280impl Tags {
281    pub fn new(map: BTreeMap<String, String>) -> Tags {
282        Tags(map)
283    }
284
285    pub fn empty() -> Tags {
286        Tags(BTreeMap::new())
287    }
288
289    pub fn get(&self, k: &str) -> Option<&String> {
290        self.0.get(k)
291    }
292
293    pub fn get_result(&self, k: &str) -> Result<&String> {
294        self.0.get(k).ok_or_else(|| anyhow!("missing {}", k))
295    }
296
297    pub fn contains_key(&self, k: &str) -> bool {
298        self.0.contains_key(k)
299    }
300
301    pub fn has_any(&self, keys: Vec<&str>) -> bool {
302        keys.into_iter().any(|key| self.contains_key(key))
303    }
304
305    pub fn is(&self, k: &str, v: &str) -> bool {
306        self.0.get(k) == Some(&v.to_string())
307    }
308
309    pub fn is_any(&self, k: &str, values: Vec<&str>) -> bool {
310        if let Some(v) = self.0.get(k) {
311            values.contains(&v.as_ref())
312        } else {
313            false
314        }
315    }
316
317    pub fn insert<K: Into<String>, V: Into<String>>(&mut self, k: K, v: V) {
318        self.0.insert(k.into(), v.into());
319    }
320    pub fn remove(&mut self, k: &str) -> Option<String> {
321        self.0.remove(k)
322    }
323
324    pub fn is_empty(&self) -> bool {
325        self.0.is_empty()
326    }
327
328    // TODO Really just iter()
329    pub fn inner(&self) -> &BTreeMap<String, String> {
330        &self.0
331    }
332
333    pub fn into_inner(self) -> BTreeMap<String, String> {
334        self.0
335    }
336
337    /// Find all values that differ. Returns (key, value1, value2). If one set of tags is missing a
338    /// value, return a blank string.
339    pub fn diff(&self, other: &Tags) -> Vec<(String, String, String)> {
340        let mut results = Vec::new();
341        for (k, v1) in self.inner() {
342            let v2 = other.get(k).cloned().unwrap_or_else(String::new);
343            if v1 != &v2 {
344                results.push((k.clone(), v1.clone(), v2));
345            }
346        }
347        for (k, v2) in other.inner() {
348            if !self.contains_key(k) {
349                results.push((k.clone(), String::new(), v2.clone()));
350            }
351        }
352        results
353    }
354}
355
356/// Use with `FixedMap`. From a particular key, extract a `usize`. These values should be
357/// roughly contiguous; the space used by the `FixedMap` will be `O(n)` with respect to the largest
358/// value returned here.
359pub trait IndexableKey {
360    fn index(&self) -> usize;
361}
362
363/// A drop-in replacement for `BTreeMap`, where the keys have the property of being array indices.
364/// Some values may be missing. Much more efficient at operations on individual objects, because
365/// it just becomes a simple array lookup.
366#[derive(Serialize, Deserialize, Clone)]
367pub struct FixedMap<K: IndexableKey, V> {
368    inner: Vec<Option<V>>,
369    key_type: PhantomData<K>,
370}
371
372impl<K: IndexableKey, V> FixedMap<K, V> {
373    pub fn new() -> FixedMap<K, V> {
374        FixedMap {
375            inner: Vec::new(),
376            key_type: PhantomData,
377        }
378    }
379
380    pub fn insert(&mut self, key: K, value: V) {
381        let idx = key.index();
382        // Depending on the order of calls, this could wind up pushing one value at a time. It may
383        // be more efficient to resize less times and allocate more, but it'll require the caller
384        // to know about how many values it'll need.
385        if idx >= self.inner.len() {
386            self.inner.resize_with(idx + 1, || None);
387        }
388        self.inner[idx] = Some(value);
389    }
390
391    pub fn get(&self, key: &K) -> Option<&V> {
392        self.inner.get(key.index())?.as_ref()
393    }
394
395    pub fn get_mut(&mut self, key: &K) -> Option<&mut V> {
396        self.inner.get_mut(key.index())?.as_mut()
397    }
398
399    pub fn contains_key(&self, key: &K) -> bool {
400        self.inner
401            .get(key.index())
402            .map(|x| x.is_some())
403            .unwrap_or(false)
404    }
405
406    pub fn remove(&mut self, key: &K) -> Option<V> {
407        self.inner.get_mut(key.index())?.take()
408    }
409
410    pub fn values(&self) -> std::iter::Flatten<std::slice::Iter<'_, std::option::Option<V>>> {
411        self.inner.iter().flatten()
412    }
413}
414
415impl<K: IndexableKey, V> Default for FixedMap<K, V> {
416    fn default() -> Self {
417        FixedMap::new()
418    }
419}
420
421impl<K: IndexableKey, V> std::ops::Index<&K> for FixedMap<K, V> {
422    type Output = V;
423
424    fn index(&self, key: &K) -> &Self::Output {
425        self.inner[key.index()].as_ref().unwrap()
426    }
427}
abstutil/collections.rs

abstutil/
collections.rs
