rustolio_utils/bytes/encoding/

mod.rs

//
// SPDX-License-Identifier: MPL-2.0
//
// Copyright (c) 2026 Tobias Binnewies. All rights reserved.
//
// This Source Code Form is subject to the terms of the Mozilla Public
// License, v. 2.0. If a copy of the MPL was not distributed with this
// file, You can obtain one at http://mozilla.org/MPL/2.0/.
//

mod decode;
mod encode;

use std::{any::type_name, mem};

use crate::{
    bytes::{Bytes, BytesMut},
    prelude::*,
};

use bytes::{Buf, BufMut};
pub use decode::{AsyncDecoder, Decode, Decoder};
pub use encode::{AsyncEncoder, Encode, Encoder};
pub use rustolio_utils_macros::{Decode, Encode, async_impl};

pub fn encode_to_bytes(v: &impl Encode) -> crate::Result<Bytes> {
    struct BytesWriter(BytesMut);

    impl Encoder for BytesWriter {
        fn write_u8(&mut self, byte: u8) -> crate::Result<()> {
            self.0.put_u8(byte);
            Ok(())
        }
        fn write_all(&mut self, bytes: &[u8]) -> crate::Result<()> {
            self.0.extend_from_slice(bytes);
            Ok(())
        }
    }

    let size = v.encode_size();
    let mut writer = BytesWriter(BytesMut::with_capacity(size));
    v.encode(&mut writer)?;
    let bytes = writer.0.freeze();

    #[cfg(test)]
    assert_eq!(bytes.len(), size);

    Ok(bytes)
}

pub fn decode_from_bytes<T: Decode>(bytes: Bytes) -> crate::Result<T> {
    struct BytesReader(Bytes);

    impl Decoder for BytesReader {
        fn read_u8(&mut self) -> crate::Result<u8> {
            Ok(self.0.get_u8())
        }
        fn read_exact(&mut self, buf: &mut [u8]) -> crate::Result<()> {
            self.0.try_copy_to_slice(buf).context("UnexpectedEof")
        }
    }

    let mut reader = BytesReader(bytes);
    let v: T = T::decode(&mut reader)?;
    Ok(v)
}

pub fn encode_to_std<W>(v: &impl Encode, writer: &mut W) -> crate::Result<()>
where
    W: std::io::Write + Unpin,
{
    struct StdWriter<W>(W);

    impl<W> Encoder for StdWriter<W>
    where
        W: std::io::Write + Unpin,
    {
        fn write_u8(&mut self, byte: u8) -> crate::Result<()> {
            std::io::Write::write_all(&mut self.0, &[byte]).context("Failed to write u8")
        }
        fn write_all(&mut self, bytes: &[u8]) -> crate::Result<()> {
            std::io::Write::write_all(&mut self.0, bytes).context("Failed to write all")
        }
    }

    let mut writer = StdWriter(writer);
    v.encode(&mut writer)?;
    Ok(())
}

pub fn decode_from_std<R, T: Decode>(reader: &mut R) -> crate::Result<T>
where
    R: std::io::Read + Unpin,
{
    struct StdReader<R>(R);

    impl<R> Decoder for StdReader<R>
    where
        R: std::io::Read + Unpin,
    {
        fn read_u8(&mut self) -> crate::Result<u8> {
            let mut buf = [0];
            std::io::Read::read_exact(&mut self.0, &mut buf).context("Failed to read u8")?;
            Ok(buf[0])
        }
        fn read_exact(&mut self, buf: &mut [u8]) -> crate::Result<()> {
            std::io::Read::read_exact(&mut self.0, buf).context("Failed to read exact")
        }
    }

    let mut reader = StdReader(reader);
    let v: T = T::decode(&mut reader)?;
    Ok(v)
}

