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 super::*;
6use crate::ule::{AsULE, EncodeAsVarULE, VarULE};
7use crate::{VarZeroVec, ZeroSlice, ZeroVec};
8use alloc::borrow::Borrow;
9use alloc::boxed::Box;
10use core::cmp::Ordering;
11use core::fmt;
12use core::iter::FromIterator;
13
14/// A zero-copy map datastructure, built on sorted binary-searchable [`ZeroVec`]
15/// and [`VarZeroVec`].
16///
17/// This type, like [`ZeroVec`] and [`VarZeroVec`], is able to zero-copy
18/// deserialize from appropriately formatted byte buffers. It is internally copy-on-write, so it can be mutated
19/// afterwards as necessary.
20///
21/// Internally, a `ZeroMap` is a zero-copy vector for keys paired with a zero-copy vector for
22/// values, sorted by the keys. Therefore, all types used in `ZeroMap` need to work with either
23/// [`ZeroVec`] or [`VarZeroVec`].
24///
25/// This does mean that for fixed-size data, one must use the regular type (`u32`, `u8`, `char`, etc),
26/// whereas for variable-size data, `ZeroMap` will use the dynamically sized version (`str` not `String`,
27/// `ZeroSlice` not `ZeroVec`, `FooULE` not `Foo` for custom types)
28///
29/// # Examples
30///
31/// ```
32/// use zerovec::ZeroMap;
33///
34/// #[derive(serde::Serialize, serde::Deserialize)]
35/// struct Data<'a> {
36/// #[serde(borrow)]
37/// map: ZeroMap<'a, u32, str>,
38/// }
39///
40/// let mut map = ZeroMap::new();
41/// map.insert(&1, "one");
42/// map.insert(&2, "two");
43/// map.insert(&4, "four");
44///
45/// let data = Data { map };
46///
47/// let bincode_bytes =
48/// bincode::serialize(&data).expect("Serialization should be successful");
49///
50/// // Will deserialize without any allocations
51/// let deserialized: Data = bincode::deserialize(&bincode_bytes)
52/// .expect("Deserialization should be successful");
53///
54/// assert_eq!(data.map.get(&1), Some("one"));
55/// assert_eq!(data.map.get(&2), Some("two"));
56/// ```
57///
58/// [`VarZeroVec`]: crate::VarZeroVec
59// ZeroMap has only one invariant: keys.len() == values.len()
60// It is also expected that the keys are sorted, but this is not an invariant. See #1433
61pub struct ZeroMap<'a, K, V>
62where
63 K: ZeroMapKV<'a> + ?Sized,
64 V: ZeroMapKV<'a> + ?Sized,
65{
66 pub(crate) keys: K::Container,
67 pub(crate) values: V::Container,
68}
69
70impl<'a, K, V> Default for ZeroMap<'a, K, V>
71where
72 K: ZeroMapKV<'a> + ?Sized,
73 V: ZeroMapKV<'a> + ?Sized,
74{
75 fn default() -> Self {
76 Self::new()
77 }
78}
79
80impl<'a, K, V> ZeroMap<'a, K, V>
81where
82 K: ZeroMapKV<'a> + ?Sized,
83 V: ZeroMapKV<'a> + ?Sized,
84{
85 /// Creates a new, empty `ZeroMap<K, V>`.
86 ///
87 /// # Examples
88 ///
89 /// ```
90 /// use zerovec::ZeroMap;
91 ///
92 /// let zm: ZeroMap<u16, str> = ZeroMap::new();
93 /// assert!(zm.is_empty());
94 /// ```
95 pub fn new() -> Self {
96 Self {
97 keys: K::Container::zvl_with_capacity(0),
98 values: V::Container::zvl_with_capacity(0),
99 }
100 }
101
102 #[doc(hidden)] // databake internal
103 pub const unsafe fn from_parts_unchecked(keys: K::Container, values: V::Container) -> Self {
104 Self { keys, values }
105 }
106
107 /// Construct a new [`ZeroMap`] with a given capacity
108 pub fn with_capacity(capacity: usize) -> Self {
109 Self {
110 keys: K::Container::zvl_with_capacity(capacity),
111 values: V::Container::zvl_with_capacity(capacity),
112 }
113 }
114
115 /// Obtain a borrowed version of this map
116 pub fn as_borrowed(&'a self) -> ZeroMapBorrowed<'a, K, V> {
117 ZeroMapBorrowed {
118 keys: self.keys.zvl_as_borrowed(),
119 values: self.values.zvl_as_borrowed(),
120 }
121 }
122
123 /// The number of elements in the [`ZeroMap`]
124 pub fn len(&self) -> usize {
125 self.values.zvl_len()
126 }
127
128 /// Whether the [`ZeroMap`] is empty
129 pub fn is_empty(&self) -> bool {
130 self.values.zvl_len() == 0
131 }
132
133 /// Remove all elements from the [`ZeroMap`]
134 pub fn clear(&mut self) {
135 self.keys.zvl_clear();
136 self.values.zvl_clear();
137 }
138
139 /// Reserve capacity for `additional` more elements to be inserted into
140 /// the [`ZeroMap`] to avoid frequent reallocations.
141 ///
142 /// See [`Vec::reserve()`](alloc::vec::Vec::reserve) for more information.
143 pub fn reserve(&mut self, additional: usize) {
144 self.keys.zvl_reserve(additional);
145 self.values.zvl_reserve(additional);
146 }
147}
148impl<'a, K, V> ZeroMap<'a, K, V>
149where
150 K: ZeroMapKV<'a> + ?Sized + Ord,
151 V: ZeroMapKV<'a> + ?Sized,
152{
153 /// Get the value associated with `key`, if it exists.
154 ///
155 /// For fixed-size ([`AsULE`]) `V` types, this _will_ return
156 /// their corresponding [`AsULE::ULE`] type. If you wish to work with the `V`
157 /// type directly, [`Self::get_copied()`] exists for convenience.
158 ///
159 /// ```rust
160 /// use zerovec::ZeroMap;
161 ///
162 /// let mut map = ZeroMap::new();
163 /// map.insert(&1, "one");
164 /// map.insert(&2, "two");
165 /// assert_eq!(map.get(&1), Some("one"));
166 /// assert_eq!(map.get(&3), None);
167 /// ```
168 pub fn get(&self, key: &K) -> Option<&V::GetType> {
169 let index = self.keys.zvl_binary_search(key).ok()?;
170 self.values.zvl_get(index)
171 }
172
173 /// Binary search the map with `predicate` to find a key, returning the value.
174 ///
175 /// ```rust
176 /// use zerovec::ZeroMap;
177 ///
178 /// let mut map = ZeroMap::new();
179 /// map.insert(&1, "one");
180 /// map.insert(&2, "two");
181 /// assert_eq!(map.get_by(|probe| probe.cmp(&1)), Some("one"));
182 /// assert_eq!(map.get_by(|probe| probe.cmp(&3)), None);
183 /// ```
184 pub fn get_by(&self, predicate: impl FnMut(&K) -> Ordering) -> Option<&V::GetType> {
185 let index = self.keys.zvl_binary_search_by(predicate).ok()?;
186 self.values.zvl_get(index)
187 }
188
189 /// Returns whether `key` is contained in this map
190 ///
191 /// ```rust
192 /// use zerovec::ZeroMap;
193 ///
194 /// let mut map = ZeroMap::new();
195 /// map.insert(&1, "one");
196 /// map.insert(&2, "two");
197 /// assert!(map.contains_key(&1));
198 /// assert!(!map.contains_key(&3));
199 /// ```
200 pub fn contains_key(&self, key: &K) -> bool {
201 self.keys.zvl_binary_search(key).is_ok()
202 }
203
204 /// Insert `value` with `key`, returning the existing value if it exists.
205 ///
206 /// ```rust
207 /// use zerovec::ZeroMap;
208 ///
209 /// let mut map = ZeroMap::new();
210 /// map.insert(&1, "one");
211 /// map.insert(&2, "two");
212 /// assert_eq!(map.get(&1), Some("one"));
213 /// assert_eq!(map.get(&3), None);
214 /// ```
215 pub fn insert(&mut self, key: &K, value: &V) -> Option<V::OwnedType> {
216 match self.keys.zvl_binary_search(key) {
217 Ok(index) => Some(self.values.zvl_replace(index, value)),
218 Err(index) => {
219 self.keys.zvl_insert(index, key);
220 self.values.zvl_insert(index, value);
221 None
222 }
223 }
224 }
225
226 /// Remove the value at `key`, returning it if it exists.
227 ///
228 /// ```rust
229 /// use zerovec::ZeroMap;
230 ///
231 /// let mut map = ZeroMap::new();
232 /// map.insert(&1, "one");
233 /// map.insert(&2, "two");
234 /// assert_eq!(map.remove(&1), Some("one".to_owned().into_boxed_str()));
235 /// assert_eq!(map.get(&1), None);
236 /// ```
237 pub fn remove(&mut self, key: &K) -> Option<V::OwnedType> {
238 let idx = self.keys.zvl_binary_search(key).ok()?;
239 self.keys.zvl_remove(idx);
240 Some(self.values.zvl_remove(idx))
241 }
242
243 /// Appends `value` with `key` to the end of the underlying vector, returning
244 /// `key` and `value` _if it failed_. Useful for extending with an existing
245 /// sorted list.
246 /// ```rust
247 /// use zerovec::ZeroMap;
248 ///
249 /// let mut map = ZeroMap::new();
250 /// assert!(map.try_append(&1, "uno").is_none());
251 /// assert!(map.try_append(&3, "tres").is_none());
252 ///
253 /// let unsuccessful = map.try_append(&3, "tres-updated");
254 /// assert!(unsuccessful.is_some(), "append duplicate of last key");
255 ///
256 /// let unsuccessful = map.try_append(&2, "dos");
257 /// assert!(unsuccessful.is_some(), "append out of order");
258 ///
259 /// assert_eq!(map.get(&1), Some("uno"));
260 ///
261 /// // contains the original value for the key: 3
262 /// assert_eq!(map.get(&3), Some("tres"));
263 ///
264 /// // not appended since it wasn't in order
265 /// assert_eq!(map.get(&2), None);
266 /// ```
267 #[must_use]
268 pub fn try_append<'b>(&mut self, key: &'b K, value: &'b V) -> Option<(&'b K, &'b V)> {
269 if self.keys.zvl_len() != 0 {
270 if let Some(last) = self.keys.zvl_get(self.keys.zvl_len() - 1) {
271 if K::Container::t_cmp_get(key, last) != Ordering::Greater {
272 return Some((key, value));
273 }
274 }
275 }
276
277 self.keys.zvl_push(key);
278 self.values.zvl_push(value);
279 None
280 }
281}
282
283impl<'a, K, V> ZeroMap<'a, K, V>
284where
285 K: ZeroMapKV<'a> + ?Sized,
286 V: ZeroMapKV<'a> + ?Sized,
287{
288 /// Produce an ordered iterator over key-value pairs
289 pub fn iter<'b>(
290 &'b self,
291 ) -> impl ExactSizeIterator<
292 Item = (
293 &'b <K as ZeroMapKV<'a>>::GetType,
294 &'b <V as ZeroMapKV<'a>>::GetType,
295 ),
296 > {
297 (0..self.keys.zvl_len()).map(move |idx| {
298 (
299 #[allow(clippy::unwrap_used)] // idx is in-range
300 self.keys.zvl_get(idx).unwrap(),
301 #[allow(clippy::unwrap_used)] // idx is in-range
302 self.values.zvl_get(idx).unwrap(),
303 )
304 })
305 }
306
307 /// Produce an ordered iterator over keys
308 pub fn iter_keys<'b>(
309 &'b self,
310 ) -> impl ExactSizeIterator<Item = &'b <K as ZeroMapKV<'a>>::GetType> {
311 #[allow(clippy::unwrap_used)] // idx is in-range
312 (0..self.keys.zvl_len()).map(move |idx| self.keys.zvl_get(idx).unwrap())
313 }
314
315 /// Produce an iterator over values, ordered by keys
316 pub fn iter_values<'b>(
317 &'b self,
318 ) -> impl ExactSizeIterator<Item = &'b <V as ZeroMapKV<'a>>::GetType> {
319 #[allow(clippy::unwrap_used)] // idx is in-range
320 (0..self.values.zvl_len()).map(move |idx| self.values.zvl_get(idx).unwrap())
321 }
322}
323
324impl<'a, K, V> ZeroMap<'a, K, V>
325where
326 K: ZeroMapKV<'a> + ?Sized + Ord,
327 V: ZeroMapKV<'a, Container = VarZeroVec<'a, V>> + ?Sized,
328 V: VarULE,
329{
330 /// Same as `insert()`, but allows using [EncodeAsVarULE](crate::ule::EncodeAsVarULE)
331 /// types with the value to avoid an extra allocation when dealing with custom ULE types.
332 ///
333 /// ```rust
334 /// use std::borrow::Cow;
335 /// use zerovec::ZeroMap;
336 ///
337 /// #[zerovec::make_varule(PersonULE)]
338 /// #[derive(Clone, Eq, PartialEq, Ord, PartialOrd)]
339 /// struct Person<'a> {
340 /// age: u8,
341 /// name: Cow<'a, str>,
342 /// }
343 ///
344 /// let mut map: ZeroMap<u32, PersonULE> = ZeroMap::new();
345 /// map.insert_var_v(
346 /// &1,
347 /// &Person {
348 /// age: 20,
349 /// name: "Joseph".into(),
350 /// },
351 /// );
352 /// map.insert_var_v(
353 /// &1,
354 /// &Person {
355 /// age: 35,
356 /// name: "Carla".into(),
357 /// },
358 /// );
359 /// assert_eq!(&map.get(&1).unwrap().name, "Carla");
360 /// assert!(map.get(&3).is_none());
361 /// ```
362 pub fn insert_var_v<VE: EncodeAsVarULE<V>>(&mut self, key: &K, value: &VE) -> Option<Box<V>> {
363 match self.keys.zvl_binary_search(key) {
364 Ok(index) => {
365 #[allow(clippy::unwrap_used)] // binary search
366 let ret = self.values.get(index).unwrap().to_boxed();
367 self.values.make_mut().replace(index, value);
368 Some(ret)
369 }
370 Err(index) => {
371 self.keys.zvl_insert(index, key);
372 self.values.make_mut().insert(index, value);
373 None
374 }
375 }
376 }
377
378 // insert_var_k, insert_var_kv are not possible since one cannot perform the binary search with EncodeAsVarULE
379 // though we might be able to do it in the future if we add a trait for cross-Ord requirements
380}
381
382impl<'a, K, V> ZeroMap<'a, K, V>
383where
384 K: ZeroMapKV<'a> + ?Sized + Ord,
385 V: ZeroMapKV<'a> + ?Sized,
386 V: Copy,
387{
388 /// For cases when `V` is fixed-size, obtain a direct copy of `V` instead of `V::ULE`.
389 ///
390 /// # Examples
391 ///
392 /// ```rust
393 /// use zerovec::ZeroMap;
394 ///
395 /// let mut map = ZeroMap::new();
396 /// map.insert(&1, &'a');
397 /// map.insert(&2, &'b');
398 /// assert_eq!(map.get_copied(&1), Some('a'));
399 /// assert_eq!(map.get_copied(&3), None);
400 #[inline]
401 pub fn get_copied(&self, key: &K) -> Option<V> {
402 let index = self.keys.zvl_binary_search(key).ok()?;
403 self.get_copied_at(index)
404 }
405
406 /// Binary search the map with `predicate` to find a key, returning the value.
407 ///
408 /// For cases when `V` is fixed-size, use this method to obtain a direct copy of `V`
409 /// instead of `V::ULE`.
410 ///
411 /// # Examples
412 ///
413 /// ```rust
414 /// use zerovec::ZeroMap;
415 ///
416 /// let mut map = ZeroMap::new();
417 /// map.insert(&1, &'a');
418 /// map.insert(&2, &'b');
419 /// assert_eq!(map.get_copied_by(|probe| probe.cmp(&1)), Some('a'));
420 /// assert_eq!(map.get_copied_by(|probe| probe.cmp(&3)), None);
421 /// ```
422 #[inline]
423 pub fn get_copied_by(&self, predicate: impl FnMut(&K) -> Ordering) -> Option<V> {
424 let index = self.keys.zvl_binary_search_by(predicate).ok()?;
425 self.get_copied_at(index)
426 }
427
428 fn get_copied_at(&self, index: usize) -> Option<V> {
429 let ule = self.values.zvl_get(index)?;
430 let mut result = Option::<V>::None;
431 V::Container::zvl_get_as_t(ule, |v| result.replace(*v));
432 #[allow(clippy::unwrap_used)] // `zvl_get_as_t` guarantees that the callback is invoked
433 Some(result.unwrap())
434 }
435}
436
437impl<'a, K, V> ZeroMap<'a, K, V>
438where
439 K: ZeroMapKV<'a> + ?Sized,
440 V: ZeroMapKV<'a, Container = ZeroVec<'a, V>> + ?Sized,
441 V: AsULE + Copy,
442{
443 /// Similar to [`Self::iter()`] except it returns a direct copy of the values instead of references
444 /// to `V::ULE`, in cases when `V` is fixed-size
445 pub fn iter_copied_values<'b>(
446 &'b self,
447 ) -> impl Iterator<Item = (&'b <K as ZeroMapKV<'a>>::GetType, V)> {
448 (0..self.keys.zvl_len()).map(move |idx: usize| {
449 (
450 #[allow(clippy::unwrap_used)] // idx in 0..keys.zvl_len()
451 self.keys.zvl_get(index:idx).unwrap(),
452 #[allow(clippy::unwrap_used)] // idx in 0..keys.zvl_len() = values.zvl_len()
453 ZeroSlice::get(&*self.values, index:idx).unwrap(),
454 )
455 })
456 }
457}
458
459impl<'a, K, V> ZeroMap<'a, K, V>
460where
461 K: ZeroMapKV<'a, Container = ZeroVec<'a, K>> + ?Sized,
462 V: ZeroMapKV<'a, Container = ZeroVec<'a, V>> + ?Sized,
463 K: AsULE + Copy,
464 V: AsULE + Copy,
465{
466 /// Similar to [`Self::iter()`] except it returns a direct copy of the keys values instead of references
467 /// to `K::ULE` and `V::ULE`, in cases when `K` and `V` are fixed-size
468 #[allow(clippy::needless_lifetimes)] // Lifetime is necessary in impl Trait
469 pub fn iter_copied<'b>(&'b self) -> impl Iterator<Item = (K, V)> + 'b {
470 let keys: &ZeroVec<'_, K> = &self.keys;
471 let values: &ZeroVec<'_, V> = &self.values;
472 (0..keys.len()).map(move |idx: usize| {
473 (
474 #[allow(clippy::unwrap_used)] // idx in 0..keys.zvl_len()
475 ZeroSlice::get(&**keys, index:idx).unwrap(),
476 #[allow(clippy::unwrap_used)] // idx in 0..keys.zvl_len() = values.zvl_len()
477 ZeroSlice::get(&**values, index:idx).unwrap(),
478 )
479 })
480 }
481}
482
483impl<'a, K, V> From<ZeroMapBorrowed<'a, K, V>> for ZeroMap<'a, K, V>
484where
485 K: ZeroMapKV<'a>,
486 V: ZeroMapKV<'a>,
487 K: ?Sized,
488 V: ?Sized,
489{
490 fn from(other: ZeroMapBorrowed<'a, K, V>) -> Self {
491 Self {
492 keys: K::Container::zvl_from_borrowed(other.keys),
493 values: V::Container::zvl_from_borrowed(other.values),
494 }
495 }
496}
497
498// We can't use the default PartialEq because ZeroMap is invariant
499// so otherwise rustc will not automatically allow you to compare ZeroMaps
500// with different lifetimes
501impl<'a, 'b, K, V> PartialEq<ZeroMap<'b, K, V>> for ZeroMap<'a, K, V>
502where
503 K: for<'c> ZeroMapKV<'c> + ?Sized,
504 V: for<'c> ZeroMapKV<'c> + ?Sized,
505 <K as ZeroMapKV<'a>>::Container: PartialEq<<K as ZeroMapKV<'b>>::Container>,
506 <V as ZeroMapKV<'a>>::Container: PartialEq<<V as ZeroMapKV<'b>>::Container>,
507{
508 fn eq(&self, other: &ZeroMap<'b, K, V>) -> bool {
509 self.keys.eq(&other.keys) && self.values.eq(&other.values)
510 }
511}
512
513impl<'a, K, V> fmt::Debug for ZeroMap<'a, K, V>
514where
515 K: ZeroMapKV<'a> + ?Sized,
516 V: ZeroMapKV<'a> + ?Sized,
517 <K as ZeroMapKV<'a>>::Container: fmt::Debug,
518 <V as ZeroMapKV<'a>>::Container: fmt::Debug,
519{
520 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> Result<(), fmt::Error> {
521 f&mut DebugStruct<'_, '_>.debug_struct("ZeroMap")
522 .field("keys", &self.keys)
523 .field(name:"values", &self.values)
524 .finish()
525 }
526}
527
528impl<'a, K, V> Clone for ZeroMap<'a, K, V>
529where
530 K: ZeroMapKV<'a> + ?Sized,
531 V: ZeroMapKV<'a> + ?Sized,
532 <K as ZeroMapKV<'a>>::Container: Clone,
533 <V as ZeroMapKV<'a>>::Container: Clone,
534{
535 fn clone(&self) -> Self {
536 Self {
537 keys: self.keys.clone(),
538 values: self.values.clone(),
539 }
540 }
541}
542
543impl<'a, A, B, K, V> FromIterator<(A, B)> for ZeroMap<'a, K, V>
544where
545 A: Borrow<K>,
546 B: Borrow<V>,
547 K: ZeroMapKV<'a> + ?Sized + Ord,
548 V: ZeroMapKV<'a> + ?Sized,
549{
550 fn from_iter<T>(iter: T) -> Self
551 where
552 T: IntoIterator<Item = (A, B)>,
553 {
554 let iter: ::IntoIter = iter.into_iter();
555 let mut map: ZeroMap<'_, K, V> = match iter.size_hint() {
556 (_, Some(upper: usize)) => Self::with_capacity(upper),
557 (lower: usize, None) => Self::with_capacity(lower),
558 };
559
560 for (key: A, value: B) in iter {
561 if let Some((key: &K, value: &V)) = map.try_append(key:key.borrow(), value:value.borrow()) {
562 map.insert(key, value);
563 }
564 }
565 map
566 }
567}
568