Added composition decoders and encoders

Author: tomDev5

tokio-util/src/codec/composition_codec.rs

@@ -0,0 +1,210 @@
+use bytes::BytesMut;
+
+use crate::codec::{Decoder, Encoder};
+
+/// A [`Decoder`] that composes two decoders into a single decoder.
+///
+/// The first decoder must output [`BytesMut`], to be connected to the second one.
+///
+/// This decoder will hold the intermediate result of the first decoder in a buffer.
+///
+/// [`Decoder`]: crate::codec::Decoder
+#[derive(Debug)]
+pub struct CompositionDecoder<D1, D2>
+where
+    D1: Decoder<Item = BytesMut>,
+    D2: Decoder,
+{
+    first: D1,
+    second: D2,
+    intermediate_buffer: BytesMut,
+}
+
+impl<D1, D2> CompositionDecoder<D1, D2>
+where
+    D1: Decoder<Item = BytesMut>,
+    D2: Decoder,
+{
+    /// Creates a new `CompositionDecoder` that will use `first` to decode data
+    /// and then pass the result to `second`.
+    pub fn new(first: D1, second: D2) -> Self {
+        Self {
+            first,
+            second,
+            intermediate_buffer: BytesMut::new(),
+        }
+    }
+}
+
+impl<D1, D2> Decoder for CompositionDecoder<D1, D2>
+where
+    D1: Decoder<Item = BytesMut>,
+    D2: Decoder,
+{
+    type Item = D2::Item;
+
+    type Error = CompositionDecoderError<D1, D2>;
+
+    fn decode(&mut self, src: &mut BytesMut) -> Result<Option<Self::Item>, Self::Error> {
+        // First, check if we can decode using our buffer
+        if !self.intermediate_buffer.is_empty() {
+            if let Some(result) = self
+                .second
+                .decode(&mut self.intermediate_buffer)
+                .map_err(CompositionDecoderError::D2)?
+            {
+                return Ok(Some(result));
+            }
+        }
+
+        // Then, try to load more data into the second decoder's buffer using the first decoder
+        if let Some(intermediate) = self
+            .first
+            .decode(src)
+            .map_err(CompositionDecoderError::D1)?
+        {
+            // Add the intermediate result to our buffer
+            self.intermediate_buffer.extend_from_slice(&intermediate);
+        }
+
+        // Now retry to decode using the second decoder
+        self.second
+            .decode(&mut self.intermediate_buffer)
+            .map_err(CompositionDecoderError::D2)
+    }
+}
+
+#[derive(Debug)]
+pub enum CompositionDecoderError<D1, D2>
+where
+    D1: Decoder,
+    D2: Decoder,
+{
+    Io(std::io::Error),
+    D1(D1::Error),
+    D2(D2::Error),
+}
+
+impl<D1, D2> From<std::io::Error> for CompositionDecoderError<D1, D2>
+where
+    D1: Decoder,
+    D2: Decoder,
+{
+    fn from(err: std::io::Error) -> Self {
+        CompositionDecoderError::Io(err)
+    }
+}
+
+/// Extension trait for [`Decoder`] that allows composing a decoder after `self`.
+pub trait DecoderCompositionExt: Decoder<Item = BytesMut> {
+    /// Compose a decoder after `self`.
+    fn compose_decoder<D>(self, other: D) -> CompositionDecoder<Self, D>
+    where
+        Self: Sized,
+        D: Decoder,
+    {
+        CompositionDecoder::new(self, other)
+    }
+}
+
+impl<D> DecoderCompositionExt for D
+where
+    D: Decoder<Item = BytesMut>,
+{
+    fn compose_decoder<D2>(self, other: D2) -> CompositionDecoder<Self, D2>
+    where
+        Self: Sized,
+        D2: Decoder,
+    {
+        CompositionDecoder::new(self, other)
+    }
+}
+
+/// An [`Encoder`] that composes two encoders into a single encoder.
+///
+/// The second encoder must take [`BytesMut`], to be connected to the first one.
+///
+/// [`Encoder`]: crate::codec::Encoder
+#[derive(Debug)]
+pub struct CompositionEncoder<E1, E2> {
+    first: E1,
+    second: E2,
+}
+
+impl<E1, E2> CompositionEncoder<E1, E2> {
+    /// Creates a new `CompositionEncoder` that will use `first` to encode data
+    /// and then pass the result to `second`.
+    pub fn new(first: E1, second: E2) -> Self {
+        Self { first, second }
+    }
+}
+
+impl<E1, E2, Item> Encoder<Item> for CompositionEncoder<E1, E2>
+where
+    E1: Encoder<Item>,
+    E2: Encoder<BytesMut>,
+{
+    type Error = CompositionEncoderError<E1, E2, Item>;
+
+    fn encode(&mut self, item: Item, dst: &mut BytesMut) -> Result<(), Self::Error> {
+        let mut intermediate = BytesMut::new();
+
+        // First, try to encode using our first encoder
+        self.first
+            .encode(item, &mut intermediate)
+            .map_err(CompositionEncoderError::D1)?;
+
+        // Then, try to encode using our second encoder
+        self.second
+            .encode(intermediate, dst)
+            .map_err(CompositionEncoderError::D2)?;
+
+        Ok(())
+    }
+}
+
+#[derive(Debug)]
+pub enum CompositionEncoderError<D1, D2, Item>
+where
+    D1: Encoder<Item>,
+    D2: Encoder<BytesMut>,
+{
+    Io(std::io::Error),
+    D1(D1::Error),
+    D2(D2::Error),
+}
+
+impl<D1, D2, Item> From<std::io::Error> for CompositionEncoderError<D1, D2, Item>
+where
+    D1: Encoder<Item>,
+    D2: Encoder<BytesMut>,
+{
+    fn from(err: std::io::Error) -> Self {
+        CompositionEncoderError::Io(err)
+    }
+}
+
+/// Extension trait for [`Encoder`] that allows composing a encoder after `self`.
+pub trait EncoderCompositionExt<E>: Encoder<E> {
+    /// Compose a encoder after `self`.
+    fn compose_encoder<D>(self, other: D) -> CompositionEncoder<Self, D>
+    where
+        Self: Sized,
+        D: Encoder<BytesMut>,
+    {
+        CompositionEncoder::new(self, other)
+    }
+}
+
+impl<D, E> EncoderCompositionExt<E> for D
+where
+    D: Encoder<E>,
+{
+    fn compose_encoder<D2>(self, other: D2) -> CompositionEncoder<Self, D2>
+    where
+        Self: Sized,
+        D2: Encoder<BytesMut>,
+    {
+        CompositionEncoder::new(self, other)
+    }
+}

