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::DeserializableTraceData;
6use rmp::{decode, decode::RmpRead, Marker};
7use std::collections::HashMap;
8
9/// Reads a map from the buffer and returns it as a `HashMap`.
10///
11/// This function is generic over the key and value types of the map, and it uses a provided
12/// function to read key-value pairs from the buffer.
13///
14/// # Arguments
15///
16/// * `len` - The number of key-value pairs to read from the buffer.
17/// * `buf` - A reference to the slice containing the encoded map data.
18/// * `read_pair` - A function that reads a key-value pair from the buffer and returns it as a
19///   `Result<(K, V), DecodeError>`.
20///
21/// # Returns
22///
23/// * `Ok(HashMap<K, V>)` - A `HashMap` containing the decoded key-value pairs if successful.
24/// * `Err(DecodeError)` - An error if the decoding process fails.
25///
26/// # Errors
27///
28/// This function will return an error if:
29/// - The `read_pair` function returns an error while reading a key-value pair.
30///
31/// # Type Parameters
32///
33/// * `K` - The type of the keys in the map. Must implement `std::hash::Hash` and `Eq`.
34/// * `V` - The type of the values in the map.
35/// * `F` - The type of the function used to read key-value pairs from the buffer.
36#[inline]
37pub fn read_map<K, V, F, B>(
38    len: usize,
39    buf: &mut B,
40    read_pair: F,
41) -> Result<HashMap<K, V>, DecodeError>
42where
43    K: std::hash::Hash + Eq,
44    F: Fn(&mut B) -> Result<(K, V), DecodeError>,
45{
46    let mut map = HashMap::with_capacity(len);
47    for _ in 0..len {
48        let (k, v) = read_pair(buf)?;
49        map.insert(k, v);
50    }
51    Ok(map)
52}
53
54/// Reads map length from the buffer
55///
56/// # Arguments
57///
58/// * `buf` - A reference to the Bytes containing the encoded map data.
59///
60/// # Returns
61///
62/// * `Ok(usize)` - Map length.
63/// * `Err(DecodeError)` - An error if the decoding process fails.
64///
65/// # Errors
66///
67/// This function will return an error if:
68/// - The buffer does not contain a map.
69/// - There is an error reading from the buffer.
70#[inline]
71pub fn read_map_len<T: DeserializableTraceData>(buf: &mut Buffer<T>) -> Result<usize, DecodeError> {
72    let buf = buf.as_mut_slice();
73    match decode::read_marker(buf)
74        .map_err(|_| DecodeError::InvalidFormat("Unable to read marker for map".to_owned()))?
75    {
76        Marker::FixMap(len) => Ok(len as usize),
77        Marker::Map16 => buf
78            .read_data_u16()
79            .map_err(|_| DecodeError::IOError)
80            .map(|len| len as usize),
81        Marker::Map32 => buf
82            .read_data_u32()
83            .map_err(|_| DecodeError::IOError)
84            .map(|len| len as usize),
85        _ => Err(DecodeError::InvalidType(
86            "Unable to read map from buffer".to_owned(),
87        )),
88    }
89}
libdd_trace_utils/msgpack_decoder/decode/map.rs

libdd_trace_utils/msgpack_decoder/decode/
map.rs
