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

libdd_trace_utils/msgpack_decoder/decode/
map.rs