tokio-util/src/codec/mod.rs

@@ -354,3 +354,8 @@ pub use self::lines_codec::{LinesCodec, LinesCodecError};
 
 mod any_delimiter_codec;
 pub use self::any_delimiter_codec::{AnyDelimiterCodec, AnyDelimiterCodecError};
+
+mod composition_codec;
+pub use self::composition_codec::{
+    CompositionDecoder, CompositionEncoder, DecoderCompositionExt, EncoderCompositionExt,
+};

tokio-util/tests/composition_decoder.rs

@@ -0,0 +1,83 @@
+use bytes::{Buf, BufMut, BytesMut};
+use tokio_util::codec::{Decoder, DecoderCompositionExt};
+
+#[derive(Debug)]
+struct StringDecoder {}
+
+impl Decoder for StringDecoder {
+    type Item = String;
+
+    type Error = std::io::Error;
+
+    fn decode(&mut self, src: &mut BytesMut) -> Result<Option<Self::Item>, Self::Error> {
+        if src.len() < size_of::<u32>() {
+            return Ok(None);
+        }
+
+        let length: usize = src.get_u32().try_into().expect("length is too large");
+        if src.len() < length.into() {
+            return Ok(None);
+        }
+
+        let string = src.split_to(length.into());
+        Ok(Some(
+            String::from_utf8(string.to_vec()).expect("string is not valid utf8"),
+        ))
+    }
+}
+
+#[derive(Debug)]
+struct BitwiseNotDecoder {}
+
+impl Decoder for BitwiseNotDecoder {
+    type Item = BytesMut;
+
+    type Error = std::io::Error;
+
+    fn decode(&mut self, src: &mut BytesMut) -> Result<Option<Self::Item>, Self::Error> {
+        let mut buffer = src.split();
+
+        bitwise_not_in_place(&mut buffer);
+
+        Ok(Some(buffer))
+    }
+}
+
+fn bitwise_not_in_place(buffer: &mut BytesMut) {
+    for byte in buffer.iter_mut() {
+        *byte = !*byte;
+    }
+}
+
+#[test]
+fn test_simple_chain() {
+    let mut decoder = BitwiseNotDecoder {}.compose_decoder(StringDecoder {});
+
+    let string = String::from("Hello, world!");
+    let mut buffer = BytesMut::new();
+    buffer.put_u32(string.len() as u32);
+    buffer.put(string.as_bytes());
+
+    bitwise_not_in_place(&mut buffer);
+
+    let result = decoder.decode(&mut buffer).unwrap().unwrap();
+    assert_eq!(result, string);
+}
+
+#[test]
+fn test_not_not_not() {
+    let mut decoder = BitwiseNotDecoder {}
+        .compose_decoder(BitwiseNotDecoder {})
+        .compose_decoder(BitwiseNotDecoder {})
+        .compose_decoder(StringDecoder {});
+
+    let string = String::from("Hello, world!");
+    let mut buffer = BytesMut::new();
+    buffer.put_u32(string.len() as u32);
+    buffer.put(string.as_bytes());
+
+    bitwise_not_in_place(&mut buffer);
+
+    let result = decoder.decode(&mut buffer).unwrap().unwrap();
+    assert_eq!(result, string);
+}

