Skip to main content

libdd_trace_utils/span/v05/
dict.rs

1// Copyright 2021-Present Datadog, Inc. https://www.datadoghq.com/
2// SPDX-License-Identifier: Apache-2.0
3
4use crate::span::SpanText;
5
6/// This struct represents the shared dictionary used for interning all the strings belonging to a
7/// v05 trace chunk.
8#[derive(Debug, Clone)]
9pub struct SharedDict<T> {
10    /// Map strings with their index and keep insertion order(O(1) retrieval complexity).
11    pub(crate) map: indexmap::IndexMap<T, ()>,
12}
13
14impl<T: SpanText> serde::Serialize for SharedDict<T> {
15    fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
16    where
17        S: serde::Serializer,
18    {
19        serializer.collect_seq(
20            self.map
21                .iter()
22                .map(|(entry, ())| -> &str { entry.borrow() }),
23        )
24    }
25}
26
27impl<T: SpanText> SharedDict<T> {
28    /// Gets the index of the interned string. If the string is not part of the dictionary it is
29    /// added and its corresponding index returned.
30    ///
31    /// # Arguments:
32    ///
33    /// * `str`: string to look up in the dictionary.
34    pub fn get_or_insert(&mut self, s: T) -> Result<u32, std::num::TryFromIntError> {
35        if let Some(index) = self.map.get_index_of(s.borrow()) {
36            (index).try_into()
37        } else {
38            let index = self.map.len();
39            self.map.insert(s, ());
40            index.try_into()
41        }
42    }
43
44    #[allow(clippy::len_without_is_empty)]
45    pub fn len(&self) -> usize {
46        self.map.len()
47    }
48
49    pub fn iter(&self) -> impl Iterator<Item = &T> {
50        self.map.keys()
51    }
52}
53
54impl<T: SpanText> Default for SharedDict<T> {
55    fn default() -> Self {
56        Self {
57            map: indexmap::indexmap! {T::default() => ()},
58        }
59    }
60}
61
62#[cfg(test)]
63mod tests {
64    use libdd_tinybytes::{Bytes, BytesString};
65
66    use super::*;
67
68    #[test]
69    fn default_test() {
70        let dict: SharedDict<BytesString> = SharedDict::default();
71
72        assert_eq!(dict.map.len(), 1);
73    }
74
75    #[test]
76    fn get_or_insert_test() {
77        let mut dict = SharedDict::default();
78        unsafe {
79            let _ = dict.get_or_insert(BytesString::from_bytes_unchecked(Bytes::from_static(
80                b"foo",
81            )));
82        };
83        unsafe {
84            let _ = dict.get_or_insert(BytesString::from_bytes_unchecked(Bytes::from_static(
85                b"bar",
86            )));
87        };
88
89        assert_eq!(dict.map.len(), 3);
90
91        assert_eq!(dict.map.get_index(0).unwrap().0.as_str(), "");
92        assert_eq!(dict.map.get_index(1).unwrap().0.as_str(), "foo");
93        assert_eq!(dict.map.get_index(2).unwrap().0.as_str(), "bar");
94    }
95}