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> dynZeroMapKV<'c> + ?Sized,
325 K1: for<'c> dynZeroMapKV<'c> + ?Sized,
326 V: for<'c> dynZeroMapKV<'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&mut DebugStruct<'_, '_>.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(name:"key0_index", &self.key0_index)
356 .finish()
357 }
358}
359