libdd_trace_utils/msgpack_decoder/decode/
map.rs

1// Copyright 2024-Present Datadog, Inc. https://www.datadoghq.com/
2// SPDX-License-Identifier: Apache-2.0
3
4use crate::msgpack_decoder::decode::{buffer::Buffer, error::DecodeError};
5use crate::span::vec_map::VecMap;
6use crate::span::DeserializableTraceData;
7use rmp::{decode, decode::RmpRead, Marker};
8use std::collections::HashMap;
9
10/// Reads a map from the buffer and returns it as a `HashMap`.
11///
12/// This function is generic over the key and value types of the map, and it uses a provided
13/// function to read key-value pairs from the buffer.
14///
15/// # Arguments
16///
17/// * `len` - The number of key-value pairs to read from the buffer.
18/// * `buf` - A reference to the slice containing the encoded map data.
19/// * `read_pair` - A function that reads a key-value pair from the buffer and returns it as a
20///   `Result<(K, V), DecodeError>`.
21///
22/// # Returns
23///
24/// * `Ok(HashMap<K, V>)` - A `HashMap` containing the decoded key-value pairs if successful.
25/// * `Err(DecodeError)` - An error if the decoding process fails.
26///
27/// # Errors
28///
29/// This function will return an error if:
30/// - The `read_pair` function returns an error while reading a key-value pair.
31///
32/// # Type Parameters
33///
34/// * `K` - The type of the keys in the map. Must implement `std::hash::Hash` and `Eq`.
35/// * `V` - The type of the values in the map.
36/// * `F` - The type of the function used to read key-value pairs from the buffer.
37#[inline]
38pub fn read_map<K, V, F, B>(
39    len: usize,
40    buf: &mut B,
41    read_pair: F,
42) -> Result<HashMap<K, V>, DecodeError>
43where
44    K: std::hash::Hash + Eq,
45    F: Fn(&mut B) -> Result<(K, V), DecodeError>,
46{
47    let mut map = HashMap::with_capacity(len);
48    for _ in 0..len {
49        let (k, v) = read_pair(buf)?;
50        map.insert(k, v);
51    }
52    Ok(map)
53}
54
55/// Reads a map from the buffer and returns it as a `VecMap<K, V>`.
56///
57/// Like `read_map` but returns pairs in a VecMap instead of HashMap for better
58/// cache locality when iteration order doesn't matter.
59#[inline]
60pub fn read_map_vec<K, V, F, B>(
61    len: usize,
62    buf: &mut B,
63    read_pair: F,
64) -> Result<VecMap<K, V>, DecodeError>
65where
66    F: Fn(&mut B) -> Result<(K, V), DecodeError>,
67{
68    let mut map = VecMap::with_capacity(len);
69    for _ in 0..len {
70        let (k, v) = read_pair(buf)?;
71        map.insert(k, v);
72    }
73    Ok(map)
74}
75
76/// Reads map length from the buffer
77///
78/// # Arguments
79///
80/// * `buf` - A reference to the Bytes containing the encoded map data.
81///
82/// # Returns
83///
84/// * `Ok(usize)` - Map length.
85/// * `Err(DecodeError)` - An error if the decoding process fails.
86///
87/// # Errors
88///
89/// This function will return an error if:
90/// - The buffer does not contain a map.
91/// - There is an error reading from the buffer.
92#[inline]
93pub fn read_map_len<T: DeserializableTraceData>(buf: &mut Buffer<T>) -> Result<usize, DecodeError> {
94    let buf = buf.as_mut_slice();
95    match decode::read_marker(buf)
96        .map_err(|_| DecodeError::InvalidFormat("Unable to read marker for map".to_owned()))?
97    {
98        Marker::FixMap(len) => Ok(len as usize),
99        Marker::Map16 => buf
100            .read_data_u16()
101            .map_err(|_| DecodeError::IOError)
102            .map(|len| len as usize),
103        Marker::Map32 => buf
104            .read_data_u32()
105            .map_err(|_| DecodeError::IOError)
106            .map(|len| len as usize),
107        _ => Err(DecodeError::InvalidType(
108            "Unable to read map from buffer".to_owned(),
109        )),
110    }
111}
libdd_trace_utils/msgpack_decoder/decode/map.rs

libdd_trace_utils/msgpack_decoder/decode/
map.rs
