libdd_trace_utils/span/
vec_map.rs

1// Copyright 2026-Present Datadog, Inc. https://www.datadoghq.com/
2// SPDX-License-Identifier: Apache-2.0
3
4//! This module defines a associative map datastructure for spans data (meta, metrics, etc.) backed
5//! by a vector. Spans are mostly allocated and constructed, and more rarely read or mutated.
6//! [VecMap] is thus optimized for insertion (which is just `Vec::push`), without any hashing
7//! involved. Fetching and removing a value is, on the other hand, linear time in the size of the
8//! map. However, since meta and metrics are expected to be typically small (20ish elements or
9//! less), linear scan is usually still competitive with hashmap's `get`.
10
11use serde::ser::{Serialize, Serializer};
12use std::borrow::Borrow;
13use std::collections::{HashMap, HashSet};
14use std::hash::Hash;
15use std::slice;
16use std::sync::atomic::{AtomicBool, Ordering};
17
18/// A Vec-backed map that provides HashMap-like lookup by key.
19///
20/// # Duplicates
21///
22/// Duplicates are tolerated: [VecMap::insert] always appends, and [VecMap::get]/[VecMap::get_mut]
23/// return the *last* matching entry so that later writes shadow earlier ones. This optimizes for
24/// fast insertion and construction (that might happen on the client's application hot path),
25/// avoiding a linear scan on each insert, or a potential full re-hashing with a hashmap.
26/// Additionally, while overriding a metric or a meta definitively happens, it's assumed to be rare
27/// enough so such that the size penalty of duplication is expected to be reasonable.
28///
29/// **Important**: note that only [VecMap::get] and [VecMap::get_mut] are duplicate-aware, so to
30/// speak. [VecMap::len], [VecMap::iter], and others just delegates to the underlying `Vec`, and
31/// won't deduplicate.
32///
33/// Explicit deduplication is currently being done on-demand by [VecMap::dedup]. An internal flag is
34/// used to avoid undue deduplication (see [VecMap::dedup]). `VecMap` is automatically deduped
35/// before serialization.
36///
37/// In the future, we could trigger deduplication on other events, for example at insertion if the
38/// size is bigger than a threshold (and we haven't deduped for `x` operations).
39///
40/// # Ordering
41///
42/// As this is a map, iteration order is not defined nor guaranteed. In practice, iteration follows
43/// insertion order, but [Self::dedup] will reverse the underlying vector.
44#[derive(Clone, Debug)]
45pub struct VecMap<K, V> {
46    data: Vec<(K, V)>,
47    /// Deduped is a flag that is set after entry deduplication. It is dirtied (set to `false`)
48    /// when any modification that could create duplicates is performed (`deduped == false`
49    /// doesn't imply there are actual duplicates, just than there might be). This is useful to
50    /// avoid performing deduplication several times in a row, for example in the export
51    /// pipeline.
52    deduped: bool,
53}
54
55impl<K, V> Default for VecMap<K, V> {
56    fn default() -> Self {
57        Self {
58            data: Default::default(),
59            deduped: false,
60        }
61    }
62}
63
64// This implementation allocates, which isn't expected for equality testing (not allocating would be
65// rather tedious though). It's fine for tests (where `PartialEq` is currently needed), so we
66// cfg-gate it to avoid unintended usage in prod.
67#[cfg(any(test, feature = "test-utils"))]
68impl<K: Eq + Hash, V: PartialEq> PartialEq for VecMap<K, V> {
69    fn eq(&self, other: &Self) -> bool {
70        let lhs: HashMap<&K, &V> = self.data.iter().map(|(k, v)| (k, v)).collect();
71        let rhs: HashMap<&K, &V> = other.data.iter().map(|(k, v)| (k, v)).collect();
72        lhs == rhs
73    }
74}
75
76#[cfg(any(test, feature = "test-utils"))]
77impl<K: Eq + Hash, V: Eq> Eq for VecMap<K, V> {}
78
79impl<K, V> VecMap<K, V> {
80    #[must_use]
81    #[inline]
82    pub fn new() -> Self {
83        Self::default()
84    }
85
86    /// Dirty the `dedup` flag after a mutation that could introduce duplicates.
87    fn dirty(&mut self) {
88        self.deduped = false;
89    }
90
91    #[must_use]
92    #[inline]
93    pub fn with_capacity(capacity: usize) -> Self {
94        VecMap {
95            data: Vec::with_capacity(capacity),
96            deduped: false,
97        }
98    }
99
100    #[inline]
101    pub fn insert(&mut self, key: K, value: V) {
102        self.data.push((key, value));
103        self.dirty();
104    }
105
106    #[inline]
107    pub fn get<Q>(&self, key: &Q) -> Option<&V>
108    where
109        K: Borrow<Q>,
110        Q: ?Sized + PartialEq,
111    {
112        self.data
113            .iter()
114            .rev()
115            .find(|(k, _)| k.borrow() == key)
116            .map(|(_, v)| v)
117    }
118
119    #[inline]
120    pub fn get_mut<Q>(&mut self, key: &Q) -> Option<&mut V>
121    where
122        K: Borrow<Q>,
123        Q: ?Sized + PartialEq,
124    {
125        self.data
126            .iter_mut()
127            .rev()
128            .find(|(k, _)| (*k).borrow() == key)
129            .map(|(_, v)| v)
130    }
131
132    #[inline]
133    pub fn contains_key<Q>(&self, key: &Q) -> bool
134    where
135        K: Borrow<Q>,
136        Q: ?Sized + PartialEq,
137    {
138        self.data.iter().any(|(k, _)| k.borrow() == key)
139    }
140
141    /// Remove all entries matching this key from the map. This method uses [Vec::retain], and is
142    /// thus potentially costly (like any removal in a vector-like datastructure).
143    // Note: we might implement a tombstone or option-based deletion later, if removal is a bit too
144    // costly.
145    #[inline]
146    pub fn remove_slow<Q>(&mut self, key: &Q)
147    where
148        K: Borrow<Q>,
149        Q: ?Sized + PartialEq,
150    {
151        self.data.retain(|(k, _)| k.borrow() != key);
152    }
153
154    /// Iterate over the element, including duplicate entries.
155    #[inline]
156    pub fn iter(&self) -> slice::Iter<'_, (K, V)> {
157        self.data.iter()
158    }
159
160    /// Iterate mutably over the elements, including duplicate entries.
161    #[inline]
162    pub fn iter_mut(&mut self) -> slice::IterMut<'_, (K, V)> {
163        self.dirty();
164        self.data.iter_mut()
165    }
166
167    /// Return the length of the underlying vector, thus including duplicate entries.
168    #[inline]
169    pub fn len(&self) -> usize {
170        self.data.len()
171    }
172
173    #[inline]
174    pub fn is_empty(&self) -> bool {
175        self.data.is_empty()
176    }
177
178    /// Return `true` if the map hasn't been extended since the last call to [Self::dedup],
179    /// guaranteeing that the underlying vector doesn't have any duplicate key.
180    ///
181    /// If `is_deduped` returns `false`, the map may have duplicate keys.
182    #[inline]
183    pub fn is_deduped(&self) -> bool {
184        self.deduped
185    }
186
187    /// Assert, without scanning, that this map holds no duplicate keys, setting the `deduped` flag.
188    ///
189    /// For builders whose source guarantees key uniqueness (e.g. msgpack decoding, where the wire
190    /// format is a map), to skip the [Self::dedup] pass. A later mutation re-dirties the flag.
191    ///
192    /// **Caution**: if the source can actually contain duplicate keys, prefer [Self::dedup].
193    #[inline]
194    pub fn mark_deduped(&mut self) {
195        self.deduped = true;
196    }
197
198    #[inline]
199    pub fn clear(&mut self) {
200        self.data.clear()
201    }
202
203    #[inline]
204    pub fn drain<R: std::ops::RangeBounds<usize>>(
205        &mut self,
206        range: R,
207    ) -> std::vec::Drain<'_, (K, V)> {
208        self.data.drain(range)
209    }
210}
211
212impl<K: Eq + Hash, V> VecMap<K, V> {
213    /// Returns a deduped map, that either borrows from `self` without performing any work if the
214    /// map is already deduped, or dedup the entries in a new separate vec otherwise. As opposed to
215    /// [Self::dedup], `as_deduped_map` takes an immutable reference to `self` but might allocate.
216    /// Prefer [Self::dedup] when applicable.
217    pub fn as_deduped_map(&self) -> DedupedVecMap<'_, K, V> {
218        if self.deduped {
219            DedupedVecMap::Borrowed(self)
220        } else {
221            DedupedVecMap::Owned(self.data.iter().map(|(k, v)| (k, v)).collect())
222        }
223    }
224
225    /// This is a convenience wrapper around [Self::as_deduped_map] used in the msgpack encoder,
226    /// where we expect the map to be deduped, but call `as_deduped_map` as a defensive measure. If
227    /// the latter had to deduplicate and allocate a new vec, we log a warning (at most once).
228    pub fn defensive_dedup(&self) -> DedupedVecMap<'_, K, V> {
229        if !self.is_deduped() {
230            static WARNED: AtomicBool = AtomicBool::new(false);
231            if !WARNED.swap(true, Ordering::Relaxed) {
232                tracing::warn!(
233                    "VecMap not deduped before encoding. Performing defensive on-the-fly dedup"
234                );
235            }
236        }
237
238        self.as_deduped_map()
239    }
240
241    /// Remove entries with a duplicate key, only keeping the last one. After this, a flag is set
242    /// internally, such that as long as the map isn't extended or mutably iterated, the next
243    /// [Self::dedup] doesn't perform the work again.
244    pub fn dedup(&mut self) {
245        if self.deduped {
246            return;
247        }
248
249        self.data.reverse();
250
251        // Since we're going to shuffle elements around, it's not easy to keep references to keys in
252        // the deduping set while deleting some of them, since deletion in a vec shifts all other
253        // elements after it, invalidating references. When we finally call `retain`, we must not
254        // hold any reference to vecmap elements anymore.
255        //
256        // The following approaches are possible:
257        //
258        // - clone the keys in the hashset. Alas, we don't want a `Clone` bound on `SpanText` (which
259        //   is the type of keys in practice), as some representations can be expensive to clone,
260        //   e.g. requiring to lock the GIL in Python (python-native, reference counted strings)
261        // - a two-pass approach. In a first pass we store seen key references in a HashSet and
262        //   build a bitmap of indices to keep. Once built, we can release the set and call
263        //   `Vec::retain` without borrowing issues. It's safe but requires an additional pass over
264        //   the vecmap and an auxiliary `Vec<bool>`, in addition to the hashset.
265        // - an unsafe, one-pass approach: if we re-implement a custom `retain`, we can store key
266        //   references in an auxiliary HashSet that are guaranteed to remain valid as we move
267        //   elements: we first move an element to keep at their final location, and only then
268        //   insert a pointer to the key in the `seen` hashmap, which will remain valid.
269        //
270        // We choose the two-pass approach, which is simpler, safe and reasonably fast. If needed in
271        // the future, the unsafe one-pass approach can be implemented.
272        let keep: Vec<bool> = {
273            let mut seen = HashSet::with_capacity(self.len());
274            self.data.iter().map(|(k, _)| seen.insert(k)).collect()
275        };
276
277        let mut keep = keep.into_iter();
278        self.data.retain(|_| keep.next().unwrap_or(false));
279
280        self.deduped = true;
281    }
282}
283
284impl<K, V> From<Vec<(K, V)>> for VecMap<K, V> {
285    fn from(data: Vec<(K, V)>) -> Self {
286        Self {
287            data,
288            deduped: false,
289        }
290    }
291}
292
293impl<K, V> From<VecMap<K, V>> for Vec<(K, V)> {
294    fn from(value: VecMap<K, V>) -> Self {
295        value.data
296    }
297}
298
299impl<K, V> FromIterator<(K, V)> for VecMap<K, V> {
300    fn from_iter<I: IntoIterator<Item = (K, V)>>(iter: I) -> Self {
301        Self {
302            data: iter.into_iter().collect(),
303            deduped: false,
304        }
305    }
306}
307
308impl<K, V> IntoIterator for VecMap<K, V> {
309    type Item = (K, V);
310    type IntoIter = std::vec::IntoIter<(K, V)>;
311
312    fn into_iter(self) -> Self::IntoIter {
313        self.data.into_iter()
314    }
315}
316
317impl<'a, K, V> IntoIterator for &'a VecMap<K, V> {
318    type Item = &'a (K, V);
319    type IntoIter = slice::Iter<'a, (K, V)>;
320
321    fn into_iter(self) -> Self::IntoIter {
322        self.data.iter()
323    }
324}
325
326impl<'a, K, V> IntoIterator for &'a mut VecMap<K, V> {
327    type Item = &'a mut (K, V);
328    type IntoIter = slice::IterMut<'a, (K, V)>;
329
330    fn into_iter(self) -> Self::IntoIter {
331        // Since we iterate on keys as well, they can modified, and introduce duplicates.
332        self.dirty();
333        self.data.iter_mut()
334    }
335}
336
337impl<K, V> Extend<(K, V)> for VecMap<K, V> {
338    fn extend<I: IntoIterator<Item = (K, V)>>(&mut self, iter: I) {
339        self.dirty();
340        self.data.extend(iter);
341    }
342}
343
344impl<K: Serialize + Eq + Hash, V: Serialize> Serialize for VecMap<K, V> {
345    fn serialize<S: Serializer>(&self, serializer: S) -> Result<S::Ok, S::Error> {
346        use serde::ser::SerializeMap;
347
348        let deduped = self.as_deduped_map();
349        let mut map_ser = serializer.serialize_map(Some(deduped.len()))?;
350
351        for (k, v) in deduped.iter() {
352            map_ser.serialize_entry(k, v)?;
353        }
354
355        map_ser.end()
356    }
357}
358
359pub enum DedupedVecMap<'a, K, V> {
360    Borrowed(&'a VecMap<K, V>),
361    Owned(HashMap<&'a K, &'a V>),
362}
363
364impl<'a, K, V> DedupedVecMap<'a, K, V> {
365    #[inline]
366    pub fn iter(&self) -> DedupedVecMapIter<'_, 'a, K, V> {
367        match self {
368            DedupedVecMap::Borrowed(vec_map) => DedupedVecMapIter::Borrowed(vec_map.iter()),
369            DedupedVecMap::Owned(map) => DedupedVecMapIter::Owned(map.iter()),
370        }
371    }
372
373    #[inline]
374    pub fn len(&self) -> usize {
375        match self {
376            DedupedVecMap::Borrowed(vec_map) => vec_map.len(),
377            DedupedVecMap::Owned(map) => map.len(),
378        }
379    }
380
381    #[inline]
382    pub fn is_empty(&self) -> bool {
383        match self {
384            DedupedVecMap::Borrowed(vec_map) => vec_map.is_empty(),
385            DedupedVecMap::Owned(map) => map.is_empty(),
386        }
387    }
388}
389
390pub enum DedupedVecMapIter<'b, 'a: 'b, K, V> {
391    Borrowed(slice::Iter<'a, (K, V)>),
392    Owned(std::collections::hash_map::Iter<'b, &'a K, &'a V>),
393}
394
395impl<'b, 'a: 'b, K, V> Iterator for DedupedVecMapIter<'b, 'a, K, V> {
396    type Item = (&'a K, &'a V);
397
398    fn next(&mut self) -> Option<Self::Item> {
399        match self {
400            DedupedVecMapIter::Borrowed(iter) => iter.next().map(|(k, v)| (k, v)),
401            DedupedVecMapIter::Owned(iter) => iter.next().map(|(&k, &v)| (k, v)),
402        }
403    }
404
405    fn size_hint(&self) -> (usize, Option<usize>) {
406        match self {
407            DedupedVecMapIter::Borrowed(iter) => iter.size_hint(),
408            DedupedVecMapIter::Owned(iter) => iter.size_hint(),
409        }
410    }
411}
412
413impl<'b, 'a: 'b, K, V> ExactSizeIterator for DedupedVecMapIter<'b, 'a, K, V> {}
414
415#[cfg(test)]
416mod tests {
417    use super::*;
418
419    #[test]
420    fn get_returns_last_inserted() {
421        let mut m = VecMap::new();
422        m.insert("a", 1);
423        m.insert("a", 2);
424        assert_eq!(m.get("a"), Some(&2));
425    }
426
427    #[test]
428    fn get_mut_returns_last_inserted() {
429        let mut m = VecMap::new();
430        m.insert("a", 1);
431        m.insert("a", 2);
432        *m.get_mut("a").unwrap() = 42;
433        assert_eq!(m.get("a"), Some(&42));
434        // First entry unchanged
435        assert_eq!(m.iter().next().unwrap().1, 1);
436    }
437
438    #[test]
439    fn remove_removes_all_occurrences() {
440        let mut m = VecMap::new();
441        m.insert("a", 1);
442        m.insert("b", 2);
443        m.insert("a", 3);
444        m.remove_slow("a");
445        assert_eq!(m.get("a"), None);
446        assert!(!m.contains_key("a"));
447        assert_eq!(m.len(), 1);
448    }
449
450    #[test]
451    fn contains_key_works() {
452        let mut m = VecMap::new();
453        assert!(!m.contains_key("x"));
454        m.insert("x", 10);
455        assert!(m.contains_key("x"));
456    }
457
458    #[test]
459    fn from_iterator() {
460        let m: VecMap<&str, i32> = vec![("a", 1), ("b", 2)].into_iter().collect();
461        assert_eq!(m.len(), 2);
462        assert_eq!(m.get("b"), Some(&2));
463    }
464
465    #[test]
466    fn into_iter_consuming() {
467        let mut m = VecMap::new();
468        m.insert("a", 1);
469        m.insert("b", 2);
470        let pairs: Vec<_> = m.into_iter().collect();
471        assert_eq!(pairs, vec![("a", 1), ("b", 2)]);
472    }
473
474    #[test]
475    fn is_deduped_false_initially() {
476        let m: VecMap<&str, i32> = VecMap::new();
477        assert!(!m.is_deduped());
478    }
479
480    #[test]
481    fn is_deduped_false_after_from() {
482        let m: VecMap<&str, i32> = vec![("a", 1)].into();
483        assert!(!m.is_deduped());
484    }
485
486    #[test]
487    fn is_deduped_false_after_collect() {
488        let m: VecMap<&str, i32> = vec![("a", 1)].into_iter().collect();
489        assert!(!m.is_deduped());
490    }
491
492    #[test]
493    fn dedup_sets_flag() {
494        let mut m = VecMap::new();
495        m.insert("a", 1);
496        assert!(!m.is_deduped());
497        m.dedup();
498        assert!(m.is_deduped());
499    }
500
501    #[test]
502    fn dedup_on_empty_map() {
503        let mut m: VecMap<String, i32> = VecMap::new();
504        m.dedup();
505        assert!(m.is_deduped());
506        assert!(m.is_empty());
507    }
508
509    #[test]
510    fn dedup_no_duplicates() {
511        let mut m = VecMap::new();
512        m.insert("a", 1);
513        m.insert("b", 2);
514        m.insert("c", 3);
515        m.dedup();
516        assert_eq!(m.len(), 3);
517        assert_eq!(m.get("a"), Some(&1));
518        assert_eq!(m.get("b"), Some(&2));
519        assert_eq!(m.get("c"), Some(&3));
520    }
521
522    #[test]
523    fn dedup_keeps_last_value() {
524        let mut m = VecMap::new();
525        m.insert("a", 1);
526        m.insert("b", 10);
527        m.insert("a", 2);
528        m.insert("a", 3);
529        m.insert("b", 20);
530        m.dedup();
531        assert_eq!(m.len(), 2);
532        assert_eq!(m.get("a"), Some(&3));
533        assert_eq!(m.get("b"), Some(&20));
534    }
535
536    #[test]
537    fn dedup_is_idempotent() {
538        let mut m = VecMap::new();
539        m.insert("a", 1);
540        m.insert("a", 2);
541        m.dedup();
542        assert!(m.is_deduped());
543        assert_eq!(m.len(), 1);
544        m.dedup();
545        assert!(m.is_deduped());
546        assert_eq!(m.len(), 1);
547        assert_eq!(m.get("a"), Some(&2));
548    }
549
550    #[test]
551    fn insert_dirties_dedup_flag() {
552        let mut m = VecMap::new();
553        m.insert("a", 1);
554        m.dedup();
555        assert!(m.is_deduped());
556
557        m.insert("b", 2);
558        assert!(!m.is_deduped());
559    }
560
561    #[test]
562    fn extend_dirties_dedup_flag() {
563        let mut m = VecMap::new();
564        m.insert("a", 1);
565        m.dedup();
566        assert!(m.is_deduped());
567
568        m.extend(vec![("b", 2)]);
569        assert!(!m.is_deduped());
570    }
571
572    #[test]
573    fn iter_mut_dirties_dedup_flag() {
574        let mut m = VecMap::new();
575        m.insert("a", 1);
576        m.dedup();
577        assert!(m.is_deduped());
578
579        for (_, v) in m.iter_mut() {
580            *v += 1;
581        }
582
583        assert!(!m.is_deduped());
584    }
585
586    #[test]
587    fn deduped_vec_map_borrowed_iter_and_len() {
588        let mut m = VecMap::new();
589        m.insert("a", 1);
590        m.insert("b", 2);
591        m.dedup();
592
593        let d = m.defensive_dedup();
594        assert!(matches!(d, DedupedVecMap::Borrowed(_)));
595        assert_eq!(d.len(), 2);
596
597        let mut items: Vec<_> = d.iter().collect();
598        items.sort_by_key(|(k, _)| **k);
599        assert_eq!(items, vec![(&"a", &1), (&"b", &2)]);
600    }
601
602    #[test]
603    fn deduped_vec_map_copy_iter_and_len() {
604        let mut m = VecMap::new();
605        m.insert("a", 1);
606        m.insert("a", 2);
607        m.insert("b", 3);
608
609        let d = m.defensive_dedup();
610        assert!(matches!(d, DedupedVecMap::Owned(_)));
611        assert_eq!(d.len(), 2);
612
613        let items: HashMap<&&str, &i32> = d.iter().collect();
614        assert_eq!(items[&"a"], &2);
615        assert_eq!(items[&"b"], &3);
616    }
617
618    #[test]
619    fn deduped_vec_map_iter_exact_size() {
620        let mut m = VecMap::new();
621        m.insert("a", 1);
622        m.insert("b", 2);
623        m.insert("c", 3);
624        m.dedup();
625
626        let d = m.defensive_dedup();
627        let mut iter = d.iter();
628        assert_eq!(iter.len(), 3);
629        iter.next();
630        assert_eq!(iter.len(), 2);
631    }
632
633    #[test]
634    fn deduped_vec_map_empty() {
635        let m: VecMap<String, i32> = VecMap::new();
636        let d = DedupedVecMap::Borrowed(&m);
637        assert_eq!(d.len(), 0);
638        assert_eq!(d.iter().count(), 0);
639    }
640
641    #[test]
642    fn dedup_does_not_require_clone() {
643        #[derive(Debug, PartialEq, Eq, Hash)]
644        struct NonCloneKey(u32);
645
646        let mut m = VecMap::new();
647        m.insert(NonCloneKey(1), "a");
648        m.insert(NonCloneKey(2), "b");
649        m.insert(NonCloneKey(1), "c");
650        m.dedup();
651        assert_eq!(m.len(), 2);
652        assert_eq!(m.get(&NonCloneKey(1)), Some(&"c"));
653        assert_eq!(m.get(&NonCloneKey(2)), Some(&"b"));
654    }
655
656    #[test]
657    fn serialize_deduplicates_keeping_last() {
658        let mut m = VecMap::new();
659        m.insert("a", 0);
660        m.insert("b", 0);
661        m.insert("b", 1);
662        m.insert("a", 1);
663        m.insert("a", 3);
664        m.insert("b", 2);
665
666        let serialized: serde_json::Value = serde_json::to_value(&m).unwrap();
667        let obj = serialized.as_object().unwrap();
668
669        assert_eq!(obj.len(), 2);
670        assert_eq!(obj.get("a").unwrap(), 3);
671        assert_eq!(obj.get("b").unwrap(), 2);
672    }
673}
libdd_trace_utils/span/vec_map.rs

libdd_trace_utils/span/
vec_map.rs
