[WIP] Rearrange folder structure.

Author: hoolioh

trace-utils/src/lib.rs

@@ -12,4 +12,4 @@ pub mod trace_utils;
 pub mod tracer_header_tags;
 pub mod tracer_payload;
 
-pub mod span_v04;
+pub mod span;

trace-utils/src/msgpack_decoder/v04/error.rs renamed to trace-utils/src/msgpack_decoder/decode/error.rs

@@ -0,0 +1,121 @@
+// Copyright 2024-Present Datadog, Inc. https://www.datadoghq.com/
+// SPDX-License-Identifier: Apache-2.0
+
+pub mod error;
+pub mod number;
+pub mod span_link;
+pub mod string;
+use crate::msgpack_decoder::decode::error::DecodeError;
+use crate::msgpack_decoder::decode::number::read_number_bytes;
+use crate::msgpack_decoder::decode::string::{handle_null_marker, read_string_bytes};
+use rmp::{decode, decode::RmpRead, Marker};
+use std::{collections::HashMap, f64};
+use tinybytes::{Bytes, BytesString};
+
+#[inline]
+pub fn read_metric_pair(buf: &mut Bytes) -> Result<(BytesString, f64), DecodeError> {
+    let key = read_string_bytes(buf)?;
+    let v = read_number_bytes(buf)?;
+
+    Ok((key, v))
+}
+#[inline]
+pub fn read_metrics(buf: &mut Bytes) -> Result<HashMap<BytesString, f64>, DecodeError> {
+    if let Some(empty_map) = handle_null_marker(buf, HashMap::default) {
+        return Ok(empty_map);
+    }
+
+    let len = read_map_len(unsafe { buf.as_mut_slice() })?;
+
+    read_map(len, buf, read_metric_pair)
+}
+
+#[inline]
+pub fn read_meta_struct(buf: &mut Bytes) -> Result<HashMap<BytesString, Vec<u8>>, DecodeError> {
+    if let Some(empty_map) = handle_null_marker(buf, HashMap::default) {
+        return Ok(empty_map);
+    }
+
+    fn read_meta_struct_pair(buf: &mut Bytes) -> Result<(BytesString, Vec<u8>), DecodeError> {
+        let key = read_string_bytes(buf)?;
+        let array_len = decode::read_array_len(unsafe { buf.as_mut_slice() }).map_err(|_| {
+            DecodeError::InvalidFormat("Unable to read array len for meta_struct".to_owned())
+        })?;
+
+        let mut v = Vec::with_capacity(array_len as usize);
+
+        for _ in 0..array_len {
+            let value = read_number_bytes(buf)?;
+            v.push(value);
+        }
+        Ok((key, v))
+    }
+
+    let len = read_map_len(unsafe { buf.as_mut_slice() })?;
+    read_map(len, buf, read_meta_struct_pair)
+}
+
+/// Reads a map from the buffer and returns it as a `HashMap`.
+///
+/// This function is generic over the key and value types of the map, and it uses a provided
+/// function to read key-value pairs from the buffer.
+///
+/// # Arguments
+///
+/// * `len` - The number of key-value pairs to read from the buffer.
+/// * `buf` - A reference to the Bytes containing the encoded map data.
+/// * `read_pair` - A function that reads a key-value pair from the buffer and returns it as a
+///   `Result<(K, V), DecodeError>`.
+///
+/// # Returns
+///
+/// * `Ok(HashMap<K, V>)` - A `HashMap` containing the decoded key-value pairs if successful.
+/// * `Err(DecodeError)` - An error if the decoding process fails.
+///
+/// # Errors
+///
+/// This function will return an error if:
+/// - The `read_pair` function returns an error while reading a key-value pair.
+///
+/// # Type Parameters
+///
+/// * `K` - The type of the keys in the map. Must implement `std::hash::Hash` and `Eq`.
+/// * `V` - The type of the values in the map.
+/// * `F` - The type of the function used to read key-value pairs from the buffer.
+#[inline]
+pub fn read_map<K, V, F>(
+    len: usize,
+    buf: &mut Bytes,
+    read_pair: F,
+) -> Result<HashMap<K, V>, DecodeError>
+where
+    K: std::hash::Hash + Eq,
+    F: Fn(&mut Bytes) -> Result<(K, V), DecodeError>,
+{
+    let mut map = HashMap::with_capacity(len);
+    for _ in 0..len {
+        let (k, v) = read_pair(buf)?;
+        map.insert(k, v);
+    }
+    Ok(map)
+}
+
+#[inline]
+pub fn read_map_len(buf: &mut &[u8]) -> Result<usize, DecodeError> {
+    match decode::read_marker(buf)
+        .map_err(|_| DecodeError::InvalidFormat("Unable to read marker for map".to_owned()))?
+    {
+        Marker::FixMap(len) => Ok(len as usize),
+        Marker::Map16 => buf
+            .read_data_u16()
+            .map_err(|_| DecodeError::IOError)
+            .map(|len| len as usize),
+        Marker::Map32 => buf
+            .read_data_u32()
+            .map_err(|_| DecodeError::IOError)
+            .map(|len| len as usize),
+        _ => Err(DecodeError::InvalidType(
+            "Unable to read map from buffer".to_owned(),
+        )),
+    }
+}