#[cfg(not(target_arch = "wasm32"))]
pub async fn encode_to_tokio<W>(v: &impl Encode, writer: &mut W) -> crate::Result<()>
where
    W: tokio::io::AsyncWriteExt + Unpin,
{
    struct TokioWriter<W>(W);

    impl<W> AsyncEncoder for TokioWriter<W>
    where
        W: tokio::io::AsyncWriteExt + Unpin,
    {
        async fn write_u8(&mut self, byte: u8) -> crate::Result<()> {
            tokio::io::AsyncWriteExt::write_u8(&mut self.0, byte)
                .await
                .context("Failed to write u8")
        }
        async fn write_all(&mut self, bytes: &[u8]) -> crate::Result<()> {
            tokio::io::AsyncWriteExt::write_all(&mut self.0, bytes)
                .await
                .context("Failed to write all")
        }
    }

    let mut writer = TokioWriter(writer);
    v.encode_async(&mut writer).await?;
    Ok(())
}

#[cfg(not(target_arch = "wasm32"))]
pub async fn decode_from_tokio<R, T: Decode>(reader: &mut R) -> crate::Result<T>
where
    R: tokio::io::AsyncReadExt + Unpin,
{
    struct TokioReader<R>(R);

    impl<R> AsyncDecoder for TokioReader<R>
    where
        R: tokio::io::AsyncReadExt + Unpin,
    {
        async fn read_u8(&mut self) -> crate::Result<u8> {
            tokio::io::AsyncReadExt::read_u8(&mut self.0)
                .await
                .context("Failed to read u8")
        }
        async fn read_exact(&mut self, buf: &mut [u8]) -> crate::Result<()> {
            tokio::io::AsyncReadExt::read_exact(&mut self.0, buf)
                .await
                .map(|_| ())
                .context("Failed to read exact")
        }
    }

    let mut reader = TokioReader(reader);
    let v: T = T::decode_async(&mut reader).await?;
    Ok(v)
}

#[cfg(not(target_arch = "wasm32"))]
pub async fn decode_from_hyper_stream<T, S, B>(stream: S) -> crate::Result<T>
where
    T: Decode,
    S: futures::StreamExt<Item = std::result::Result<B::Data, B::Error>> + Unpin,
    B: hyper::body::Body + Unpin,
{
    struct StreamReader<S, B>(S, BytesMut, std::marker::PhantomData<B>);

    impl<S, B> StreamReader<S, B>
    where
        S: futures::StreamExt<Item = std::result::Result<B::Data, B::Error>> + Unpin,
        B: hyper::body::Body + Unpin,
    {
        async fn fill_buf(&mut self, until: usize) -> crate::Result<()> {
            while self.1.remaining() < until {
                let Some(n) = self.0.next().await else {
                    return Err(crate::Error::new("UnexpectedEof"));
                };
                let mut buf = match n {
                    Ok(n) => n,
                    Err(_) => {
                        return Err(crate::Error::new("UnexpectedEof"));
                    }
                };
                if buf.remaining() == 0 {
                    return Err(crate::Error::new("UnexpectedEof"));
                }
                while buf.remaining() > 0 {
                    let chunk = buf.chunk();
                    self.1.put_slice(chunk);
                    buf.advance(chunk.len());
                }
            }
            Ok(())
        }
    }

    impl<S, B> AsyncDecoder for StreamReader<S, B>
    where
        S: futures::StreamExt<Item = std::result::Result<B::Data, B::Error>> + Unpin,
        B: hyper::body::Body + Unpin,
    {
        async fn read_exact(&mut self, buf: &mut [u8]) -> crate::Result<()> {
            self.fill_buf(buf.len()).await?;
            self.1.copy_to_slice(buf);
            Ok(())
        }
        async fn read_u8(&mut self) -> crate::Result<u8> {
            self.fill_buf(1).await?;
            Ok(self.1.get_u8())
        }
    }

    let mut reader: StreamReader<S, B> =
        StreamReader(stream, BytesMut::new(), std::marker::PhantomData);
    T::decode_async(&mut reader).await
}

fn type_eq<T1, T2>() -> bool {
    type_name::<T1>() == type_name::<T2>()
}

unsafe fn transmute_slice<F, T>(slice: &[F]) -> &[T] { unsafe {
    assert_eq!(mem::size_of::<F>(), mem::size_of::<T>());
    assert_eq!(mem::align_of::<F>(), mem::align_of::<T>());
    mem::transmute(slice)
}}

unsafe fn transmute_vec<F, T>(vec: Vec<F>) -> Vec<T> { unsafe {
    assert_eq!(mem::size_of::<F>(), mem::size_of::<T>());
    assert_eq!(mem::align_of::<F>(), mem::align_of::<T>());
    mem::transmute(vec)
}}

