mz_ore/
collections.rs

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    #[track_caller]
46    fn into_element(self) -> T::Item {
47        self.expect_element(|| "into_element called on collection without exactly one element")
48    }
49
50    /// Consumes the collection and returns its only element.
51    ///
52    /// This method panics with the given error function's return value if the collection does not
53    /// have exactly one element.
54    fn expect_element<Err: Display>(self, msg_fn: impl FnOnce() -> Err) -> T::Item;
55}
56
57impl<T> CollectionExt<T> for T
58where
59    T: IntoIterator,
60{
61    #[track_caller]
62    fn into_first(self) -> T::Item {
63        self.into_iter().next().unwrap()
64    }
65
66    #[track_caller]
67    fn into_last(self) -> T::Item {
68        self.into_iter().last().unwrap()
69    }
70
71    #[track_caller]
72    fn expect_element<Err: Display>(self, msg_fn: impl FnOnce() -> Err) -> T::Item {
73        let mut iter = self.into_iter();
74        match (iter.next(), iter.next()) {
75            (Some(el), None) => el,
76            _ => panic!("{}", msg_fn()),
77        }
78    }
79}
80
81/// Extension methods for associative collections.
82pub trait AssociativeExt<K, V> {
83    /// Inserts a key and value, panicking with
84    /// a given message if a true
85    /// insert (as opposed to an update) cannot be done
86    /// because the key already existed in the collection.
87    fn expect_insert(&mut self, k: K, v: V, msg: &str);
88    /// Inserts a key and value, panicking if a true
89    /// insert (as opposed to an update) cannot be done
90    /// because the key already existed in the collection.
91    fn unwrap_insert(&mut self, k: K, v: V) {
92        self.expect_insert(k, v, "called `unwrap_insert` for an already-existing key")
93    }
94
95    /// Removes a key, panicking with
96    /// a given message 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 expect_remove(&mut self, k: &K, msg: &str) -> V;
100    /// Removes a key, panicking if a true
101    /// removal (as opposed to a no-op) cannot be done
102    /// because the key does not exist in the collection.
103    fn unwrap_remove(&mut self, k: &K) -> V {
104        self.expect_remove(k, "called `unwrap_remove` for a non-existing key")
105    }
106}
107
108impl<K, V> AssociativeExt<K, V> for HashMap<K, V>
109where
110    K: Eq + std::hash::Hash + Debug,
111    V: Debug,
112{
113    fn expect_insert(&mut self, k: K, v: V, msg: &str) {
114        match self.entry(k) {
115            HEntry::Vacant(e) => {
116                e.insert(v);
117            }
118            HEntry::Occupied(e) => {
119                panic!(
120                    "{} (key: {:?}, old value: {:?}, new value: {:?})",
121                    msg,
122                    e.key(),
123                    e.get(),
124                    v
125                )
126            }
127        }
128    }
129
130    fn expect_remove(&mut self, k: &K, msg: &str) -> V {
131        match self.remove(k) {
132            Some(v) => v,
133            None => panic!("{} (key: {:?})", msg, k),
134        }
135    }
136}
137
138impl<K, V> AssociativeExt<K, V> for BTreeMap<K, V>
139where
140    K: Ord + Debug,
141    V: Debug,
142{
143    fn expect_insert(&mut self, k: K, v: V, msg: &str) {
144        match self.entry(k) {
145            BEntry::Vacant(e) => {
146                e.insert(v);
147            }
148            BEntry::Occupied(e) => {
149                panic!(
150                    "{} (key: {:?}, old value: {:?}, new value: {:?})",
151                    msg,
152                    e.key(),
153                    e.get(),
154                    v
155                )
156            }
157        }
158    }
159
160    fn expect_remove(&mut self, k: &K, msg: &str) -> V {
161        match self.remove(k) {
162            Some(v) => v,
163            None => panic!("{} (key: {:?})", msg, k),
164        }
165    }
166}
mz_ore/collections.rs

mz_ore/
collections.rs