trace-utils/src/msgpack_decoder/decode/mod.rs

@@ -1,7 +1,7 @@
 // Copyright 2024-Present Datadog, Inc. https://www.datadoghq.com/
 // SPDX-License-Identifier: Apache-2.0
 
-use super::error::DecodeError;
+use crate::msgpack_decoder::decode::error::DecodeError;
 use rmp::{decode::RmpRead, Marker};
 use std::fmt;
 use tinybytes::Bytes;

trace-utils/src/msgpack_decoder/v04/number.rs renamed to trace-utils/src/msgpack_decoder/decode/number.rs

@@ -1,12 +1,12 @@
 // Copyright 2024-Present Datadog, Inc. https://www.datadoghq.com/
 // SPDX-License-Identifier: Apache-2.0
 
-use crate::msgpack_decoder::v04::decoder::{
+use crate::msgpack_decoder::decode::error::DecodeError;
+use crate::msgpack_decoder::decode::number::read_number_bytes;
+use crate::msgpack_decoder::decode::string::{
     handle_null_marker, read_str_map_to_bytes_strings, read_string_bytes, read_string_ref,
 };
-use crate::msgpack_decoder::v04::error::DecodeError;
-use crate::msgpack_decoder::v04::number::read_number_bytes;
-use crate::span_v04::SpanLink;
+use crate::span::v04::SpanLink;
 use rmp::Marker;
 use std::str::FromStr;
 use tinybytes::Bytes;
@@ -98,7 +98,7 @@ fn decode_span_link(buf: &mut Bytes) -> Result<SpanLink, DecodeError> {
 #[cfg(test)]
 mod tests {
     use super::SpanLinkKey;
-    use crate::msgpack_decoder::v04::error::DecodeError;
+    use crate::msgpack_decoder::decode::error::DecodeError;
     use std::str::FromStr;
 
     #[test]

trace-utils/src/msgpack_decoder/v04/decoder/span_link.rs renamed to trace-utils/src/msgpack_decoder/decode/span_link.rs

@@ -0,0 +1,97 @@
+// Copyright 2024-Present Datadog, Inc. https://www.datadoghq.com/
+// SPDX-License-Identifier: Apache-2.0
+
+use crate::msgpack_decoder::decode::error::DecodeError;
+use rmp::decode;
+use rmp::decode::DecodeStringError;
+use std::collections::HashMap;
+use tinybytes::{Bytes, BytesString};
+
+// https://docs.rs/rmp/latest/rmp/enum.Marker.html#variant.Null (0xc0 == 192)
+const NULL_MARKER: &u8 = &0xc0;
+
+#[inline]
+pub fn read_string_ref_nomut(buf: &[u8]) -> Result<(&str, &[u8]), DecodeError> {
+    decode::read_str_from_slice(buf).map_err(|e| match e {
+        DecodeStringError::InvalidMarkerRead(e) => DecodeError::InvalidFormat(e.to_string()),
+        DecodeStringError::InvalidDataRead(e) => DecodeError::InvalidConversion(e.to_string()),
+        DecodeStringError::TypeMismatch(marker) => {
+            DecodeError::InvalidType(format!("Type mismatch at marker {:?}", marker))
+        }
+        DecodeStringError::InvalidUtf8(_, e) => DecodeError::Utf8Error(e.to_string()),
+        _ => DecodeError::IOError,
+    })
+}
+
+#[inline]
+pub fn read_string_ref<'a>(buf: &mut &'a [u8]) -> Result<&'a str, DecodeError> {
+    read_string_ref_nomut(buf).map(|(str, newbuf)| {
+        *buf = newbuf;
+        str
+    })
+}
+
+#[inline]
+pub fn read_string_bytes(buf: &mut Bytes) -> Result<BytesString, DecodeError> {
+    // Note: we need to pass a &'static lifetime here, otherwise it'll complain
+    read_string_ref_nomut(unsafe { buf.as_mut_slice() }).map(|(str, newbuf)| {
+        let string = BytesString::from_bytes_slice(buf, str);
+        *unsafe { buf.as_mut_slice() } = newbuf;
+        string
+    })
+}
+
+#[inline]
+pub fn read_nullable_string_bytes(buf: &mut Bytes) -> Result<BytesString, DecodeError> {
+    if let Some(empty_string) = handle_null_marker(buf, BytesString::default) {
+        Ok(empty_string)
+    } else {
+        read_string_bytes(buf)
+    }
+}
+
+#[inline]
+// Safety: read_string_ref checks utf8 validity, so we don't do it again when creating the
+// BytesStrings.
+pub fn read_str_map_to_bytes_strings(
+    buf: &mut Bytes,
+) -> Result<HashMap<BytesString, BytesString>, DecodeError> {
+    let len = decode::read_map_len(unsafe { buf.as_mut_slice() })
+        .map_err(|_| DecodeError::InvalidFormat("Unable to get map len for str map".to_owned()))?;
+
+    let mut map = HashMap::with_capacity(len.try_into().expect("Unable to cast map len to usize"));
+    for _ in 0..len {
+        let key = read_string_bytes(buf)?;
+        let value = read_string_bytes(buf)?;
+        map.insert(key, value);
+    }
+    Ok(map)
+}
+
+#[inline]
+pub fn read_nullable_str_map_to_bytes_strings(
+    buf: &mut Bytes,
+) -> Result<HashMap<BytesString, BytesString>, DecodeError> {
+    if let Some(empty_map) = handle_null_marker(buf, HashMap::default) {
+        return Ok(empty_map);
+    }
+
+    read_str_map_to_bytes_strings(buf)
+}
+
+/// When you want to "peek" if the next value is a null marker, and only advance the buffer if it is
+/// null and return the default value. If it is not null, you can continue to decode as expected.
+#[inline]
+pub fn handle_null_marker<T, F>(buf: &mut Bytes, default: F) -> Option<T>
+where
+    F: FnOnce() -> T,
+{
+    let slice = unsafe { buf.as_mut_slice() };
+
+    if slice.first() == Some(NULL_MARKER) {
+        *slice = &slice[1..];
+        Some(default())
+    } else {
+        None
+    }
+}

trace-utils/src/msgpack_decoder/decode/string.rs

@@ -1,4 +1,5 @@
 // Copyright 2024-Present Datadog, Inc. https://www.datadoghq.com/
 // SPDX-License-Identifier: Apache-2.0
 
+pub mod decode;
 pub mod v04;