unsafe fn transmute_array<F, T, const N: usize>(arr: [F; N]) -> [T; N] { unsafe {
    assert_eq!(mem::size_of::<F>(), mem::size_of::<T>());
    assert_eq!(mem::align_of::<F>(), mem::align_of::<T>());
    let ptr = &arr as *const [F; N] as *const [T; N];
    ptr.read()
}}

#[cfg(test)]
mod tests {
    use std::{
        collections::{HashMap, HashSet},
        io::Write,
    };

    use crate::prelude::*;

    use super::*;

    #[test]
    fn test_encoding_byte_vec() {
        let v = vec![1, 2, 3];

        let encoded = encode_to_bytes(&v).unwrap();
        let vec: Vec<u8> = decode_from_bytes(encoded).unwrap();

        assert_eq!(vec, v);
    }

    #[test]
    fn test_encoding_vec() {
        let v = vec!["test"];

        let encoded = encode_to_bytes(&v).unwrap();
        let vec: Vec<String> = decode_from_bytes(encoded).unwrap();

        assert_eq!(vec, v);
    }

    #[test]
    fn test_encoding_byte_array() {
        let a = [1, 2, 3];

        let encoded = encode_to_bytes(&a).unwrap();
        let array: [u8; 3] = decode_from_bytes(encoded).unwrap();

        assert_eq!(array, a);
    }

    #[test]
    fn test_encoding_array() {
        let a = [String::from("a")];

        let encoded = encode_to_bytes(&a).unwrap();
        let array: [String; 1] = decode_from_bytes(encoded).unwrap();

        assert_eq!(array, a);
    }

    #[test]
    fn test_encoding_vec_generic() {
        let v: Vec<i32> = vec![1, 2, 3];

        let encoded = encode_to_bytes(&v).unwrap();
        let vec: Vec<i32> = decode_from_bytes(encoded).unwrap();

        assert_eq!(vec, v);
    }

    #[derive(Debug, Clone, PartialEq, Eq, Encode, Decode)]
    struct TestStruct {
        f1: [u8; 32],
        f2: Bytes,
    }

    #[test]
    fn test_encoding_custom_struct() {
        let t = TestStruct {
            f1: [6; 32],
            f2: Bytes::new(),
        };

        let encoded = encode_to_bytes(&t).unwrap();
        let test: TestStruct = decode_from_bytes(encoded).unwrap();

        assert_eq!(test, t);
    }

    #[derive(Debug, Clone, PartialEq, Eq, Encode, Decode)]
    struct TestUnnamed([u8; 32], Bytes);

    #[test]
    fn test_encoding_custom_unnamed() {
        let t = TestUnnamed([1; 32], Bytes::new());

        let encoded = encode_to_bytes(&t).unwrap();
        let test: TestUnnamed = decode_from_bytes(encoded).unwrap();

        assert_eq!(test, t);
    }

    #[derive(Debug, Clone, PartialEq, Eq, Encode, Decode)]
    enum TestEnum {
        Unit,
        Unnamed(TestUnnamed),
        Named(TestStruct),
    }

    #[test]
    fn test_encoding_custom_enum() {
        let t = TestEnum::Unnamed(TestUnnamed([0; 32], Bytes::new()));

        let encoded = encode_to_bytes(&t).unwrap();
        let test: TestEnum = decode_from_bytes(encoded).unwrap();

        assert_eq!(test, t);
    }

    #[derive(Debug, Clone, PartialEq, Eq, Encode, Decode)]
    struct TestGeneric<D, const N: usize> {
        f1: [u8; N],
        f2: Bytes,
        f3: D,
    }

    #[test]
    fn test_encoding_custom_generic() {
        let t = TestGeneric {
            f1: [3; 32],
            f2: Bytes::new(),
            f3: String::from("test"),
        };

        let encoded = encode_to_bytes(&t).unwrap();
        let test: TestGeneric<String, 32> = decode_from_bytes(encoded).unwrap();

        assert_eq!(test, t);
    }

    #[derive(Debug, Clone, PartialEq, Eq, Encode, Decode)]
    struct TestUnsend<D, const N: usize> {
        f1: [u8; N],
        f2: Bytes,
        f3: D,
        _maker: std::marker::PhantomData<std::rc::Rc<()>>, // !Send
    }

