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
188impl<K: Eq + Hash, V> VecMap<K, V> {
189    /// Returns a deduped map, that either borrows from `self` without performing any work if the
190    /// map is already deduped, or dedup the entries in a new separate vec otherwise. As opposed to
191    /// [Self::dedup], `as_deduped_map` takes an immutable reference to `self` but might allocate.
192    /// Prefer [Self::dedup] when applicable.
193    pub fn as_deduped_map(&self) -> DedupedVecMap<'_, K, V> {
194        if self.deduped {
195            DedupedVecMap::Borrowed(self)
196        } else {
197            DedupedVecMap::Owned(self.data.iter().map(|(k, v)| (k, v)).collect())
198        }
199    }
200
201    /// This is a convenience wrapper around [Self::as_deduped_map] used in the msgpack encoder,
202    /// where we expect the map to be deduped, but call `as_deduped_map` as a defensive measure. If
203    /// the latter had to deduplicate and allocate a new vec, we log a warning (at most once).
204    #[allow(unused)]
205    pub(crate) fn defensive_dedup(&self) -> DedupedVecMap<'_, K, V> {
206        if !self.is_deduped() {
207            static WARNED: AtomicBool = AtomicBool::new(false);
208            if !WARNED.swap(true, Ordering::Relaxed) {
209                tracing::warn!(
210                    "VecMap not deduped before encoding. Performing defensive on-the-fly dedup"
211                );
212            }
213        }
214
215        self.as_deduped_map()
216    }
217
218    /// Remove entries with a duplicate key, only keeping the last one. After this, a flag is set
219    /// internally, such that as long as the map isn't extended or mutably iterated, the next
220    /// [Self::dedup] doesn't perform the work again.
221    pub fn dedup(&mut self) {
222        if self.deduped {
223            return;
224        }
225
226        self.data.reverse();
227
228        // Since we're going to shuffle elements around, it's not easy to keep references to keys in
229        // the deduping set while deleting some of them, since deletion in a vec shifts all other
230        // elements after it, invalidating references. When we finally call `retain`, we must not
231        // hold any reference to vecmap elements anymore.
232        //
233        // The following approaches are possible:
234        //
235        // - clone the keys in the hashset. Alas, we don't want a `Clone` bound on `SpanText` (which
236        //   is the type of keys in practice), as some representations can be expensive to clone,
237        //   e.g. requiring to lock the GIL in Python (python-native, reference counted strings)
238        // - a two-pass approach. In a first pass we store seen key references in a HashSet and
239        //   build a bitmap of indices to keep. Once built, we can release the set and call
240        //   `Vec::retain` without borrowing issues. It's safe but requires an additional pass over
241        //   the vecmap and an auxiliary `Vec<bool>`, in addition to the hashset.
242        // - an unsafe, one-pass approach: if we re-implement a custom `retain`, we can store key
243        //   references in an auxiliary HashSet that are guaranteed to remain valid as we move
244        //   elements: we first move an element to keep at their final location, and only then
245        //   insert a pointer to the key in the `seen` hashmap, which will remain valid.
246        //
247        // We choose the two-pass approach, which is simpler, safe and reasonably fast. If needed in
248        // the future, the unsafe one-pass approach can be implemented.
249        let keep: Vec<bool> = {
250            let mut seen = HashSet::with_capacity(self.len());
251            self.data.iter().map(|(k, _)| seen.insert(k)).collect()
252        };
253
254        let mut keep = keep.into_iter();
255        self.data.retain(|_| keep.next().unwrap_or(false));
256
257        self.deduped = true;
258    }
259}
260
261impl<K, V> From<Vec<(K, V)>> for VecMap<K, V> {
262    fn from(data: Vec<(K, V)>) -> Self {
263        Self {
264            data,
265            deduped: false,
266        }
267    }
268}
269
270impl<K, V> From<VecMap<K, V>> for Vec<(K, V)> {
271    fn from(value: VecMap<K, V>) -> Self {
272        value.data
273    }
274}
275
276impl<K, V> FromIterator<(K, V)> for VecMap<K, V> {
277    fn from_iter<I: IntoIterator<Item = (K, V)>>(iter: I) -> Self {
278        Self {
279            data: iter.into_iter().collect(),
280            deduped: false,
281        }
282    }
283}
284
285impl<K, V> IntoIterator for VecMap<K, V> {
286    type Item = (K, V);
287    type IntoIter = std::vec::IntoIter<(K, V)>;
288
289    fn into_iter(self) -> Self::IntoIter {
290        self.data.into_iter()
291    }
292}
293
294impl<'a, K, V> IntoIterator for &'a VecMap<K, V> {
295    type Item = &'a (K, V);
296    type IntoIter = slice::Iter<'a, (K, V)>;
297
298    fn into_iter(self) -> Self::IntoIter {
299        self.data.iter()
300    }
301}
302
303impl<'a, K, V> IntoIterator for &'a mut VecMap<K, V> {
304    type Item = &'a mut (K, V);
305    type IntoIter = slice::IterMut<'a, (K, V)>;
306
307    fn into_iter(self) -> Self::IntoIter {
308        // Since we iterate on keys as well, they can modified, and introduce duplicates.
309        self.dirty();
310        self.data.iter_mut()
311    }
312}
313
314impl<K, V> Extend<(K, V)> for VecMap<K, V> {
315    fn extend<I: IntoIterator<Item = (K, V)>>(&mut self, iter: I) {
316        self.dirty();
317        self.data.extend(iter);
318    }
319}
320
321impl<K: Serialize + Eq + Hash, V: Serialize> Serialize for VecMap<K, V> {
322    fn serialize<S: Serializer>(&self, serializer: S) -> Result<S::Ok, S::Error> {
323        use serde::ser::SerializeMap;
324
325        let deduped = self.as_deduped_map();
326        let mut map_ser = serializer.serialize_map(Some(deduped.len()))?;
327
328        for (k, v) in deduped.iter() {
329            map_ser.serialize_entry(k, v)?;
330        }
331
332        map_ser.end()
333    }
334}
335
336pub enum DedupedVecMap<'a, K, V> {
337    Borrowed(&'a VecMap<K, V>),
338    Owned(HashMap<&'a K, &'a V>),
339}
340
341impl<'a, K, V> DedupedVecMap<'a, K, V> {
342    #[inline]
343    pub fn iter(&self) -> DedupedVecMapIter<'_, 'a, K, V> {
344        match self {
345            DedupedVecMap::Borrowed(vec_map) => DedupedVecMapIter::Borrowed(vec_map.iter()),
346            DedupedVecMap::Owned(map) => DedupedVecMapIter::Owned(map.iter()),
347        }
348    }
349
350    #[inline]
351    pub fn len(&self) -> usize {
352        match self {
353            DedupedVecMap::Borrowed(vec_map) => vec_map.len(),
354            DedupedVecMap::Owned(map) => map.len(),
355        }
356    }
357
358    #[inline]
359    pub fn is_empty(&self) -> bool {
360        match self {
361            DedupedVecMap::Borrowed(vec_map) => vec_map.is_empty(),
362            DedupedVecMap::Owned(map) => map.is_empty(),
363        }
364    }
365}
366
367pub enum DedupedVecMapIter<'b, 'a: 'b, K, V> {
368    Borrowed(slice::Iter<'a, (K, V)>),
369    Owned(std::collections::hash_map::Iter<'b, &'a K, &'a V>),
370}
371
372impl<'b, 'a: 'b, K, V> Iterator for DedupedVecMapIter<'b, 'a, K, V> {
373    type Item = (&'a K, &'a V);
374
375    fn next(&mut self) -> Option<Self::Item> {
376        match self {
377            DedupedVecMapIter::Borrowed(iter) => iter.next().map(|(k, v)| (k, v)),
378            DedupedVecMapIter::Owned(iter) => iter.next().map(|(&k, &v)| (k, v)),
379        }
380    }
381
382    fn size_hint(&self) -> (usize, Option<usize>) {
383        match self {
384            DedupedVecMapIter::Borrowed(iter) => iter.size_hint(),
385            DedupedVecMapIter::Owned(iter) => iter.size_hint(),
386        }
387    }
388}
389
390impl<'b, 'a: 'b, K, V> ExactSizeIterator for DedupedVecMapIter<'b, 'a, K, V> {}
391
392#[cfg(test)]
393mod tests {
394    use super::*;
395
396    #[test]
397    fn get_returns_last_inserted() {
398        let mut m = VecMap::new();
399        m.insert("a", 1);
400        m.insert("a", 2);
401        assert_eq!(m.get("a"), Some(&2));
402    }
403
404    #[test]
405    fn get_mut_returns_last_inserted() {
406        let mut m = VecMap::new();
407        m.insert("a", 1);
408        m.insert("a", 2);
409        *m.get_mut("a").unwrap() = 42;
410        assert_eq!(m.get("a"), Some(&42));
411        // First entry unchanged
412        assert_eq!(m.iter().next().unwrap().1, 1);
413    }
414
415    #[test]
416    fn remove_removes_all_occurrences() {
417        let mut m = VecMap::new();
418        m.insert("a", 1);
419        m.insert("b", 2);
420        m.insert("a", 3);
421        m.remove_slow("a");
422        assert_eq!(m.get("a"), None);
423        assert!(!m.contains_key("a"));
424        assert_eq!(m.len(), 1);
425    }
426
427    #[test]
428    fn contains_key_works() {
429        let mut m = VecMap::new();
430        assert!(!m.contains_key("x"));
431        m.insert("x", 10);
432        assert!(m.contains_key("x"));
433    }
434
435    #[test]
436    fn from_iterator() {
437        let m: VecMap<&str, i32> = vec![("a", 1), ("b", 2)].into_iter().collect();
438        assert_eq!(m.len(), 2);
439        assert_eq!(m.get("b"), Some(&2));
440    }
441
442    #[test]
443    fn into_iter_consuming() {
444        let mut m = VecMap::new();
445        m.insert("a", 1);
446        m.insert("b", 2);
447        let pairs: Vec<_> = m.into_iter().collect();
448        assert_eq!(pairs, vec![("a", 1), ("b", 2)]);
449    }
450
451    #[test]
452    fn is_deduped_false_initially() {
453        let m: VecMap<&str, i32> = VecMap::new();
454        assert!(!m.is_deduped());
455    }
456
457    #[test]
458    fn is_deduped_false_after_from() {
459        let m: VecMap<&str, i32> = vec![("a", 1)].into();
460        assert!(!m.is_deduped());
461    }
462
463    #[test]
464    fn is_deduped_false_after_collect() {
465        let m: VecMap<&str, i32> = vec![("a", 1)].into_iter().collect();
466        assert!(!m.is_deduped());
467    }
468
469    #[test]
470    fn dedup_sets_flag() {
471        let mut m = VecMap::new();
472        m.insert("a", 1);
473        assert!(!m.is_deduped());
474        m.dedup();
475        assert!(m.is_deduped());
476    }
477
478    #[test]
479    fn dedup_on_empty_map() {
480        let mut m: VecMap<String, i32> = VecMap::new();
481        m.dedup();
482        assert!(m.is_deduped());
483        assert!(m.is_empty());
484    }
485
486    #[test]
487    fn dedup_no_duplicates() {
488        let mut m = VecMap::new();
489        m.insert("a", 1);
490        m.insert("b", 2);
491        m.insert("c", 3);
492        m.dedup();
493        assert_eq!(m.len(), 3);
494        assert_eq!(m.get("a"), Some(&1));
495        assert_eq!(m.get("b"), Some(&2));
496        assert_eq!(m.get("c"), Some(&3));
497    }
498
499    #[test]
500    fn dedup_keeps_last_value() {
501        let mut m = VecMap::new();
502        m.insert("a", 1);
503        m.insert("b", 10);
504        m.insert("a", 2);
505        m.insert("a", 3);
506        m.insert("b", 20);
507        m.dedup();
508        assert_eq!(m.len(), 2);
509        assert_eq!(m.get("a"), Some(&3));
510        assert_eq!(m.get("b"), Some(&20));
511    }
512
513    #[test]
514    fn dedup_is_idempotent() {
515        let mut m = VecMap::new();
516        m.insert("a", 1);
517        m.insert("a", 2);
518        m.dedup();
519        assert!(m.is_deduped());
520        assert_eq!(m.len(), 1);
521        m.dedup();
522        assert!(m.is_deduped());
523        assert_eq!(m.len(), 1);
524        assert_eq!(m.get("a"), Some(&2));
525    }
526
527    #[test]
528    fn insert_dirties_dedup_flag() {
529        let mut m = VecMap::new();
530        m.insert("a", 1);
531        m.dedup();
532        assert!(m.is_deduped());
533
534        m.insert("b", 2);
535        assert!(!m.is_deduped());
536    }
537
538    #[test]
539    fn extend_dirties_dedup_flag() {
540        let mut m = VecMap::new();
541        m.insert("a", 1);
542        m.dedup();
543        assert!(m.is_deduped());
544
545        m.extend(vec![("b", 2)]);
546        assert!(!m.is_deduped());
547    }
548
549    #[test]
550    fn iter_mut_dirties_dedup_flag() {
551        let mut m = VecMap::new();
552        m.insert("a", 1);
553        m.dedup();
554        assert!(m.is_deduped());
555
556        for (_, v) in m.iter_mut() {
557            *v += 1;
558        }
559
560        assert!(!m.is_deduped());
561    }
562
563    #[test]
564    fn deduped_vec_map_borrowed_iter_and_len() {
565        let mut m = VecMap::new();
566        m.insert("a", 1);
567        m.insert("b", 2);
568        m.dedup();
569
570        let d = m.defensive_dedup();
571        assert!(matches!(d, DedupedVecMap::Borrowed(_)));
572        assert_eq!(d.len(), 2);
573
574        let mut items: Vec<_> = d.iter().collect();
575        items.sort_by_key(|(k, _)| **k);
576        assert_eq!(items, vec![(&"a", &1), (&"b", &2)]);
577    }
578
579    #[test]
580    fn deduped_vec_map_copy_iter_and_len() {
581        let mut m = VecMap::new();
582        m.insert("a", 1);
583        m.insert("a", 2);
584        m.insert("b", 3);
585
586        let d = m.defensive_dedup();
587        assert!(matches!(d, DedupedVecMap::Owned(_)));
588        assert_eq!(d.len(), 2);
589
590        let items: HashMap<&&str, &i32> = d.iter().collect();
591        assert_eq!(items[&"a"], &2);
592        assert_eq!(items[&"b"], &3);
593    }
594
595    #[test]
596    fn deduped_vec_map_iter_exact_size() {
597        let mut m = VecMap::new();
598        m.insert("a", 1);
599        m.insert("b", 2);
600        m.insert("c", 3);
601        m.dedup();
602
603        let d = m.defensive_dedup();
604        let mut iter = d.iter();
605        assert_eq!(iter.len(), 3);
606        iter.next();
607        assert_eq!(iter.len(), 2);
608    }
609
610    #[test]
611    fn deduped_vec_map_empty() {
612        let m: VecMap<String, i32> = VecMap::new();
613        let d = DedupedVecMap::Borrowed(&m);
614        assert_eq!(d.len(), 0);
615        assert_eq!(d.iter().count(), 0);
616    }
617
618    #[test]
619    fn dedup_does_not_require_clone() {
620        #[derive(Debug, PartialEq, Eq, Hash)]
621        struct NonCloneKey(u32);
622
623        let mut m = VecMap::new();
624        m.insert(NonCloneKey(1), "a");
625        m.insert(NonCloneKey(2), "b");
626        m.insert(NonCloneKey(1), "c");
627        m.dedup();
628        assert_eq!(m.len(), 2);
629        assert_eq!(m.get(&NonCloneKey(1)), Some(&"c"));
630        assert_eq!(m.get(&NonCloneKey(2)), Some(&"b"));
631    }
632
633    #[test]
634    fn serialize_deduplicates_keeping_last() {
635        let mut m = VecMap::new();
636        m.insert("a", 0);
637        m.insert("b", 0);
638        m.insert("b", 1);
639        m.insert("a", 1);
640        m.insert("a", 3);
641        m.insert("b", 2);
642
643        let serialized: serde_json::Value = serde_json::to_value(&m).unwrap();
644        let obj = serialized.as_object().unwrap();
645
646        assert_eq!(obj.len(), 2);
647        assert_eq!(obj.get("a").unwrap(), 3);
648        assert_eq!(obj.get("b").unwrap(), 2);
649    }
650}
libdd_trace_utils/span/vec_map.rs

libdd_trace_utils/span/
vec_map.rs