tokio-util/tests/composition_encoder.rs

@@ -0,0 +1,70 @@
+use bytes::{BufMut, BytesMut};
+use tokio_util::codec::{Encoder, EncoderCompositionExt};
+
+#[derive(Debug)]
+struct StringEncoder {}
+
+impl Encoder<String> for StringEncoder {
+    type Error = std::io::Error;
+
+    fn encode(&mut self, item: String, dst: &mut BytesMut) -> Result<(), Self::Error> {
+        let length: u32 = item.len().try_into().expect("length is too large");
+        dst.put_u32(length);
+        dst.put(item.as_bytes());
+        Ok(())
+    }
+}
+
+#[derive(Debug)]
+struct BitwiseNotEncoder {}
+
+impl Encoder<BytesMut> for BitwiseNotEncoder {
+    type Error = std::io::Error;
+
+    fn encode(&mut self, mut item: BytesMut, dst: &mut BytesMut) -> Result<(), Self::Error> {
+        bitwise_not_in_place(&mut item);
+        dst.extend_from_slice(&item);
+        Ok(())
+    }
+}
+
+fn bitwise_not_in_place(buffer: &mut BytesMut) {
+    for byte in buffer.iter_mut() {
+        *byte = !*byte;
+    }
+}
+
+#[test]
+fn test_simple_chain() {
+    let mut encoder = StringEncoder {}.compose_encoder(BitwiseNotEncoder {});
+
+    let string = String::from("Hello, world!");
+    let mut buffer = BytesMut::new();
+    buffer.put_u32(string.len() as u32);
+    buffer.put(string.as_bytes());
+
+    bitwise_not_in_place(&mut buffer);
+    let mut dst = BytesMut::new();
+
+    encoder.encode(string, &mut dst).unwrap();
+    assert_eq!(dst, buffer);
+}
+
+#[test]
+fn test_not_not() {
+    let mut encoder = StringEncoder {}
+        .compose_encoder(BitwiseNotEncoder {})
+        .compose_encoder(BitwiseNotEncoder {})
+        .compose_encoder(BitwiseNotEncoder {});
+
+    let string = String::from("Hello, world!");
+    let mut buffer = BytesMut::new();
+    buffer.put_u32(string.len() as u32);
+    buffer.put(string.as_bytes());
+
+    bitwise_not_in_place(&mut buffer);
+    let mut dst = BytesMut::new();
+
+    encoder.encode(string, &mut dst).unwrap();
+    assert_eq!(dst, buffer);
+}