    #[tokio::test]
    async fn test_encoding_tokio_unsend() {
        // let handle = tokio::spawn(async {
        let t = TestUnsend {
            f1: [3; 32],
            f2: Bytes::new(),
            f3: String::from("test"),
            _maker: std::marker::PhantomData,
        };

        let mut writer = std::io::Cursor::new(Vec::new());
        encode_to_tokio(&t, &mut writer).await.unwrap();
        writer.flush().unwrap();
        let encoded = writer.into_inner();

        let mut reader = std::io::Cursor::new(encoded);
        let test: TestUnsend<String, 32> = decode_from_tokio(&mut reader).await.unwrap();

        assert_eq!(test, t);
        // });

        // handle.await.unwrap();
    }

    #[tokio::test]
    async fn test_encoding_tokio_send() {
        let handle = tokio::spawn(async {
            let t = TestGeneric {
                f1: [3; 32],
                f2: Bytes::new(),
                f3: String::from("test"),
            };

            let mut writer = std::io::Cursor::new(Vec::new());
            encode_to_tokio(&t, &mut writer).await.unwrap();
            writer.flush().unwrap();
            let encoded = writer.into_inner();

            let mut reader = std::io::Cursor::new(encoded);
            let test: TestGeneric<String, 32> = decode_from_tokio(&mut reader).await.unwrap();

            assert_eq!(test, t);
        });

        handle.await.unwrap();
    }

    #[tokio::test]
    async fn test_encoding_hash_map() {
        let m = HashMap::from([(1u8, String::from("test")), (2u8, String::from("test1"))]);

        let encoded = encode_to_bytes(&m).unwrap();
        let map: HashMap<u8, String> = decode_from_bytes(encoded).unwrap();

        assert_eq!(map, m);
    }

    #[tokio::test]
    async fn test_encoding_hash_set() {
        let s = HashSet::from([1u8, 2u8]);

        let encoded = encode_to_bytes(&s).unwrap();
        let set: HashSet<u8> = decode_from_bytes(encoded).unwrap();

        assert_eq!(set, s);

        let s = HashSet::from([String::from("test"), String::from("test1")]);

        let encoded = encode_to_bytes(&s).unwrap();
        let set: HashSet<String> = decode_from_bytes(encoded).unwrap();

        assert_eq!(set, s);
    }

    #[derive(Encode, Decode)]
    enum NoVariant {}

    #[derive(Debug, PartialEq, Eq, Encode, Decode)]
    enum Skipped {
        First(String),

        #[encode(skip)]
        Second(String),

        Third(String, #[encode(skip)] String, String),

        Forth {
            first: String,
            #[encode(skip)]
            second: String,
            third: String,
        },
    }

    #[test]
    fn test_encoding_skip() {
        let s1 = Skipped::First(String::from("first"));
        let s2 = Skipped::Second(String::from("second"));
        let s3 = Skipped::Third(
            String::from("third1"),
            String::from("third2"),
            String::from("third3"),
        );
        let s4 = Skipped::Forth {
            first: String::from("forth1"),
            second: String::from("forth2"),
            third: String::from("forth3"),
        };

        let en1 = encode_to_bytes(&s1).unwrap();
        let en2 = encode_to_bytes(&s2).unwrap();
        let en3 = encode_to_bytes(&s3).unwrap();
        let en4 = encode_to_bytes(&s4).unwrap();

        let skip1: Skipped = decode_from_bytes(en1).unwrap();
        let skip2 = decode_from_bytes::<Skipped>(en2).unwrap_err();
        let skip3: Skipped = decode_from_bytes(en3).unwrap();
        let skip4: Skipped = decode_from_bytes(en4).unwrap();

        assert_eq!(skip1, s1);
        assert_eq!(skip2.context(), "Variant is skipped and cannot be decoded");
        assert_eq!(
            skip3,
            Skipped::Third(
                String::from("third1"),
                String::from(""),
                String::from("third3"),
            )
        );
        assert_eq!(
            skip4,
            Skipped::Forth {
                first: String::from("forth1"),
                second: String::from(""),
                third: String::from("forth3"),
            }
        );
    }
}