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
219impl<K: Hash + Eq + Clone, V> VecMap<K, V> {
220    /// Remove entries with a duplicate key, only keeping the last one. After this, a flag is set
221    /// internally, such that as long as the map isn't extended or mutably iterated, the next
222    /// [Self::dedup] doesn't perform the work again.
223    pub fn dedup(&mut self) {
224        if self.deduped {
225            return;
226        }
227
228        // Since we're going to shuffle elements around, it's not easy to keep references to keys in
229        // the deduping set. The simplest is to clone them.
230        let mut seen = HashSet::with_capacity(self.len());
231
232        self.data.reverse();
233        self.data.retain(|(k, _)| seen.insert(k.clone()));
234        self.deduped = true;
235    }
236}
237
238impl<K, V> From<Vec<(K, V)>> for VecMap<K, V> {
239    fn from(data: Vec<(K, V)>) -> Self {
240        Self {
241            data,
242            deduped: false,
243        }
244    }
245}
246
247impl<K, V> From<VecMap<K, V>> for Vec<(K, V)> {
248    fn from(value: VecMap<K, V>) -> Self {
249        value.data
250    }
251}
252
253impl<K, V> FromIterator<(K, V)> for VecMap<K, V> {
254    fn from_iter<I: IntoIterator<Item = (K, V)>>(iter: I) -> Self {
255        Self {
256            data: iter.into_iter().collect(),
257            deduped: false,
258        }
259    }
260}
261
262impl<K, V> IntoIterator for VecMap<K, V> {
263    type Item = (K, V);
264    type IntoIter = std::vec::IntoIter<(K, V)>;
265
266    fn into_iter(self) -> Self::IntoIter {
267        self.data.into_iter()
268    }
269}
270
271impl<'a, K, V> IntoIterator for &'a VecMap<K, V> {
272    type Item = &'a (K, V);
273    type IntoIter = slice::Iter<'a, (K, V)>;
274
275    fn into_iter(self) -> Self::IntoIter {
276        self.data.iter()
277    }
278}
279
280impl<'a, K, V> IntoIterator for &'a mut VecMap<K, V> {
281    type Item = &'a mut (K, V);
282    type IntoIter = slice::IterMut<'a, (K, V)>;
283
284    fn into_iter(self) -> Self::IntoIter {
285        // Since we iterate on keys as well, they can modified, and introduce duplicates.
286        self.dirty();
287        self.data.iter_mut()
288    }
289}
290
291impl<K, V> Extend<(K, V)> for VecMap<K, V> {
292    fn extend<I: IntoIterator<Item = (K, V)>>(&mut self, iter: I) {
293        self.dirty();
294        self.data.extend(iter);
295    }
296}
297
298impl<K: Serialize + Eq + Hash, V: Serialize> Serialize for VecMap<K, V> {
299    fn serialize<S: Serializer>(&self, serializer: S) -> Result<S::Ok, S::Error> {
300        use serde::ser::SerializeMap;
301
302        let deduped = self.as_deduped_map();
303        let mut map_ser = serializer.serialize_map(Some(deduped.len()))?;
304
305        for (k, v) in deduped.iter() {
306            map_ser.serialize_entry(k, v)?;
307        }
308
309        map_ser.end()
310    }
311}
312
313pub enum DedupedVecMap<'a, K, V> {
314    Borrowed(&'a VecMap<K, V>),
315    Owned(HashMap<&'a K, &'a V>),
316}
317
318impl<'a, K, V> DedupedVecMap<'a, K, V> {
319    #[inline]
320    pub fn iter(&self) -> DedupedVecMapIter<'_, 'a, K, V> {
321        match self {
322            DedupedVecMap::Borrowed(vec_map) => DedupedVecMapIter::Borrowed(vec_map.iter()),
323            DedupedVecMap::Owned(map) => DedupedVecMapIter::Owned(map.iter()),
324        }
325    }
326
327    #[inline]
328    pub fn len(&self) -> usize {
329        match self {
330            DedupedVecMap::Borrowed(vec_map) => vec_map.len(),
331            DedupedVecMap::Owned(map) => map.len(),
332        }
333    }
334
335    #[inline]
336    pub fn is_empty(&self) -> bool {
337        match self {
338            DedupedVecMap::Borrowed(vec_map) => vec_map.is_empty(),
339            DedupedVecMap::Owned(map) => map.is_empty(),
340        }
341    }
342}
343
344pub enum DedupedVecMapIter<'b, 'a: 'b, K, V> {
345    Borrowed(slice::Iter<'a, (K, V)>),
346    Owned(std::collections::hash_map::Iter<'b, &'a K, &'a V>),
347}
348
349impl<'b, 'a: 'b, K, V> Iterator for DedupedVecMapIter<'b, 'a, K, V> {
350    type Item = (&'a K, &'a V);
351
352    fn next(&mut self) -> Option<Self::Item> {
353        match self {
354            DedupedVecMapIter::Borrowed(iter) => iter.next().map(|(k, v)| (k, v)),
355            DedupedVecMapIter::Owned(iter) => iter.next().map(|(&k, &v)| (k, v)),
356        }
357    }
358
359    fn size_hint(&self) -> (usize, Option<usize>) {
360        match self {
361            DedupedVecMapIter::Borrowed(iter) => iter.size_hint(),
362            DedupedVecMapIter::Owned(iter) => iter.size_hint(),
363        }
364    }
365}
366
367impl<'b, 'a: 'b, K, V> ExactSizeIterator for DedupedVecMapIter<'b, 'a, K, V> {}
368
369#[cfg(test)]
370mod tests {
371    use super::*;
372
373    #[test]
374    fn get_returns_last_inserted() {
375        let mut m = VecMap::new();
376        m.insert("a", 1);
377        m.insert("a", 2);
378        assert_eq!(m.get("a"), Some(&2));
379    }
380
381    #[test]
382    fn get_mut_returns_last_inserted() {
383        let mut m = VecMap::new();
384        m.insert("a", 1);
385        m.insert("a", 2);
386        *m.get_mut("a").unwrap() = 42;
387        assert_eq!(m.get("a"), Some(&42));
388        // First entry unchanged
389        assert_eq!(m.iter().next().unwrap().1, 1);
390    }
391
392    #[test]
393    fn remove_removes_all_occurrences() {
394        let mut m = VecMap::new();
395        m.insert("a", 1);
396        m.insert("b", 2);
397        m.insert("a", 3);
398        m.remove_slow("a");
399        assert_eq!(m.get("a"), None);
400        assert!(!m.contains_key("a"));
401        assert_eq!(m.len(), 1);
402    }
403
404    #[test]
405    fn contains_key_works() {
406        let mut m = VecMap::new();
407        assert!(!m.contains_key("x"));
408        m.insert("x", 10);
409        assert!(m.contains_key("x"));
410    }
411
412    #[test]
413    fn from_iterator() {
414        let m: VecMap<&str, i32> = vec![("a", 1), ("b", 2)].into_iter().collect();
415        assert_eq!(m.len(), 2);
416        assert_eq!(m.get("b"), Some(&2));
417    }
418
419    #[test]
420    fn into_iter_consuming() {
421        let mut m = VecMap::new();
422        m.insert("a", 1);
423        m.insert("b", 2);
424        let pairs: Vec<_> = m.into_iter().collect();
425        assert_eq!(pairs, vec![("a", 1), ("b", 2)]);
426    }
427
428    #[test]
429    fn is_deduped_false_initially() {
430        let m: VecMap<&str, i32> = VecMap::new();
431        assert!(!m.is_deduped());
432    }
433
434    #[test]
435    fn is_deduped_false_after_from() {
436        let m: VecMap<&str, i32> = vec![("a", 1)].into();
437        assert!(!m.is_deduped());
438    }
439
440    #[test]
441    fn is_deduped_false_after_collect() {
442        let m: VecMap<&str, i32> = vec![("a", 1)].into_iter().collect();
443        assert!(!m.is_deduped());
444    }
445
446    #[test]
447    fn dedup_sets_flag() {
448        let mut m = VecMap::new();
449        m.insert("a", 1);
450        assert!(!m.is_deduped());
451        m.dedup();
452        assert!(m.is_deduped());
453    }
454
455    #[test]
456    fn dedup_on_empty_map() {
457        let mut m: VecMap<String, i32> = VecMap::new();
458        m.dedup();
459        assert!(m.is_deduped());
460        assert!(m.is_empty());
461    }
462
463    #[test]
464    fn dedup_no_duplicates() {
465        let mut m = VecMap::new();
466        m.insert("a", 1);
467        m.insert("b", 2);
468        m.insert("c", 3);
469        m.dedup();
470        assert_eq!(m.len(), 3);
471        assert_eq!(m.get("a"), Some(&1));
472        assert_eq!(m.get("b"), Some(&2));
473        assert_eq!(m.get("c"), Some(&3));
474    }
475
476    #[test]
477    fn dedup_keeps_last_value() {
478        let mut m = VecMap::new();
479        m.insert("a", 1);
480        m.insert("b", 10);
481        m.insert("a", 2);
482        m.insert("a", 3);
483        m.insert("b", 20);
484        m.dedup();
485        assert_eq!(m.len(), 2);
486        assert_eq!(m.get("a"), Some(&3));
487        assert_eq!(m.get("b"), Some(&20));
488    }
489
490    #[test]
491    fn dedup_is_idempotent() {
492        let mut m = VecMap::new();
493        m.insert("a", 1);
494        m.insert("a", 2);
495        m.dedup();
496        assert!(m.is_deduped());
497        assert_eq!(m.len(), 1);
498        m.dedup();
499        assert!(m.is_deduped());
500        assert_eq!(m.len(), 1);
501        assert_eq!(m.get("a"), Some(&2));
502    }
503
504    #[test]
505    fn insert_dirties_dedup_flag() {
506        let mut m = VecMap::new();
507        m.insert("a", 1);
508        m.dedup();
509        assert!(m.is_deduped());
510
511        m.insert("b", 2);
512        assert!(!m.is_deduped());
513    }
514
515    #[test]
516    fn extend_dirties_dedup_flag() {
517        let mut m = VecMap::new();
518        m.insert("a", 1);
519        m.dedup();
520        assert!(m.is_deduped());
521
522        m.extend(vec![("b", 2)]);
523        assert!(!m.is_deduped());
524    }
525
526    #[test]
527    fn iter_mut_dirties_dedup_flag() {
528        let mut m = VecMap::new();
529        m.insert("a", 1);
530        m.dedup();
531        assert!(m.is_deduped());
532
533        for (_, v) in m.iter_mut() {
534            *v += 1;
535        }
536
537        assert!(!m.is_deduped());
538    }
539
540    #[test]
541    fn deduped_vec_map_borrowed_iter_and_len() {
542        let mut m = VecMap::new();
543        m.insert("a", 1);
544        m.insert("b", 2);
545        m.dedup();
546
547        let d = m.defensive_dedup();
548        assert!(matches!(d, DedupedVecMap::Borrowed(_)));
549        assert_eq!(d.len(), 2);
550
551        let mut items: Vec<_> = d.iter().collect();
552        items.sort_by_key(|(k, _)| **k);
553        assert_eq!(items, vec![(&"a", &1), (&"b", &2)]);
554    }
555
556    #[test]
557    fn deduped_vec_map_copy_iter_and_len() {
558        let mut m = VecMap::new();
559        m.insert("a", 1);
560        m.insert("a", 2);
561        m.insert("b", 3);
562
563        let d = m.defensive_dedup();
564        assert!(matches!(d, DedupedVecMap::Owned(_)));
565        assert_eq!(d.len(), 2);
566
567        let items: HashMap<&&str, &i32> = d.iter().collect();
568        assert_eq!(items[&"a"], &2);
569        assert_eq!(items[&"b"], &3);
570    }
571
572    #[test]
573    fn deduped_vec_map_iter_exact_size() {
574        let mut m = VecMap::new();
575        m.insert("a", 1);
576        m.insert("b", 2);
577        m.insert("c", 3);
578        m.dedup();
579
580        let d = m.defensive_dedup();
581        let mut iter = d.iter();
582        assert_eq!(iter.len(), 3);
583        iter.next();
584        assert_eq!(iter.len(), 2);
585    }
586
587    #[test]
588    fn deduped_vec_map_empty() {
589        let m: VecMap<String, i32> = VecMap::new();
590        let d = DedupedVecMap::Borrowed(&m);
591        assert_eq!(d.len(), 0);
592        assert_eq!(d.iter().count(), 0);
593    }
594
595    #[test]
596    fn serialize_deduplicates_keeping_last() {
597        let mut m = VecMap::new();
598        m.insert("a", 0);
599        m.insert("b", 0);
600        m.insert("b", 1);
601        m.insert("a", 1);
602        m.insert("a", 3);
603        m.insert("b", 2);
604
605        let serialized: serde_json::Value = serde_json::to_value(&m).unwrap();
606        let obj = serialized.as_object().unwrap();
607
608        assert_eq!(obj.len(), 2);
609        assert_eq!(obj.get("a").unwrap(), 3);
610        assert_eq!(obj.get("b").unwrap(), 2);
611    }
612}
libdd_trace_utils/span/vec_map.rs

libdd_trace_utils/span/
vec_map.rs
