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    map: indexmap::IndexSet<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(self.map.iter().map(|entry| -> &str { entry.borrow() }))
20    }
21}
22
23impl<T: SpanText> SharedDict<T> {
24    /// Gets the index of the interned string. If the string is not part of the dictionary it is
25    /// added and its corresponding index returned.
26    ///
27    /// # Arguments:
28    ///
29    /// * `str`: string to look up in the dictionary.
30    pub fn get_or_insert(&mut self, s: T) -> Result<u32, std::num::TryFromIntError> {
31        if let Some(index) = self.map.get_index_of(s.borrow()) {
32            (index).try_into()
33        } else {
34            let index = self.map.len();
35            self.map.insert(s);
36            index.try_into()
37        }
38    }
39
40    #[allow(clippy::len_without_is_empty)]
41    pub fn len(&self) -> usize {
42        self.map.len()
43    }
44
45    pub fn iter(&self) -> impl Iterator<Item = &T> {
46        self.map.iter()
47    }
48}
49
50impl<T: SpanText> Default for SharedDict<T> {
51    fn default() -> Self {
52        Self {
53            map: indexmap::indexset! {T::default()},
54        }
55    }
56}
57
58#[cfg(test)]
59mod tests {
60    use libdd_tinybytes::{Bytes, BytesString};
61
62    use super::*;
63
64    #[test]
65    fn default_test() {
66        let dict: SharedDict<BytesString> = SharedDict::default();
67
68        assert_eq!(dict.map.len(), 1);
69    }
70
71    #[test]
72    fn get_or_insert_test() {
73        let mut dict = SharedDict::default();
74        unsafe {
75            let _ = dict.get_or_insert(BytesString::from_bytes_unchecked(Bytes::from_static(
76                b"foo",
77            )));
78        };
79        unsafe {
80            let _ = dict.get_or_insert(BytesString::from_bytes_unchecked(Bytes::from_static(
81                b"bar",
82            )));
83        };
84
85        assert_eq!(dict.map.len(), 3);
86
87        assert_eq!(dict.map[0].as_str(), "");
88        assert_eq!(dict.map[1].as_str(), "foo");
89        assert_eq!(dict.map[2].as_str(), "bar");
90    }
91}