starlark_map/
sorted_map.rs

1/*
2 * Copyright 2019 The Starlark in Rust Authors.
3 * Copyright (c) Facebook, Inc. and its affiliates.
4 *
5 * Licensed under the Apache License, Version 2.0 (the "License");
6 * you may not use this file except in compliance with the License.
7 * You may obtain a copy of the License at
8 *
9 *     https://www.apache.org/licenses/LICENSE-2.0
10 *
11 * Unless required by applicable law or agreed to in writing, software
12 * distributed under the License is distributed on an "AS IS" BASIS,
13 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
14 * See the License for the specific language governing permissions and
15 * limitations under the License.
16 */
17
18//! `SmallMap` which asserts that its elements are sorted.
19
20use std::hash::Hash;
21
22use allocative::Allocative;
23use serde::Deserialize;
24use serde::Serialize;
25
26use crate::ordered_map::OrderedMap;
27use crate::small_map;
28use crate::small_map::SmallMap;
29use crate::Equivalent;
30
31/// `IndexMap` but with keys sorted.
32#[derive(Debug, Clone, Eq, PartialEq, Hash, Ord, PartialOrd, Allocative)]
33pub struct SortedMap<K, V> {
34    map: OrderedMap<K, V>,
35}
36
37impl<K: Ord + Hash, V> Default for SortedMap<K, V> {
38    #[inline]
39    fn default() -> Self {
40        SortedMap {
41            map: OrderedMap::default(),
42        }
43    }
44}
45
46impl<K, V> SortedMap<K, V>
47where
48    K: Ord + Hash,
49{
50    /// Construct an empty `SortedMap`.
51    #[inline]
52    pub const fn new() -> SortedMap<K, V> {
53        SortedMap {
54            map: OrderedMap::new(),
55        }
56    }
57
58    /// Iterate over the entries.
59    #[inline]
60    pub fn iter(&self) -> impl ExactSizeIterator<Item = (&K, &V)> {
61        self.map.iter()
62    }
63
64    /// Iterate over the entries, with mutable values.
65    #[inline]
66    pub fn iter_mut(&mut self) -> impl ExactSizeIterator<Item = (&K, &mut V)> {
67        self.map.iter_mut()
68    }
69
70    /// Iterate over the keys.
71    #[inline]
72    pub fn keys(&self) -> impl ExactSizeIterator<Item = &K> {
73        self.map.keys()
74    }
75
76    /// Iterate over the values.
77    #[inline]
78    pub fn values(&self) -> impl ExactSizeIterator<Item = &V> {
79        self.map.values()
80    }
81
82    /// Iterate over the values mutably.
83    #[inline]
84    pub fn values_mut(&mut self) -> impl ExactSizeIterator<Item = &mut V> {
85        self.map.values_mut()
86    }
87
88    /// Return the number of elements in the map.
89    #[inline]
90    pub fn len(&self) -> usize {
91        self.map.len()
92    }
93
94    /// Check if the map is empty.
95    #[inline]
96    pub fn is_empty(&self) -> bool {
97        self.map.is_empty()
98    }
99
100    /// Get a reference to the value associated with the given key.
101    #[inline]
102    pub fn get<Q: ?Sized>(&self, key: &Q) -> Option<&V>
103    where
104        Q: Hash + Equivalent<K>,
105    {
106        self.map.get(key)
107    }
108
109    /// Get a mutable reference to the value associated with the given key.
110    #[inline]
111    pub fn get_mut<Q: ?Sized>(&mut self, key: &Q) -> Option<&mut V>
112    where
113        Q: Hash + Equivalent<K>,
114    {
115        self.map.get_mut(key)
116    }
117
118    /// Check if the map contains the given key.
119    #[inline]
120    pub fn contains_key<Q>(&self, k: &Q) -> bool
121    where
122        Q: Hash + Equivalent<K> + ?Sized,
123    {
124        self.map.contains_key(k)
125    }
126
127    /// Iterate over the map with hashes.
128    #[inline]
129    pub fn iter_hashed(&self) -> small_map::IterHashed<K, V> {
130        self.map.iter_hashed()
131    }
132}
133
134impl<K: Ord + Hash, V> FromIterator<(K, V)> for SortedMap<K, V> {
135    #[inline]
136    fn from_iter<T: IntoIterator<Item = (K, V)>>(iter: T) -> Self {
137        let map = OrderedMap::from_iter(iter);
138        SortedMap::from(map)
139    }
140}
141
142impl<K: Ord + Hash, V> From<OrderedMap<K, V>> for SortedMap<K, V> {
143    #[inline]
144    fn from(mut map: OrderedMap<K, V>) -> SortedMap<K, V> {
145        map.sort_keys();
146        SortedMap { map }
147    }
148}
149
150impl<K: Ord + Hash, V> From<SmallMap<K, V>> for SortedMap<K, V> {
151    #[inline]
152    fn from(map: SmallMap<K, V>) -> SortedMap<K, V> {
153        // `OrderedMap: From<SmallMap>` is trivial, so this does not do any extra work
154        SortedMap::from(OrderedMap::from(map))
155    }
156}
157
158impl<K: Ord + Hash, V> IntoIterator for SortedMap<K, V> {
159    type Item = (K, V);
160    type IntoIter = small_map::IntoIter<K, V>;
161
162    #[inline]
163    fn into_iter(self) -> Self::IntoIter {
164        self.map.into_iter()
165    }
166}
167
168impl<'a, K: Ord + Hash, V> IntoIterator for &'a SortedMap<K, V> {
169    type Item = (&'a K, &'a V);
170    type IntoIter = small_map::Iter<'a, K, V>;
171
172    #[inline]
173    fn into_iter(self) -> Self::IntoIter {
174        self.map.iter()
175    }
176}
177
178impl<'a, K: Ord + Hash, V> IntoIterator for &'a mut SortedMap<K, V> {
179    type Item = (&'a K, &'a mut V);
180    type IntoIter = small_map::IterMut<'a, K, V>;
181
182    #[inline]
183    fn into_iter(self) -> Self::IntoIter {
184        self.map.iter_mut()
185    }
186}
187
188impl<K: Serialize, V: Serialize> Serialize for SortedMap<K, V> {
189    fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
190    where
191        S: serde::Serializer,
192    {
193        self.map.serialize(serializer)
194    }
195}
196
197impl<'de, K, V> Deserialize<'de> for SortedMap<K, V>
198where
199    K: Deserialize<'de> + Hash + Eq,
200    V: Deserialize<'de>,
201{
202    fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
203    where
204        D: serde::Deserializer<'de>,
205    {
206        Ok(Self {
207            map: OrderedMap::deserialize(deserializer)?,
208        })
209    }
210}
211
212#[cfg(test)]
213mod tests {
214    use crate::sorted_map::SortedMap;
215
216    #[test]
217    fn test_from_iter() {
218        let map = SortedMap::from_iter([(1, 2), (5, 6), (3, 4)]);
219        assert_eq!(
220            vec![(&1, &2), (&3, &4), (&5, &6)],
221            map.iter().collect::<Vec<_>>()
222        );
223    }
224
225    #[test]
226    fn test_value_modification() {
227        let mut map = SortedMap::from_iter([(1, vec![1, 2, 3]), (2, vec![4]), (3, vec![5])]);
228        let mut keys = map.keys().collect::<Vec<_>>();
229        keys.sort();
230        assert_eq!(keys, map.keys().collect::<Vec<_>>(),);
231        // Support insertion for existing keys
232        map.get_mut(&1).unwrap().push(11);
233        map.get_mut(&2).unwrap().push(22);
234        map.get_mut(&3).unwrap().push(33);
235
236        assert_eq!(
237            vec![
238                (&1, &vec![1, 2, 3, 11]),
239                (&2, &vec![4, 22]),
240                (&3, &vec![5, 33])
241            ],
242            map.iter().collect::<Vec<_>>()
243        );
244        let mut keys = map.keys().collect::<Vec<_>>();
245        keys.sort();
246        assert_eq!(map.keys().collect::<Vec<_>>(), keys,);
247    }
248}
starlark_map/sorted_map.rs

starlark_map/
sorted_map.rs
