zerovec/map2d/
cursor.rs

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        #[expect(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 DoubleEndedIterator<
103        Item = (
104            &'l <K1 as ZeroMapKV<'a>>::GetType,
105            &'l <V as ZeroMapKV<'a>>::GetType,
106        ),
107    > + ExactSizeIterator
108           + '_ {
109        let range = self.get_range();
110        #[expect(clippy::unwrap_used)] // `self.get_range()` returns a valid range
111        range.map(move |idx| {
112            (
113                self.keys1.zvl_get(idx).unwrap(),
114                self.values.zvl_get(idx).unwrap(),
115            )
116        })
117    }
118
119    /// Transform this cursor into an ordered iterator over keys1 for a particular key0.
120    pub fn into_iter1(
121        self,
122    ) -> impl DoubleEndedIterator<
123        Item = (
124            &'l <K1 as ZeroMapKV<'a>>::GetType,
125            &'l <V as ZeroMapKV<'a>>::GetType,
126        ),
127    > + ExactSizeIterator {
128        let range = self.get_range();
129        #[expect(clippy::unwrap_used)] // `self.get_range()` returns a valid range
130        range.map(move |idx| {
131            (
132                self.keys1.zvl_get(idx).unwrap(),
133                self.values.zvl_get(idx).unwrap(),
134            )
135        })
136    }
137
138    /// Given key0_index, returns the corresponding range of keys1, which will be valid
139    pub(super) fn get_range(&self) -> Range<usize> {
140        debug_assert!(self.key0_index < self.joiner.len());
141        let start = if self.key0_index == 0 {
142            0
143        } else {
144            #[expect(clippy::unwrap_used)] // protected by the debug_assert above
145            self.joiner.get(self.key0_index - 1).unwrap()
146        };
147        #[expect(clippy::unwrap_used)] // protected by the debug_assert above
148        let limit = self.joiner.get(self.key0_index).unwrap();
149        // These two assertions are true based on the invariants of ZeroMap2d
150        debug_assert!(start < limit);
151        debug_assert!((limit as usize) <= self.values.zvl_len());
152        (start as usize)..(limit as usize)
153    }
154}
155
156impl<'l, 'a, K0, K1, V> ZeroMap2dCursor<'l, 'a, K0, K1, V>
157where
158    K0: ZeroMapKV<'a>,
159    K1: ZeroMapKV<'a>,
160    V: ZeroMapKV<'a>,
161    K0: ?Sized,
162    K1: ?Sized,
163    V: Copy,
164{
165    /// Borrow an ordered iterator over keys1 and values for a particular key0.
166    ///
167    /// The values are returned as copy types.
168    ///
169    /// # Examples
170    ///
171    /// ```
172    /// use zerovec::ZeroMap2d;
173    ///
174    /// let zm2d: ZeroMap2d<str, u8, u16> =
175    ///     [("a", 0u8, 1u16), ("b", 1u8, 1000u16), ("b", 2u8, 2000u16)]
176    ///         .into_iter()
177    ///         .collect();
178    ///
179    /// let mut total_value = 0;
180    ///
181    /// for cursor in zm2d.iter0() {
182    ///     for (_, value) in cursor.iter1_copied() {
183    ///         total_value += value;
184    ///     }
185    /// }
186    ///
187    /// assert_eq!(total_value, 3001);
188    /// ```
189    pub fn iter1_copied(
190        &self,
191    ) -> impl DoubleEndedIterator<Item = (&'l <K1 as ZeroMapKV<'a>>::GetType, V)> + ExactSizeIterator + '_
192    {
193        let range = self.get_range();
194        #[expect(clippy::unwrap_used)] // `self.get_range()` returns a valid range
195        range.map(move |idx| {
196            (
197                self.keys1.zvl_get(idx).unwrap(),
198                self.get1_copied_at(idx).unwrap(),
199            )
200        })
201    }
202    /// Transform this cursor into an ordered iterator over keys1 for a particular key0.
203    ///
204    /// The values are returned as copy types.
205    ///
206    /// # Examples
207    ///
208    /// ```
209    /// use zerovec::ZeroMap2d;
210    ///
211    /// let zm2d: ZeroMap2d<str, u8, u16> =
212    ///     [("a", 0u8, 1u16), ("b", 1u8, 1000u16), ("b", 2u8, 2000u16)]
213    ///         .into_iter()
214    ///         .collect();
215    ///
216    /// let mut total_value = 0;
217    ///
218    /// for cursor in zm2d.iter0() {
219    ///     for (_, value) in cursor.into_iter1_copied() {
220    ///         total_value += value;
221    ///     }
222    /// }
223    ///
224    /// assert_eq!(total_value, 3001);
225    /// ```
226    pub fn into_iter1_copied(
227        self,
228    ) -> impl DoubleEndedIterator<Item = (&'l <K1 as ZeroMapKV<'a>>::GetType, V)> + ExactSizeIterator
229    {
230        let range = self.get_range();
231        #[expect(clippy::unwrap_used)] // `self.get_range()` returns a valid range
232        range.map(move |idx| {
233            (
234                self.keys1.zvl_get(idx).unwrap(),
235                self.get1_copied_at(idx).unwrap(),
236            )
237        })
238    }
239
240    fn get1_copied_at(&self, index: usize) -> Option<V> {
241        let ule = self.values.zvl_get(index)?;
242        let mut result = Option::<V>::None;
243        V::Container::zvl_get_as_t(ule, |v| result.replace(*v));
244        #[expect(clippy::unwrap_used)] // `zvl_get_as_t` guarantees that the callback is invoked
245        Some(result.unwrap())
246    }
247}
248
249impl<'l, 'a, K0, K1, V> ZeroMap2dCursor<'l, 'a, K0, K1, V>
250where
251    K0: ZeroMapKV<'a>,
252    K1: ZeroMapKV<'a> + Ord,
253    V: ZeroMapKV<'a>,
254    K0: ?Sized,
255    K1: ?Sized,
256    V: ?Sized,
257{
258    /// Gets the value for a key1 from this cursor, or `None` if key1 is not in the map.
259    ///
260    /// ```rust
261    /// use zerovec::ZeroMap2d;
262    ///
263    /// let mut map = ZeroMap2d::new();
264    /// map.insert("one", &1u32, "foo");
265    /// assert_eq!(map.get0("one").unwrap().get1(&1), Some("foo"));
266    /// assert_eq!(map.get0("one").unwrap().get1(&2), None);
267    /// ```
268    pub fn get1(&self, key1: &K1) -> Option<&'l V::GetType> {
269        let key1_index = self.get_key1_index(key1)?;
270        #[expect(clippy::unwrap_used)] // key1_index is valid
271        Some(self.values.zvl_get(key1_index).unwrap())
272    }
273
274    /// Gets the value for a predicate from this cursor, or `None` if key1 is not in the map.
275    ///
276    /// ```rust
277    /// use zerovec::ZeroMap2d;
278    ///
279    /// let mut map = ZeroMap2d::new();
280    /// map.insert("one", &1u32, "foo");
281    /// assert_eq!(map.get0("one").unwrap().get1_by(|v| v.cmp(&1)), Some("foo"));
282    /// assert_eq!(map.get0("one").unwrap().get1_by(|v| v.cmp(&2)), None);
283    /// ```
284    pub fn get1_by(&self, predicate: impl FnMut(&K1) -> Ordering) -> Option<&'l V::GetType> {
285        let key1_index = self.get_key1_index_by(predicate)?;
286        #[expect(clippy::unwrap_used)] // key1_index is valid
287        Some(self.values.zvl_get(key1_index).unwrap())
288    }
289
290    /// Given key0_index and predicate, returns the index into the values array
291    fn get_key1_index_by(&self, predicate: impl FnMut(&K1) -> Ordering) -> Option<usize> {
292        let range = self.get_range();
293        debug_assert!(range.start < range.end); // '<' because every key0 should have a key1
294        debug_assert!(range.end <= self.keys1.zvl_len());
295        let start = range.start;
296        #[expect(clippy::expect_used)] // protected by the debug_assert above
297        let binary_search_result = self
298            .keys1
299            .zvl_binary_search_in_range_by(predicate, range)
300            .expect("in-bounds range");
301        binary_search_result.ok().map(move |s| s + start)
302    }
303
304    /// Given key0_index and key1, returns the index into the values array
305    fn get_key1_index(&self, key1: &K1) -> Option<usize> {
306        let range = self.get_range();
307        debug_assert!(range.start < range.end); // '<' because every key0 should have a key1
308        debug_assert!(range.end <= self.keys1.zvl_len());
309        let start = range.start;
310        #[expect(clippy::expect_used)] // protected by the debug_assert above
311        let binary_search_result = self
312            .keys1
313            .zvl_binary_search_in_range(key1, range)
314            .expect("in-bounds range");
315        binary_search_result.ok().map(move |s| s + start)
316    }
317}
318
319impl<'l, 'a, K0, K1, V> ZeroMap2dCursor<'l, 'a, K0, K1, V>
320where
321    K0: ZeroMapKV<'a>,
322    K1: ZeroMapKV<'a> + Ord,
323    V: ZeroMapKV<'a>,
324    V: Copy,
325    K0: ?Sized,
326    K1: ?Sized,
327{
328    /// For cases when `V` is fixed-size, obtain a direct copy of `V` instead of `V::ULE`
329    ///
330    /// ```rust
331    /// use zerovec::ZeroMap2d;
332    ///
333    /// let mut map: ZeroMap2d<u16, u16, u16> = ZeroMap2d::new();
334    /// map.insert(&1, &2, &3);
335    /// map.insert(&1, &4, &5);
336    /// map.insert(&6, &7, &8);
337    ///
338    /// assert_eq!(map.get0(&6).unwrap().get1_copied(&7), Some(8));
339    /// ```
340    #[inline]
341    pub fn get1_copied(&self, key1: &K1) -> Option<V> {
342        let key1_index = self.get_key1_index(key1)?;
343        self.get1_copied_at(key1_index)
344    }
345
346    /// For cases when `V` is fixed-size, obtain a direct copy of `V` instead of `V::ULE`
347    #[inline]
348    pub fn get1_copied_by(&self, predicate: impl FnMut(&K1) -> Ordering) -> Option<V> {
349        let key1_index = self.get_key1_index_by(predicate)?;
350        self.get1_copied_at(key1_index)
351    }
352}
353
354// We can't use the default PartialEq because ZeroMap2d is invariant
355// so otherwise rustc will not automatically allow you to compare ZeroMaps
356// with different lifetimes
357impl<'m, 'n, 'a, 'b, K0, K1, V> PartialEq<ZeroMap2dCursor<'n, 'b, K0, K1, V>>
358    for ZeroMap2dCursor<'m, 'a, K0, K1, V>
359where
360    K0: for<'c> ZeroMapKV<'c> + ?Sized,
361    K1: for<'c> ZeroMapKV<'c> + ?Sized,
362    V: for<'c> ZeroMapKV<'c> + ?Sized,
363    <K0 as ZeroMapKV<'a>>::Slice: PartialEq<<K0 as ZeroMapKV<'b>>::Slice>,
364    <K1 as ZeroMapKV<'a>>::Slice: PartialEq<<K1 as ZeroMapKV<'b>>::Slice>,
365    <V as ZeroMapKV<'a>>::Slice: PartialEq<<V as ZeroMapKV<'b>>::Slice>,
366{
367    fn eq(&self, other: &ZeroMap2dCursor<'n, 'b, K0, K1, V>) -> bool {
368        self.keys0.eq(other.keys0)
369            && self.joiner.eq(other.joiner)
370            && self.keys1.eq(other.keys1)
371            && self.values.eq(other.values)
372            && self.key0_index.eq(&other.key0_index)
373    }
374}
375
376impl<'l, 'a, K0, K1, V> fmt::Debug for ZeroMap2dCursor<'l, 'a, K0, K1, V>
377where
378    K0: ZeroMapKV<'a> + ?Sized,
379    K1: ZeroMapKV<'a> + ?Sized,
380    V: ZeroMapKV<'a> + ?Sized,
381    K0::Slice: fmt::Debug,
382    K1::Slice: fmt::Debug,
383    V::Slice: fmt::Debug,
384{
385    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> Result<(), fmt::Error> {
386        f.debug_struct("ZeroMap2d")
387            .field("keys0", &self.keys0)
388            .field("joiner", &self.joiner)
389            .field("keys1", &self.keys1)
390            .field("values", &self.values)
391            .field("key0_index", &self.key0_index)
392            .finish()
393    }
394}
zerovec/map2d/cursor.rs

zerovec/map2d/
cursor.rs
