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

libdd_trace_utils/span/
vec_map.rs
