zerovec/map2d/
cursor.rs

Help
1// This file is part of ICU4X. For terms of use, please see the file
2// called LICENSE at the top level of the ICU4X source tree
3// (online at: https://github.com/unicode-org/icu4x/blob/main/LICENSE ).
4
5use crate::{ZeroMap2d, ZeroSlice};
6
7use core::cmp::Ordering;
8use core::fmt;
9use core::ops::Range;
10
11use crate::map::ZeroMapKV;
12use crate::map::ZeroVecLike;
13
14use super::ZeroMap2dBorrowed;
15
16/// An intermediate state of queries over [`ZeroMap2d`] and [`ZeroMap2dBorrowed`].
17pub struct ZeroMap2dCursor<'l, 'a, K0, K1, V>
18where
19    K0: ZeroMapKV<'a>,
20    K1: ZeroMapKV<'a>,
21    V: ZeroMapKV<'a>,
22    K0: ?Sized,
23    K1: ?Sized,
24    V: ?Sized,
25{
26    // Invariant: these fields have the same invariants as they do in ZeroMap2d
27    keys0: &'l K0::Slice,
28    joiner: &'l ZeroSlice<u32>,
29    keys1: &'l K1::Slice,
30    values: &'l V::Slice,
31    // Invariant: key0_index is in range
32    key0_index: usize,
33}
34
35impl<'a, K0, K1, V> ZeroMap2dCursor<'a, 'a, K0, K1, V>
36where
37    K0: ZeroMapKV<'a>,
38    K1: ZeroMapKV<'a>,
39    V: ZeroMapKV<'a>,
40    K0: ?Sized,
41    K1: ?Sized,
42    V: ?Sized,
43{
44    /// `key0_index` must be in range
45    pub(crate) fn from_borrowed(
46        borrowed: &ZeroMap2dBorrowed<'a, K0, K1, V>,
47        key0_index: usize,
48    ) -> Self {
49        debug_assert!(key0_index < borrowed.joiner.len());
50        ZeroMap2dCursor {
51            keys0: borrowed.keys0,
52            joiner: borrowed.joiner,
53            keys1: borrowed.keys1,
54            values: borrowed.values,
55            key0_index,
56        }
57    }
58}
59
60impl<'l, 'a, K0, K1, V> ZeroMap2dCursor<'l, 'a, K0, K1, V>
61where
62    K0: ZeroMapKV<'a>,
63    K1: ZeroMapKV<'a>,
64    V: ZeroMapKV<'a>,
65    K0: ?Sized,
66    K1: ?Sized,
67    V: ?Sized,
68{
69    /// `key0_index` must be in range
70    pub(crate) fn from_cow(cow: &'l ZeroMap2d<'a, K0, K1, V>, key0_index: usize) -> Self {
71        debug_assert!(key0_index < cow.joiner.len());
72        Self {
73            keys0: cow.keys0.zvl_as_borrowed(),
74            joiner: &cow.joiner,
75            keys1: cow.keys1.zvl_as_borrowed(),
76            values: cow.values.zvl_as_borrowed(),
77            key0_index,
78        }
79    }
80
81    /// Returns the key0 corresponding to the cursor position.
82    ///
83    /// ```rust
84    /// use zerovec::ZeroMap2d;
85    ///
86    /// let mut map = ZeroMap2d::new();
87    /// map.insert("one", &1u32, "foo");
88    /// assert_eq!(map.get0("one").unwrap().key0(), "one");
89    /// ```
90    pub fn key0(&self) -> &'l K0::GetType {
91        #[allow(clippy::unwrap_used)] // safe by invariant on `self.key0_index`
92        self.keys0.zvl_get(self.key0_index).unwrap()
93    }
94
95    /// Borrow an ordered iterator over keys1 and values for a particular key0.
96    ///
97    /// To get the values as copy types, see [`Self::iter1_copied`].
98    ///
99    /// For an example, see [`ZeroMap2d::iter0()`].
100    pub fn iter1(
101        &self,
102    ) -> impl Iterator<
103        Item = (
104            &'l <K1 as ZeroMapKV<'a>>::GetType,
105            &'l <V as ZeroMapKV<'a>>::GetType,
106        ),
107    > + '_ {
108        let range = self.get_range();
109        #[allow(clippy::unwrap_used)] // `self.get_range()` returns a valid range
110        range.map(move |idx| {
111            (
112                self.keys1.zvl_get(idx).unwrap(),
113                self.values.zvl_get(idx).unwrap(),
114            )
115        })
116    }
117
118    /// Transform this cursor into an ordered iterator over keys1 for a particular key0.
119    pub fn into_iter1(
120        self,
121    ) -> impl Iterator<
122        Item = (
123            &'l <K1 as ZeroMapKV<'a>>::GetType,
124            &'l <V as ZeroMapKV<'a>>::GetType,
125        ),
126    > {
127        let range = self.get_range();
128        #[allow(clippy::unwrap_used)] // `self.get_range()` returns a valid range
129        range.map(move |idx| {
130            (
131                self.keys1.zvl_get(idx).unwrap(),
132                self.values.zvl_get(idx).unwrap(),
133            )
134        })
135    }
136
137    /// Given key0_index, returns the corresponding range of keys1, which will be valid
138    pub(super) fn get_range(&self) -> Range<usize> {
139        debug_assert!(self.key0_index < self.joiner.len());
140        let start = if self.key0_index == 0 {
141            0
142        } else {
143            #[allow(clippy::unwrap_used)] // protected by the debug_assert above
144            self.joiner.get(self.key0_index - 1).unwrap()
145        };
146        #[allow(clippy::unwrap_used)] // protected by the debug_assert above
147        let limit = self.joiner.get(self.key0_index).unwrap();
148        // These two assertions are true based on the invariants of ZeroMap2d
149        debug_assert!(start < limit);
150        debug_assert!((limit as usize) <= self.values.zvl_len());
151        (start as usize)..(limit as usize)
152    }
153}
154
155impl<'l, 'a, K0, K1, V> ZeroMap2dCursor<'l, 'a, K0, K1, V>
156where
157    K0: ZeroMapKV<'a>,
158    K1: ZeroMapKV<'a>,
159    V: ZeroMapKV<'a>,
160    K0: ?Sized,
161    K1: ?Sized,
162    V: Copy,
163{
164    /// Borrow an ordered iterator over keys1 and values for a particular key0.
165    ///
166    /// The values are returned as copy types.
167    ///
168    /// # Examples
169    ///
170    /// ```
171    /// use zerovec::ZeroMap2d;
172    ///
173    /// let zm2d: ZeroMap2d<str, u8, usize> = [
174    ///     ("a", 0u8, 1usize),
175    ///     ("b", 1u8, 1000usize),
176    ///     ("b", 2u8, 2000usize),
177    /// ]
178    /// .into_iter()
179    /// .collect();
180    ///
181    /// let mut total_value = 0;
182    ///
183    /// for cursor in zm2d.iter0() {
184    ///     for (_, value) in cursor.iter1_copied() {
185    ///         total_value += value;
186    ///     }
187    /// }
188    ///
189    /// assert_eq!(total_value, 3001);
190    /// ```
191    pub fn iter1_copied(
192        &self,
193    ) -> impl Iterator<Item = (&'l <K1 as ZeroMapKV<'a>>::GetType, V)> + '_ {
194        let range = self.get_range();
195        #[allow(clippy::unwrap_used)] // `self.get_range()` returns a valid range
196        range.map(move |idx| {
197            (
198                self.keys1.zvl_get(idx).unwrap(),
199                self.get1_copied_at(idx).unwrap(),
200            )
201        })
202    }
203
204    fn get1_copied_at(&self, index: usize) -> Option<V> {
205        let ule = self.values.zvl_get(index)?;
206        let mut result = Option::<V>::None;
207        V::Container::zvl_get_as_t(ule, |v| result.replace(*v));
208        #[allow(clippy::unwrap_used)] // `zvl_get_as_t` guarantees that the callback is invoked
209        Some(result.unwrap())
210    }
211}
212
213impl<'l, 'a, K0, K1, V> ZeroMap2dCursor<'l, 'a, K0, K1, V>
214where
215    K0: ZeroMapKV<'a>,
216    K1: ZeroMapKV<'a> + Ord,
217    V: ZeroMapKV<'a>,
218    K0: ?Sized,
219    K1: ?Sized,
220    V: ?Sized,
221{
222    /// Gets the value for a key1 from this cursor, or `None` if key1 is not in the map.
223    ///
224    /// ```rust
225    /// use zerovec::ZeroMap2d;
226    ///
227    /// let mut map = ZeroMap2d::new();
228    /// map.insert("one", &1u32, "foo");
229    /// assert_eq!(map.get0("one").unwrap().get1(&1), Some("foo"));
230    /// assert_eq!(map.get0("one").unwrap().get1(&2), None);
231    /// ```
232    pub fn get1(&self, key1: &K1) -> Option<&'l V::GetType> {
233        let key1_index = self.get_key1_index(key1)?;
234        #[allow(clippy::unwrap_used)] // key1_index is valid
235        Some(self.values.zvl_get(key1_index).unwrap())
236    }
237
238    /// Gets the value for a predicate from this cursor, or `None` if key1 is not in the map.
239    ///
240    /// ```rust
241    /// use zerovec::ZeroMap2d;
242    ///
243    /// let mut map = ZeroMap2d::new();
244    /// map.insert("one", &1u32, "foo");
245    /// assert_eq!(map.get0("one").unwrap().get1_by(|v| v.cmp(&1)), Some("foo"));
246    /// assert_eq!(map.get0("one").unwrap().get1_by(|v| v.cmp(&2)), None);
247    /// ```
248    pub fn get1_by(&self, predicate: impl FnMut(&K1) -> Ordering) -> Option<&'l V::GetType> {
249        let key1_index = self.get_key1_index_by(predicate)?;
250        #[allow(clippy::unwrap_used)] // key1_index is valid
251        Some(self.values.zvl_get(key1_index).unwrap())
252    }
253
254    /// Given key0_index and predicate, returns the index into the values array
255    fn get_key1_index_by(&self, predicate: impl FnMut(&K1) -> Ordering) -> Option<usize> {
256        let range = self.get_range();
257        debug_assert!(range.start < range.end); // '<' because every key0 should have a key1
258        debug_assert!(range.end <= self.keys1.zvl_len());
259        let start = range.start;
260        #[allow(clippy::expect_used)] // protected by the debug_assert above
261        let binary_search_result = self
262            .keys1
263            .zvl_binary_search_in_range_by(predicate, range)
264            .expect("in-bounds range");
265        binary_search_result.ok().map(move |s| s + start)
266    }
267
268    /// Given key0_index and key1, returns the index into the values array
269    fn get_key1_index(&self, key1: &K1) -> Option<usize> {
270        let range = self.get_range();
271        debug_assert!(range.start < range.end); // '<' because every key0 should have a key1
272        debug_assert!(range.end <= self.keys1.zvl_len());
273        let start = range.start;
274        #[allow(clippy::expect_used)] // protected by the debug_assert above
275        let binary_search_result = self
276            .keys1
277            .zvl_binary_search_in_range(key1, range)
278            .expect("in-bounds range");
279        binary_search_result.ok().map(move |s| s + start)
280    }
281}
282
283impl<'l, 'a, K0, K1, V> ZeroMap2dCursor<'l, 'a, K0, K1, V>
284where
285    K0: ZeroMapKV<'a>,
286    K1: ZeroMapKV<'a> + Ord,
287    V: ZeroMapKV<'a>,
288    V: Copy,
289    K0: ?Sized,
290    K1: ?Sized,
291{
292    /// For cases when `V` is fixed-size, obtain a direct copy of `V` instead of `V::ULE`
293    ///
294    /// ```rust
295    /// use zerovec::ZeroMap2d;
296    ///
297    /// let mut map: ZeroMap2d<u16, u16, u16> = ZeroMap2d::new();
298    /// map.insert(&1, &2, &3);
299    /// map.insert(&1, &4, &5);
300    /// map.insert(&6, &7, &8);
301    ///
302    /// assert_eq!(map.get0(&6).unwrap().get1_copied(&7), Some(8));
303    /// ```
304    #[inline]
305    pub fn get1_copied(&self, key1: &K1) -> Option<V> {
306        let key1_index = self.get_key1_index(key1)?;
307        self.get1_copied_at(key1_index)
308    }
309
310    /// For cases when `V` is fixed-size, obtain a direct copy of `V` instead of `V::ULE`
311    #[inline]
312    pub fn get1_copied_by(&self, predicate: impl FnMut(&K1) -> Ordering) -> Option<V> {
313        let key1_index = self.get_key1_index_by(predicate)?;
314        self.get1_copied_at(key1_index)
315    }
316}
317
318// We can't use the default PartialEq because ZeroMap2d is invariant
319// so otherwise rustc will not automatically allow you to compare ZeroMaps
320// with different lifetimes
321impl<'m, 'n, 'a, 'b, K0, K1, V> PartialEq<ZeroMap2dCursor<'n, 'b, K0, K1, V>>
322    for ZeroMap2dCursor<'m, 'a, K0, K1, V>
323where
324    K0: for<'c> ZeroMapKV<'c> + ?Sized,
325    K1: for<'c> ZeroMapKV<'c> + ?Sized,
326    V: for<'c> ZeroMapKV<'c> + ?Sized,
327    <K0 as ZeroMapKV<'a>>::Slice: PartialEq<<K0 as ZeroMapKV<'b>>::Slice>,
328    <K1 as ZeroMapKV<'a>>::Slice: PartialEq<<K1 as ZeroMapKV<'b>>::Slice>,
329    <V as ZeroMapKV<'a>>::Slice: PartialEq<<V as ZeroMapKV<'b>>::Slice>,
330{
331    fn eq(&self, other: &ZeroMap2dCursor<'n, 'b, K0, K1, V>) -> bool {
332        self.keys0.eq(other.keys0)
333            && self.joiner.eq(other.joiner)
334            && self.keys1.eq(other.keys1)
335            && self.values.eq(other.values)
336            && self.key0_index.eq(&other.key0_index)
337    }
338}
339
340impl<'l, 'a, K0, K1, V> fmt::Debug for ZeroMap2dCursor<'l, 'a, K0, K1, V>
341where
342    K0: ZeroMapKV<'a> + ?Sized,
343    K1: ZeroMapKV<'a> + ?Sized,
344    V: ZeroMapKV<'a> + ?Sized,
345    K0::Slice: fmt::Debug,
346    K1::Slice: fmt::Debug,
347    V::Slice: fmt::Debug,
348{
349    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> Result<(), fmt::Error> {
350        f.debug_struct("ZeroMap2d")
351            .field("keys0", &self.keys0)
352            .field("joiner", &self.joiner)
353            .field("keys1", &self.keys1)
354            .field("values", &self.values)
355            .field("key0_index", &self.key0_index)
356            .finish()
357    }
358}
zerovec/map2d/cursor.rs

zerovec/map2d/
cursor.rs
