prefix_array/
shared.rs

1//! Shared implementations between [`PrefixArray`](super::PrefixArray) and
2//! [`PrefixArraySet`](super::PrefixArraySet).
3extern crate alloc;
4
5use core::{borrow::Borrow, cmp::Ordering, mem};
6
7use alloc::vec::Vec;
8
9mod vec_ext;
10use vec_ext::InsertMany;
11
12/// Sealed module to prevent `PrefixMapOrSet` leaking
13mod sealed {
14
15    use core::borrow::Borrow;
16
17    /// private-public bounding trait representing any valid `PrefixArray` or `PrefixArraySet`
18    pub trait PrefixMapOrSet {
19        /// The Data that this collection stores for scratch space
20        type Item;
21    }
22
23    impl<K: Borrow<str>, V> PrefixMapOrSet for crate::PrefixArray<K, V> {
24        type Item = (K, V);
25    }
26
27    impl<K: Borrow<str>> PrefixMapOrSet for crate::PrefixArraySet<K> {
28        type Item = K;
29    }
30}
31
32use sealed::PrefixMapOrSet;
33
34use crate::{PrefixArray, PrefixArraySet, SetSubSlice, SubSlice};
35
36/// Scratch space for the `extend_with` methods on [`PrefixArray`](super::PrefixArray) and [`PrefixArraySet`](super::PrefixArraySet)
37///
38/// Using this scratch space allows you to potentially call extend multiple times without
39/// allocating, as the default `Extend` method will allocate to avoid severe time penalties.
40pub struct ScratchSpace<T: PrefixMapOrSet>(pub(crate) Vec<(usize, T::Item)>);
41
42impl<T: PrefixMapOrSet> Default for ScratchSpace<T> {
43    fn default() -> Self {
44        Self::new()
45    }
46}
47
48impl<T: PrefixMapOrSet> ScratchSpace<T> {
49    /// Creates a new empty scratch space
50    ///
51    /// # Examples
52    /// ```rust
53    /// # use prefix_array::{PrefixArraySet, shared::ScratchSpace};
54    /// // equivalent to .extend(["foo", "bar"]) as we dont preallocate
55    /// PrefixArraySet::new().extend_with(&mut ScratchSpace::new(), ["foo", "bar"]);
56    /// ```
57    #[must_use]
58    pub const fn new() -> Self {
59        Self(Vec::new())
60    }
61
62    /// Creates a scratch space with room to insert at least n elements before reallocating
63    #[must_use]
64    pub fn with_capacity(n: usize) -> Self {
65        Self(Vec::with_capacity(n))
66    }
67}
68
69/// A trait that abstracts over `PrefixArray` and `PrefixArraySet` owned variants
70/// This allows defining implementations once
71pub(crate) trait PrefixOwned<V>: Sized {
72    /// The data that this `PrefixOwned` stores internally
73    type Data;
74
75    /// Builds a new owned instance from a Vec of backing Data
76    fn construct(v: Vec<Self::Data>) -> Self;
77
78    /// Replaces only the value portion in a Data, may be a no op for valueless data
79    fn replace_value(dst: &mut Self::Data, src: Self::Data) -> V;
80
81    /// Returns the indexable string from this Data
82    fn as_str(v: &Self::Data) -> &str;
83
84    /// Gets mutable reference to inner Vec
85    fn get_vec_mut(&mut self) -> &mut Vec<Self::Data>;
86
87    /// Convenience fn to compare two Data by their string repr
88    fn cmp(a: &Self::Data, b: &Self::Data) -> Ordering {
89        Self::as_str(a).cmp(Self::as_str(b))
90    }
91
92    /// Convenience fn to compare two Data by their string repr
93    fn eq(a: &Self::Data, b: &Self::Data) -> bool {
94        Self::cmp(a, b).is_eq()
95    }
96
97    /// Implements the `from_vec_lossy` impl that sorts and discards duplicates
98    fn from_vec_lossy_impl(mut v: Vec<Self::Data>) -> Self {
99        v.sort_unstable_by(|f, s| Self::as_str(f).cmp(Self::as_str(s)));
100        v.dedup_by(|f, s| Self::as_str(f) == Self::as_str(s));
101
102        Self::construct(v)
103    }
104
105    /// Implements an `insert` impl that calls a replacer function when a value was already in the
106    /// collection
107    fn insert_internal<F, T>(&mut self, replacer: F, data: Self::Data) -> Option<T>
108    where
109        F: Fn(&mut Self::Data, Self::Data) -> T,
110    {
111        let vec = self.get_vec_mut();
112
113        match vec.binary_search_by_key(&Self::as_str(&data), |s| Self::as_str(s)) {
114            Err(idx) => {
115                vec.insert(idx, data);
116                None
117            }
118            Ok(idx) => Some(replacer(&mut vec[idx], data)),
119        }
120    }
121
122    /// Implements the `insert` impl that inserts new key value pairs, and on a preexisting key only
123    /// updates the value
124    fn insert_impl(&mut self, data: Self::Data) -> Option<V> {
125        self.insert_internal(Self::replace_value, data)
126    }
127
128    /// Implements the `insert_replace` impl that inserts a new value and always replaces all of
129    /// the previous value if there was one
130    fn insert_replace_impl(&mut self, data: Self::Data) -> Option<Self::Data> {
131        self.insert_internal(mem::replace, data)
132    }
133
134    /// Implements the `remove_entry` impl that removes something fully with the given key if it exists
135    fn remove_entry_impl(&mut self, key: &str) -> Option<Self::Data> {
136        if let Ok(idx) = self
137            .get_vec_mut()
138            .binary_search_by_key(&key, |s| Self::as_str(s))
139        {
140            Some(self.get_vec_mut().remove(idx))
141        } else {
142            None
143        }
144    }
145
146    /// Implements the `extend_with` impl that takes a scratch space and an iterator of addable data
147    /// to extend the collection
148    fn extend_with_impl<I>(&mut self, insert: &mut Vec<(usize, Self::Data)>, iter: I)
149    where
150        I: IntoIterator<Item = Self::Data>,
151    {
152        let iter = iter.into_iter();
153
154        // clear for correctness, a scratchspace should become empty after an insert_many call and
155        // there is no other way to push to it than this method, but we should prevent the
156        // possibility here anyways
157        insert.clear();
158
159        // speculative optimization, assume that most items are going to be newly inserted
160        // this will over allocate when that is untrue
161        insert.reserve(iter.size_hint().0);
162
163        for k in iter {
164            match self
165                .get_vec_mut()
166                .binary_search_by_key(&Self::as_str(&k), |s| Self::as_str(s))
167            {
168                // add to insertion set
169                Err(idx) => insert.push((idx, k)),
170                // replace old value
171                Ok(idx) => {
172                    Self::replace_value(&mut self.get_vec_mut()[idx], k);
173                }
174            }
175        }
176
177        // presort by string so that identical indexes are mapped correctly
178        insert.sort_unstable_by(|(_, a), (_, b)| Self::cmp(a, b));
179
180        // avoid duplicate K being inserted
181        insert.dedup_by(|(_, a), (_, b)| Self::eq(a, b));
182
183        self.get_vec_mut().insert_many(insert);
184    }
185
186    /// Implements the `from_unique_iter` impl that builds a collection from an iterator that is
187    /// guaranteed to have unique key items
188    fn from_unique_iter_impl<T: IntoIterator<Item = Self::Data>>(v: T) -> Self {
189        let mut unsorted = v.into_iter().collect::<Vec<Self::Data>>();
190        unsorted.sort_unstable_by(|f, s| Self::cmp(f, s));
191
192        Self::construct(unsorted)
193    }
194}
195
196impl<K: Borrow<str>, V> PrefixOwned<V> for PrefixArray<K, V> {
197    type Data = (K, V);
198
199    fn construct(v: Vec<Self::Data>) -> Self {
200        Self(v)
201    }
202
203    fn get_vec_mut(&mut self) -> &mut Vec<Self::Data> {
204        &mut self.0
205    }
206
207    fn replace_value(dst: &mut Self::Data, src: Self::Data) -> V {
208        core::mem::replace(&mut dst.1, src.1)
209    }
210
211    fn as_str(v: &Self::Data) -> &str {
212        v.0.borrow()
213    }
214}
215
216impl<K: Borrow<str>> PrefixOwned<()> for PrefixArraySet<K> {
217    type Data = K;
218
219    fn construct(v: Vec<Self::Data>) -> Self {
220        Self(v)
221    }
222
223    fn as_str(v: &Self::Data) -> &str {
224        v.borrow()
225    }
226
227    // no op
228    fn replace_value(_dst: &mut Self::Data, _src: Self::Data) {}
229
230    fn get_vec_mut(&mut self) -> &mut Vec<Self::Data> {
231        &mut self.0
232    }
233}
234
235use core::slice::SliceIndex;
236
237/// Internal `DuplicatesPresent` error, this location may become public in the future to allow
238/// map/set to both define `from_mut_slice` more seamlessly
239pub(crate) struct DuplicatesPresent<'a>(pub(crate) &'a str);
240
241/// A trait that abstracts over `PrefixArray` and `PrefixArraySet` borrowed variants
242/// This allows defining implementations once
243pub(crate) trait PrefixBorrowed {
244    /// The data that this `PrefixBorrowed` stores internally
245    type Data;
246
247    /// gets inner slice mutably
248    fn get_mut_slice(&mut self) -> &mut [Self::Data];
249    /// gets inner slice
250    fn get_slice(&self) -> &[Self::Data];
251
252    /// Builds Self from a mutable slice
253    fn cast_from_slice_mut(s: &mut [Self::Data]) -> &mut Self;
254    /// Builds Self from a slice
255    fn cast_from_slice(s: &[Self::Data]) -> &Self;
256
257    /// Gets the indexable string from the given Data
258    fn as_str(v: &Self::Data) -> &str;
259
260    /// Convenience fn to compare Data by its indexable string
261    fn cmp(a: &Self::Data, b: &Self::Data) -> Ordering {
262        Self::as_str(a).cmp(Self::as_str(b))
263    }
264
265    /// Convenience fn to compare Data by its indexable string
266    fn eq(a: &Self::Data, b: &Self::Data) -> bool {
267        Self::cmp(a, b).is_eq()
268    }
269
270    /// reslices self, panics on oob
271    fn reslice<I: SliceIndex<[Self::Data], Output = [Self::Data]>>(&self, i: I) -> &Self {
272        Self::cast_from_slice(&self.get_slice()[i])
273    }
274
275    /// reslices self, panics on oob
276    fn reslice_mut<I: SliceIndex<[Self::Data], Output = [Self::Data]>>(
277        &mut self,
278        i: I,
279    ) -> &mut Self {
280        Self::cast_from_slice_mut(&mut self.get_mut_slice()[i])
281    }
282
283    /// Implements builds a &mut Self from a slice of Data, validating that it can be a Self and
284    /// sorting internally
285    fn from_mut_slice_impl(data: &mut [Self::Data]) -> Result<&mut Self, DuplicatesPresent<'_>> {
286        data.sort_unstable_by(|a, b| Self::cmp(a, b));
287
288        if data.len() <= 1 {
289            return Ok(Self::cast_from_slice_mut(data));
290        }
291
292        let mut error = None;
293
294        for (idx, d) in data.windows(2).enumerate() {
295            if Self::eq(&d[0], &d[1]) {
296                error = Some(idx);
297                break;
298            }
299        }
300
301        match error {
302            Some(idx) => Err(DuplicatesPresent(Self::as_str(&data[idx]))),
303            None => Ok(Self::cast_from_slice_mut(data)),
304        }
305    }
306
307    /// Finds all items with the given prefix using binary search
308    fn find_all_with_prefix_idx_impl(&self, prefix: &str) -> core::ops::Range<usize> {
309        // skip comparisons if we have a unit prefix
310        if prefix.is_empty() {
311            return 0..self.get_slice().len();
312        }
313
314        if let Ok(start) = self.get_slice().binary_search_by(|s| {
315            if Self::as_str(s).starts_with(prefix) {
316                Ordering::Equal
317            } else {
318                Self::as_str(s).cmp(prefix)
319            }
320        }) {
321            let min =
322                self.get_slice()[..start].partition_point(|s| !Self::as_str(s).starts_with(prefix));
323            let max = self.get_slice()[start..]
324                .partition_point(|s| Self::as_str(s).starts_with(prefix))
325                + start;
326
327            min..max
328        } else {
329            0..0
330        }
331    }
332
333    /// Returns the common prefix of all of the keys in the collection
334    fn common_prefix_impl(&self) -> &str {
335        let Some(first) = self.get_slice().first().map(|s| Self::as_str(s)) else {
336            return "";
337        };
338
339        let Some(last) = self.get_slice().last().map(|s| Self::as_str(s)) else {
340            return "";
341        };
342
343        let last_idx = first.len().min(last.len());
344
345        let mut end_idx = 0;
346
347        for ((idx, fch), lch) in first
348            .char_indices()
349            .zip(last.chars())
350            .chain([((last_idx, ' '), ' ')])
351        {
352            end_idx = idx;
353            if fch != lch {
354                break;
355            }
356        }
357
358        &first[..end_idx]
359    }
360
361    /// Returns whether the collection contains the given key
362    fn contains_key_impl(&self, key: &str) -> bool {
363        self.get_slice()
364            .binary_search_by_key(&key, |s| Self::as_str(s))
365            .is_ok()
366    }
367
368    /// View the contents of an entry if there is one by the given key
369    fn get_impl(&self, key: &str) -> Option<&Self::Data> {
370        match self
371            .get_slice()
372            .binary_search_by_key(&key, |s| Self::as_str(s))
373        {
374            Ok(idx) => Some(&self.get_slice()[idx]),
375            Err(_) => None,
376        }
377    }
378
379    /// View the contents of an entry mutably if there is one by the given key
380    fn get_mut_impl(&mut self, key: &str) -> Option<&mut Self::Data> {
381        match self
382            .get_slice()
383            .binary_search_by_key(&key, |s| Self::as_str(s))
384        {
385            Ok(idx) => Some(&mut self.get_mut_slice()[idx]),
386            Err(_) => None,
387        }
388    }
389}
390
391impl<K: Borrow<str>, V> PrefixBorrowed for SubSlice<K, V> {
392    type Data = (K, V);
393
394    fn get_mut_slice(&mut self) -> &mut [Self::Data] {
395        &mut self.0
396    }
397    fn get_slice(&self) -> &[Self::Data] {
398        &self.0
399    }
400
401    fn cast_from_slice_mut(s: &mut [Self::Data]) -> &mut Self {
402        Self::cast_from_slice_mut_core(s)
403    }
404    fn cast_from_slice(s: &[Self::Data]) -> &Self {
405        Self::cast_from_slice_core(s)
406    }
407
408    fn as_str(v: &Self::Data) -> &str {
409        v.0.borrow()
410    }
411}
412
413impl<K: Borrow<str>> PrefixBorrowed for SetSubSlice<K> {
414    type Data = K;
415
416    fn get_mut_slice(&mut self) -> &mut [Self::Data] {
417        &mut self.0
418    }
419
420    fn get_slice(&self) -> &[Self::Data] {
421        &self.0
422    }
423
424    fn cast_from_slice_mut(s: &mut [Self::Data]) -> &mut Self {
425        Self::cast_from_slice_mut_core(s)
426    }
427
428    fn cast_from_slice(s: &[Self::Data]) -> &Self {
429        Self::cast_from_slice_core(s)
430    }
431
432    fn as_str(v: &Self::Data) -> &str {
433        v.borrow()
434    }
435}
prefix_array/shared.rs

prefix_array/
shared.rs
