Impl Extend for LiteMap and avoid quadratic behavior in from_iter and deserialize (#6132)

Adds `lm_extend` to the `StoreMut` trait. The Vec implementation is
optimized to have O(N) performance on sorted inputs and an asymptotic
worst case of O(NLog(N)), effectively avoiding quadratic worst cases.

This PR takes advantage of the new `extend` optimization and changes the
`FromIterator` and `Deserialize` implementations to also avoid quadratic
worst-case.

Before
```
litemap/from_iter_rand/small
                        time:   [916.60 ns 941.77 ns 965.82 ns]
litemap/from_iter_rand/large
                        time:   [2.3779 s 2.3863 s 2.3946 s]
litemap/from_iter_sorted/small
                        time:   [84.010 ns 84.084 ns 84.155 ns]
litemap/from_iter_sorted/large
                        time:   [725.31 µs 739.88 µs 755.20 µs]
```

After

```
litemap/from_iter_rand/small
                        time:   [354.62 ns 365.31 ns 376.02 ns]
litemap/from_iter_rand/large
                        time:   [25.129 ms 25.415 ms 25.704 ms]
litemap/from_iter_sorted/small
                        time:   [76.593 ns 76.767 ns 76.938 ns]
litemap/from_iter_sorted/large
                        time:   [662.70 µs 673.66 µs 686.83 µs]
litemap/extend_rand/large
                        time:   [66.221 ms 67.467 ms 68.814 ms]
litemap/extend_rand_dups/large
                        time:   [190.24 ms 196.05 ms 202.13 ms]
litemap/extend_from_litemap_rand/large
                        time:   [2.0401 s 2.0488 s 2.0575 s]
```

This is a followup of #6068

Author: arthurprs

Cargo.lock

@@ -12,6 +12,15 @@ impl<K, V> StoreConstEmpty<K, V> for ShortBoxSlice<(K, V)> {
     const EMPTY: ShortBoxSlice<(K, V)> = ShortBoxSlice::new();
 }
 
+impl<K, V> StoreSlice<K, V> for ShortBoxSlice<(K, V)> {
+    type Slice = [(K, V)];
+
+    #[inline]
+    fn lm_get_range(&self, range: core::ops::Range<usize>) -> Option<&Self::Slice> {
+        self.get(range)
+    }
+}
+
 impl<K, V> Store<K, V> for ShortBoxSlice<(K, V)> {
     #[inline]
     fn lm_len(&self) -> usize {
@@ -83,13 +92,60 @@ impl<K, V> StoreMut<K, V> for ShortBoxSlice<(K, V)> {
     fn lm_clear(&mut self) {
         self.clear();
     }
+}
 
+#[cfg(feature = "alloc")]
+impl<K: Ord, V> StoreBulkMut<K, V> for ShortBoxSlice<(K, V)> {
     fn lm_retain<F>(&mut self, mut predicate: F)
     where
         F: FnMut(&K, &V) -> bool,
     {
         self.retain(|(k, v)| predicate(k, v))
     }
+
+    fn lm_extend<I>(&mut self, other: I)
+    where
+        I: IntoIterator<Item = (K, V)>,
+    {
+        let mut other = other.into_iter();
+        // Use an Option to hold the first item of the map and move it to
+        // items if there are more items. Meaning that if items is not
+        // empty, first is None.
+        let mut first = None;
+        let mut items = alloc::vec::Vec::new();
+        match core::mem::take(&mut self.0) {
+            ShortBoxSliceInner::ZeroOne(zo) => {
+                first = zo;
+                // Attempt to avoid the items allocation by advancing the iterator
+                // up to two times. If we eventually find a second item, we can
+                // lm_extend the Vec and with the first, next (second) and the rest
+                // of the iterator.
+                while let Some(next) = other.next() {
+                    if let Some(first) = first.take() {
+                        // lm_extend will take care of sorting and deduplicating
+                        // first, next and the rest of the other iterator.
+                        items.lm_extend([first, next].into_iter().chain(other));
+                        break;
+                    }
+                    first = Some(next);
+                }
+            }
+            ShortBoxSliceInner::Multi(existing_items) => {
+                items.reserve_exact(existing_items.len() + other.size_hint().0);
+                // We use a plain extend with existing items, which are already valid and
+                // lm_extend will fold over rest of the iterator sorting and deduplicating as needed.
+                items.extend(existing_items);
+                items.lm_extend(other);
+            }
+        }
+        if items.is_empty() {
+            debug_assert!(items.is_empty());
+            self.0 = ShortBoxSliceInner::ZeroOne(first);
+        } else {
+            debug_assert!(first.is_none());
+            self.0 = ShortBoxSliceInner::Multi(items.into_boxed_slice());
+        }
+    }
 }
 
 impl<'a, K: 'a, V: 'a> StoreIterable<'a, K, V> for ShortBoxSlice<(K, V)> {

components/locale_core/src/shortvec/litemap.rs

@@ -28,6 +28,7 @@ serde = { workspace = true, optional = true, features = ["alloc"]}
 yoke = { workspace = true, features = ["derive"], optional = true }
 
 [dev-dependencies]
+rand = { workspace = true }
 bincode = { workspace = true }
 icu_benchmark_macros = { path = "../../tools/benchmark/macros" }
 icu_locale_core = { path = "../../components/locale_core" }
@@ -43,7 +44,7 @@ criterion = { workspace = true }
 default = ["alloc"]
 alloc = []
 databake = ["dep:databake"]
-serde = ["dep:serde"]
+serde = ["dep:serde", "alloc"]
 yoke = ["dep:yoke"]
 
 # Enables the `testing` module with tools for testing custom stores.

utils/litemap/Cargo.toml

@@ -5,6 +5,7 @@
 use criterion::{black_box, criterion_group, criterion_main, Criterion};
 
 use litemap::LiteMap;
+use rand::seq::SliceRandom;
 
 const DATA: [(&str, &str); 16] = [
     ("ar", "Arabic"),
@@ -57,6 +58,13 @@ fn overview_bench(c: &mut Criterion) {
     bench_deserialize_large(c);
     bench_lookup(c);
     bench_lookup_large(c);
+    bench_from_iter(c);
+    bench_from_iter_rand_large(c);
+    bench_from_iter_sorted(c);
+    bench_from_iter_large_sorted(c);
+    bench_extend_rand(c);
+    bench_extend_rand_dups(c);
+    bench_extend_from_litemap_rand(c);
 }
 
 fn build_litemap(large: bool) -> LiteMap<String, String> {
@@ -92,6 +100,99 @@ fn bench_deserialize_large(c: &mut Criterion) {
     });
 }
 
+fn bench_from_iter(c: &mut Criterion) {
+    c.bench_function("litemap/from_iter_rand/small", |b| {
+        let mut ff = build_litemap(false).into_iter().collect::<Vec<_>>();
+        ff[..].shuffle(&mut rand::rng());
+        b.iter(|| {
+            let map: LiteMap<&String, &String> = LiteMap::from_iter(ff.iter().map(|(k, v)| (k, v)));
+            black_box(map)
+        })
+    });
+}
+
+fn bench_from_iter_rand_large(c: &mut Criterion) {
+    c.bench_function("litemap/from_iter_rand/large", |b| {
+        let mut ff = build_litemap(true).into_iter().collect::<Vec<_>>();
+        ff[..].shuffle(&mut rand::rng());
+        b.iter(|| {
+            let map: LiteMap<&String, &String> = LiteMap::from_iter(ff.iter().map(|(k, v)| (k, v)));
+            black_box(map)
+        })
+    });
+}
+
+fn bench_from_iter_sorted(c: &mut Criterion) {
+    c.bench_function("litemap/from_iter_sorted/small", |b| {
+        let ff = build_litemap(false).into_iter().collect::<Vec<_>>();
+        b.iter(|| {
+            let map: LiteMap<&String, &String> = LiteMap::from_iter(ff.iter().map(|(k, v)| (k, v)));
+            black_box(map)
+        })
+    });
+}
+
+fn bench_from_iter_large_sorted(c: &mut Criterion) {
+    c.bench_function("litemap/from_iter_sorted/large", |b| {
+        let ff = build_litemap(true).into_iter().collect::<Vec<_>>();
+        b.iter(|| {
+            let map: LiteMap<&String, &String> = LiteMap::from_iter(ff.iter().map(|(k, v)| (k, v)));
+            black_box(map)
+        })
+    });
+}
+
+fn bench_extend_rand(c: &mut Criterion) {
+    c.bench_function("litemap/extend_rand/large", |b| {
+        let mut ff = build_litemap(true).into_iter().collect::<Vec<_>>();
+        ff[..].shuffle(&mut rand::rng());
+        b.iter(|| {
+            let mut map: LiteMap<&String, &String> = LiteMap::with_capacity(0);
+            let mut iter = ff.iter().map(|(k, v)| (k, v));
+            let step = ff.len().div_ceil(10);
+            for _ in 0..10 {
+                map.extend(iter.by_ref().take(step));
+            }
+            black_box(map)
+        })
+    });
+}
+
+fn bench_extend_rand_dups(c: &mut Criterion) {
+    c.bench_function("litemap/extend_rand_dups/large", |b| {
+        let mut ff = build_litemap(true).into_iter().collect::<Vec<_>>();
+        ff[..].shuffle(&mut rand::rng());
+        b.iter(|| {
+            let mut map: LiteMap<&String, &String> = LiteMap::with_capacity(0);
+            for _ in 0..2 {
+                let mut iter = ff.iter().map(|(k, v)| (k, v));
+                let step = ff.len().div_ceil(10);
+                for _ in 0..10 {
+                    map.extend(iter.by_ref().take(step));
+                }
+            }
+            black_box(map)
+        })
+    });
+}
+
+fn bench_extend_from_litemap_rand(c: &mut Criterion) {
+    c.bench_function("litemap/extend_from_litemap_rand/large", |b| {
+        let mut ff = build_litemap(true).into_iter().collect::<Vec<_>>();
+        ff[..].shuffle(&mut rand::rng());
+        b.iter(|| {
+            let mut map: LiteMap<&String, &String> = LiteMap::with_capacity(0);
+            let mut iter = ff.iter().map(|(k, v)| (k, v));
+            let step = ff.len().div_ceil(10);
+            for _ in 0..10 {
+                let tmp: LiteMap<&String, &String> = LiteMap::from_iter(iter.by_ref().take(step));
+                map.extend_from_litemap(tmp);
+            }
+            black_box(map)
+        })
+    });
+}
+
 fn bench_lookup(c: &mut Criterion) {
     let map: LiteMap<String, String> = postcard::from_bytes(&POSTCARD).unwrap();
     c.bench_function("litemap/lookup/small", |b| {

utils/litemap/README.md

@@ -7,14 +7,28 @@
 //! `litemap` is a crate providing [`LiteMap`], a highly simplistic "flat" key-value map
 //! based off of a single sorted vector.
 //!
-//! The goal of this crate is to provide a map that is good enough for small
+//! The main goal of this crate is to provide a map that is good enough for small
 //! sizes, and does not carry the binary size impact of [`HashMap`](std::collections::HashMap)
 //! or [`BTreeMap`](alloc::collections::BTreeMap).
 //!
 //! If binary size is not a concern, [`std::collections::BTreeMap`] may be a better choice
 //! for your use case. It behaves very similarly to [`LiteMap`] for less than 12 elements,
 //! and upgrades itself gracefully for larger inputs.
 //!
+//! ## Performance characteristics
+//!
+//! [`LiteMap`] is a data structure with similar characteristics as [`std::collections::BTreeMap`] but
+//! with slightly different trade-offs as it's implemented on top of a flat storage (e.g. Vec).
+//!
+//! * [`LiteMap`] iteration is generally faster than `BTreeMap` because of the flat storage.
+//! * [`LiteMap`] can be pre-allocated whereas `BTreeMap` can't.
+//! * [`LiteMap`] has a smaller memory footprint than `BTreeMap` for small collections (< 20 items).
+//! * Lookup is `O(log(n))` like `BTreeMap`.
+//! * Insertion is generally `O(n)`, but optimized to `O(1)` if the new item sorts greater than the current items. In `BTreeMap` it's `O(log(n))`.
+//! * Deletion is `O(n)` whereas `BTreeMap` is `O(log(n))`.
+//! * Bulk operations like `from_iter`, `extend` and deserialization have an optimized `O(n)` path
+//!    for inputs that are ordered and `O(n*log(n))` complexity otherwise.
+//!
 //! ## Pluggable Backends
 //!
 //! By default, [`LiteMap`] is backed by a [`Vec`]; however, it can be backed by any appropriate