mz_ore/
collections.rs

Help
1// Copyright Materialize, Inc. and contributors. All rights reserved.
2//
3// Licensed under the Apache License, Version 2.0 (the "License");
4// you may not use this file except in compliance with the License.
5// You may obtain a copy of the License in the LICENSE file at the
6// root of this repository, or online at
7//
8//     http://www.apache.org/licenses/LICENSE-2.0
9//
10// Unless required by applicable law or agreed to in writing, software
11// distributed under the License is distributed on an "AS IS" BASIS,
12// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13// See the License for the specific language governing permissions and
14// limitations under the License.
15
16//! Collection utilities.
17
18use std::collections::BTreeMap;
19use std::collections::btree_map::Entry as BEntry;
20use std::collections::hash_map::Entry as HEntry;
21use std::fmt::{Debug, Display};
22
23mod hash;
24
25pub use crate::collections::hash::{HashMap, HashSet};
26
27/// Extension methods for collections.
28pub trait CollectionExt<T>: Sized
29where
30    T: IntoIterator,
31{
32    /// Consumes the collection and returns its first element.
33    ///
34    /// This method panics if the collection does not have at least one element.
35    fn into_first(self) -> T::Item;
36
37    /// Consumes the collection and returns its last element.
38    ///
39    /// This method panics if the collection does not have at least one element.
40    fn into_last(self) -> T::Item;
41
42    /// Consumes the collection and returns its only element.
43    ///
44    /// This method panics if the collection does not have exactly one element.
45    fn into_element(self) -> T::Item {
46        self.expect_element(|| "into_element called on collection without exactly one element")
47    }
48
49    /// Consumes the collection and returns its only element.
50    ///
51    /// This method panics with the given error function's return value if the collection does not
52    /// have exactly one element.
53    fn expect_element<Err: Display>(self, msg_fn: impl FnOnce() -> Err) -> T::Item;
54}
55
56impl<T> CollectionExt<T> for T
57where
58    T: IntoIterator,
59{
60    fn into_first(self) -> T::Item {
61        self.into_iter().next().unwrap()
62    }
63
64    fn into_last(self) -> T::Item {
65        self.into_iter().last().unwrap()
66    }
67
68    fn expect_element<Err: Display>(self, msg_fn: impl FnOnce() -> Err) -> T::Item {
69        let mut iter = self.into_iter();
70        match (iter.next(), iter.next()) {
71            (Some(el), None) => el,
72            _ => panic!("{}", msg_fn()),
73        }
74    }
75}
76
77/// Extension methods for associative collections.
78pub trait AssociativeExt<K, V> {
79    /// Inserts a key and value, panicking with
80    /// a given message if a true
81    /// insert (as opposed to an update) cannot be done
82    /// because the key already existed in the collection.
83    fn expect_insert(&mut self, k: K, v: V, msg: &str);
84    /// Inserts a key and value, panicking if a true
85    /// insert (as opposed to an update) cannot be done
86    /// because the key already existed in the collection.
87    fn unwrap_insert(&mut self, k: K, v: V) {
88        self.expect_insert(k, v, "called `unwrap_insert` for an already-existing key")
89    }
90
91    /// Removes a key, panicking with
92    /// a given message if a true
93    /// removal (as opposed to a no-op) cannot be done
94    /// because the key does not exist in the collection.
95    fn expect_remove(&mut self, k: &K, msg: &str) -> V;
96    /// Removes a key, panicking if a true
97    /// removal (as opposed to a no-op) cannot be done
98    /// because the key does not exist in the collection.
99    fn unwrap_remove(&mut self, k: &K) -> V {
100        self.expect_remove(k, "called `unwrap_remove` for a non-existing key")
101    }
102}
103
104impl<K, V> AssociativeExt<K, V> for HashMap<K, V>
105where
106    K: Eq + std::hash::Hash + Debug,
107    V: Debug,
108{
109    fn expect_insert(&mut self, k: K, v: V, msg: &str) {
110        match self.entry(k) {
111            HEntry::Vacant(e) => {
112                e.insert(v);
113            }
114            HEntry::Occupied(e) => {
115                panic!(
116                    "{} (key: {:?}, old value: {:?}, new value: {:?})",
117                    msg,
118                    e.key(),
119                    e.get(),
120                    v
121                )
122            }
123        }
124    }
125
126    fn expect_remove(&mut self, k: &K, msg: &str) -> V {
127        match self.remove(k) {
128            Some(v) => v,
129            None => panic!("{} (key: {:?})", msg, k),
130        }
131    }
132}
133
134impl<K, V> AssociativeExt<K, V> for BTreeMap<K, V>
135where
136    K: Ord + Debug,
137    V: Debug,
138{
139    fn expect_insert(&mut self, k: K, v: V, msg: &str) {
140        match self.entry(k) {
141            BEntry::Vacant(e) => {
142                e.insert(v);
143            }
144            BEntry::Occupied(e) => {
145                panic!(
146                    "{} (key: {:?}, old value: {:?}, new value: {:?})",
147                    msg,
148                    e.key(),
149                    e.get(),
150                    v
151                )
152            }
153        }
154    }
155
156    fn expect_remove(&mut self, k: &K, msg: &str) -> V {
157        match self.remove(k) {
158            Some(v) => v,
159            None => panic!("{} (key: {:?})", msg, k),
160        }
161    }
162}
mz_ore/collections.rs

mz_ore/
collections.rs
