async-read-super-ext 0.1.0

use std::future::Future;
use std::io;
use std::marker::{PhantomPinned, Unpin};
use std::pin::Pin;
use std::task::{Context, Poll, ready};

use pin_project_lite::pin_project;
use tokio::io::AsyncBufRead;

pin_project! {
    #[derive(Debug)]
    #[must_use = "futures do nothing unless you `.await` or poll them"]
    pub struct Utf8BoundariesLossy<'a, R: ?Sized> {
        reader: &'a mut R,
        // The output buffer where valid UTF-8 data and replacement characters are written.
        // This buffer accumulates the processed data during the async reading operation.
        output: &'a mut Vec<u8>,
        // Whether the last read was incomplete UTF-8.
        incomplete_buf: [u8; 4],
        incomplete_buf_len: usize,
        // Make this future `!Unpin` for compatibility with async trait methods.
        #[pin]
        _pin: PhantomPinned,
    }
}

pub(crate) fn read_utf8_boundaries_lossy<'a, R>(
    reader: &'a mut R,
    buf: &'a mut Vec<u8>,
) -> Utf8BoundariesLossy<'a, R>
where
    R: AsyncBufRead + ?Sized + Unpin,
{
    Utf8BoundariesLossy {
        reader,
        output: buf,
        incomplete_buf: [0; 4],
        incomplete_buf_len: 0,
        _pin: PhantomPinned,
    }
}

pub(crate) fn read_utf8_boundaries_lossy_internal<R: AsyncBufRead + ?Sized>(
    mut reader: Pin<&mut R>,
    cx: &mut Context<'_>,
    output: &mut Vec<u8>,
    incomplete_buf: &mut [u8],
    incomplete_buf_len: &mut usize,
) -> Poll<io::Result<usize>> {
    const REPLACEMENT_CHARACTER_BYTES: &[u8] = &[0xEF, 0xBF, 0xBD];

    let mut read = 0;
    let mut used = 0;

    'outer: loop {
        reader.as_mut().consume(used);
        used = 0;

        if read != 0 {
            return Poll::Ready(Ok(read));
        }

        let mut available = ready!(reader.as_mut().poll_fill_buf(cx))?;

        if available.is_empty() {
            if *incomplete_buf_len != 0 {
                output.extend(REPLACEMENT_CHARACTER_BYTES.repeat(*incomplete_buf_len));
                read += REPLACEMENT_CHARACTER_BYTES.len() * *incomplete_buf_len;
                *incomplete_buf_len = 0;
                return Poll::Ready(Ok(read));
            }
            return Poll::Ready(Ok(0));
        }

        if *incomplete_buf_len != 0 {
            let mut tmp_buf = [0; 4];
            tmp_buf[..*incomplete_buf_len].copy_from_slice(&incomplete_buf[..*incomplete_buf_len]);
            tmp_buf[*incomplete_buf_len..].copy_from_slice(&available[..4 - *incomplete_buf_len]);
            match std::str::from_utf8(&tmp_buf) {
                Ok(valid) => {
                    assert!(!valid.is_empty());

                    output.extend(valid.as_bytes());
                    read += valid.len();
                    used += valid.len() - *incomplete_buf_len;
                    *incomplete_buf_len = 0;
                    continue 'outer;
                },
                Err(error) => {
                    let valid_up_to = error.valid_up_to();
                    assert!(valid_up_to > 0);

                    let (valid, _) = tmp_buf.split_at(valid_up_to);
                    output.extend(valid);
                    read += valid.len();
                    used += valid.len() - *incomplete_buf_len;
                    *incomplete_buf_len = 0;
                    continue 'outer;
                },
            }
        }

        loop {
            match std::str::from_utf8(available) {
                Ok("") => {
                    continue 'outer;
                },
                Ok(valid) => {
                    // All data is valid UTF-8
                    let valid_bytes = valid.as_bytes();
                    output.extend(valid_bytes);
                    read += valid_bytes.len();
                    used += valid_bytes.len();
                    continue 'outer;
                },
                Err(error) => {
                    let valid_up_to = error.valid_up_to();
                    let (valid, after_valid) = available.split_at(valid_up_to);
                    if valid_up_to > 0 {
                        output.extend(valid);
                        read += valid_up_to;
                        used += valid_up_to;
                        continue 'outer;
                    }

                    let Some(invalid_sequence_length) = error.error_len() else {
                        // Incomplete UTF-8 sequence at end of buffer, need more data
                        *incomplete_buf_len = after_valid.len();
                        incomplete_buf[..after_valid.len()].copy_from_slice(after_valid);
                        used += after_valid.len();
                        continue 'outer;
                    };
                    used += invalid_sequence_length;

                    // Skip the invalid bytes and count them as replacements
                    let replacement_str_bytes = REPLACEMENT_CHARACTER_BYTES.repeat(invalid_sequence_length);
                    read += replacement_str_bytes.len();
                    output.extend(replacement_str_bytes);

                    available = &after_valid[invalid_sequence_length..];
                    continue;
                },
            }
        }
    }
}

impl<R: AsyncBufRead + ?Sized + Unpin> Future for Utf8BoundariesLossy<'_, R> {
    type Output = io::Result<usize>;

    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<usize>> {
        let me = self.project();

        read_utf8_boundaries_lossy_internal(
            Pin::new(*me.reader),
            cx,
            me.output,
            me.incomplete_buf,
            me.incomplete_buf_len,
        )
    }
}

#[cfg(test)]
mod tests {
    use std::io::{self, Cursor};

    use tokio::io::{AsyncBufRead, BufReader};

    use crate::AsyncReadSuperExt;

    async fn read_utf8_boundaries_lossy_to_end<R: AsyncBufRead + ?Sized + Unpin>(
        reader: &mut R,
        result: &mut Vec<u8>,
    ) -> io::Result<usize> {
        let mut buf = Vec::new();
        loop {
            buf.clear();
            let bytes_read = reader.read_utf8_boundaries_lossy(&mut buf).await?;
            if bytes_read == 0 {
                break;
            }
            result.extend(&buf[..bytes_read]);
        }
        Ok(result.len())
    }

    #[tokio::test]
    async fn test_valid_utf8() {
        let data = "Hello, 🦀!".as_bytes();
        let mut reader = BufReader::new(Cursor::new(data));
        let mut buf = Vec::new();

        let bytes_read = read_utf8_boundaries_lossy_to_end(&mut reader, &mut buf)
            .await
            .unwrap();

        assert_eq!(bytes_read, data.len());
        assert_eq!(buf, data);
        assert_eq!(String::from_utf8(buf).unwrap(), "Hello, 🦀!");
    }

    #[tokio::test]
    async fn test_invalid_utf8() {
        // Create data with invalid UTF-8 sequences
        let mut data = Vec::new();
        data.extend_from_slice("Hello ".as_bytes());
        data.push(0xFF); // Invalid UTF-8 byte
        data.push(0xFE); // Invalid UTF-8 byte
        data.extend_from_slice(" World".as_bytes());

        let mut reader = BufReader::new(Cursor::new(data));
        let mut buf = Vec::new();

        read_utf8_boundaries_lossy_to_end(&mut reader, &mut buf)
            .await
            .unwrap();

        let result = String::from_utf8(buf).unwrap();
        println!("result: {result}");
        // Should have replacement characters for invalid bytes
        assert!(result.contains("Hello "));
        assert!(result.contains(" World"));
        assert!(result.contains(char::REPLACEMENT_CHARACTER));

        // Should have replacement characters for the 2 invalid bytes
        let replacement_count = result
            .chars()
            .filter(|&c| c == char::REPLACEMENT_CHARACTER)
            .count();
        assert_eq!(replacement_count, 2);
    }

    #[tokio::test]
    async fn test_incomplete_utf8_at_boundary() {
        // Create data with incomplete UTF-8 sequence at the end
        let mut data = Vec::new();
        data.extend_from_slice("Hello ".as_bytes());
        // Add incomplete UTF-8 sequence (first byte of a 3-byte character)
        data.push(0xE2); // Start of UTF-8 3-byte sequence but incomplete
        let data_len = data.len();

        let mut reader = BufReader::new(Cursor::new(data));
        let mut final_result = vec![];
        let mut buf = Vec::new();
        loop {
            buf.clear();
            let bytes_read = read_utf8_boundaries_lossy_to_end(&mut reader, &mut buf)
                .await
                .unwrap();
            if bytes_read == 0 {
                break;
            }

            final_result.extend(&buf[..bytes_read]);
        }
        let final_result_len = final_result.len();
        let result = String::from_utf8(final_result).unwrap();

        // Should have "Hello " plus the incomplete byte as raw data
        assert!(result.contains("Hello "));
        // The incomplete sequence should be treated as raw bytes
        assert_eq!(final_result_len, data_len + 2);
    }

    #[tokio::test]
    async fn test_empty_input() {
        let data: &[u8] = &[];
        let mut reader = BufReader::new(Cursor::new(data));
        let mut buf = Vec::new();

        let bytes_read = read_utf8_boundaries_lossy_to_end(&mut reader, &mut buf)
            .await
            .unwrap();

        assert_eq!(bytes_read, 0);
        assert!(buf.is_empty());
    }

    #[tokio::test]
    async fn test_leading_invalid_utf8() {
        let mut data = vec![0xFF, 0xFE, 0xFD]; // Invalid UTF-8 bytes
        data.extend_from_slice("Hello, 🦀!\n".as_bytes());

        let mut reader = BufReader::new(Cursor::new(data));
        let mut buf = Vec::new();

        read_utf8_boundaries_lossy_to_end(&mut reader, &mut buf)
            .await
            .unwrap();

        let result = String::from_utf8(buf).unwrap();

        // Should have replacement characters for invalid bytes followed by valid text
        assert!(result.contains("Hello, 🦀!\n"));
        assert!(result.contains(char::REPLACEMENT_CHARACTER));

        // Should have 3 replacement characters for the 3 invalid bytes
        let replacement_count = result
            .chars()
            .filter(|&c| c == char::REPLACEMENT_CHARACTER)
            .count();
        assert_eq!(replacement_count, 3);
    }

    #[tokio::test]
    async fn test_trailing_invalid_utf8() {
        let mut data = "Hello, 🦀!\n".as_bytes().to_vec();
        data.extend_from_slice(&[0xFF, 0xFE, 0xFD]); // Invalid UTF-8 bytes
        let data_len = data.len();

        let mut reader = BufReader::new(Cursor::new(data));
        let mut buf = Vec::new();

        read_utf8_boundaries_lossy_to_end(&mut reader, &mut buf)
            .await
            .unwrap();

        let result = String::from_utf8(buf).unwrap();

        // Should have valid text followed by replacement characters
        assert!(result.contains("Hello, 🦀!\n"));
        assert!(result.contains(char::REPLACEMENT_CHARACTER));

        // Should have 3 replacement characters for the 3 invalid bytes
        let replacement_count = result
            .chars()
            .filter(|&c| c == char::REPLACEMENT_CHARACTER)
            .count();
        assert_eq!(replacement_count, 3);

        // Total length should account for replacement characters
        assert_eq!(result.len(), data_len + 6); // 3 invalid bytes become 3 replacement characters (3 bytes each)
    }

    #[tokio::test]
    async fn test_mixed_invalid_utf8() {
        let mut data = "Hello".as_bytes().to_vec();
        data.extend_from_slice(&[0xFF, 0xFE]); // Invalid UTF-8
        data.extend_from_slice(", World!".as_bytes());
        let data_len = data.len();

        let mut reader = BufReader::new(Cursor::new(data));
        let mut buf = Vec::new();

        read_utf8_boundaries_lossy_to_end(&mut reader, &mut buf)
            .await
            .unwrap();

        let result = String::from_utf8(buf).unwrap();

        // Should have valid text with replacement characters in the middle
        assert!(result.contains("Hello"));
        assert!(result.contains(", World!"));
        assert!(result.contains(char::REPLACEMENT_CHARACTER));

        // Should have 2 replacement characters for the 2 invalid bytes
        let replacement_count = result
            .chars()
            .filter(|&c| c == char::REPLACEMENT_CHARACTER)
            .count();
        assert_eq!(replacement_count, 2);

        // Total length should account for replacement characters
        assert_eq!(result.len(), data_len + 4); // 2 invalid bytes become 2 replacement characters (3 bytes each)
    }

    #[tokio::test]
    async fn test_large_input() {
        let data = "Hello, 🦀!".repeat(1024 * 10);
        let data_len = data.len();
        let mut reader = BufReader::new(Cursor::new(data.as_bytes()));
        let mut buf = Vec::new();

        let bytes_read = read_utf8_boundaries_lossy_to_end(&mut reader, &mut buf)
            .await
            .unwrap();

        let result = String::from_utf8(buf).unwrap();

        let replacement_count = result
            .chars()
            .filter(|&c| c == char::REPLACEMENT_CHARACTER)
            .count();
        assert_eq!(replacement_count, 0);
        assert_eq!(bytes_read, data_len);
        assert_eq!(result, data);
    }

    #[tokio::test]
    async fn test_large_invalid_utf8() {
        let data = &[0xFF, 0xFE, 0xFD].repeat(1024 * 1024);
        let data_len = data.len();

        let mut reader = BufReader::new(Cursor::new(data));
        let mut buf = Vec::new();

        let bytes_read = read_utf8_boundaries_lossy_to_end(&mut reader, &mut buf)
            .await
            .unwrap();

        let result = String::from_utf8(buf).unwrap();

        let replacement_count = result
            .chars()
            .filter(|&c| c == char::REPLACEMENT_CHARACTER)
            .count();
        assert_eq!(replacement_count, data_len);
        assert_eq!(bytes_read, data_len * 3);
    }

    #[tokio::test]
    async fn test_large_input_with_incomplete_utf8() {
        // Create data with incomplete UTF-8 sequence at the end
        let mut data = Vec::new();
        data.extend_from_slice("Hello, 🦀".repeat(1024 * 1024).as_bytes());
        // Add incomplete UTF-8 sequence (first byte of a 3-byte character)
        data.push(0xE2); // Start of UTF-8 3-byte sequence but incomplete
        let data_len = data.len();

        let mut reader = BufReader::new(Cursor::new(data));
        let mut buf = Vec::new();

        let bytes_read = read_utf8_boundaries_lossy_to_end(&mut reader, &mut buf)
            .await
            .unwrap();

        let result = String::from_utf8(buf).unwrap();

        let replacement_count = result
            .chars()
            .filter(|&c| c == char::REPLACEMENT_CHARACTER)
            .count();
        assert_eq!(replacement_count, 1);
        assert_eq!(bytes_read, data_len + 2);
    }

    #[tokio::test]
    async fn test_large_mixed_content() {
        let mut data = "Hello".repeat(1024 * 1024).into_bytes();
        data.extend_from_slice(&[0xFF, 0xFE]); // Invalid UTF-8
        data.extend_from_slice(", World!".as_bytes());
        let data_len = data.len();

        let mut reader = BufReader::new(Cursor::new(data));
        let mut buf = Vec::new();

        read_utf8_boundaries_lossy_to_end(&mut reader, &mut buf)
            .await
            .unwrap();

        let result = String::from_utf8(buf).unwrap();

        // Should have valid text with replacement characters
        assert!(result.contains(&"Hello".repeat(1024 * 1024)));
        assert!(result.contains(", World!"));
        assert!(result.contains(char::REPLACEMENT_CHARACTER));

        // Should have 2 replacement characters for the 2 invalid bytes
        let replacement_count = result
            .chars()
            .filter(|&c| c == char::REPLACEMENT_CHARACTER)
            .count();
        assert_eq!(replacement_count, 2);

        // Total length should account for replacement characters
        assert_eq!(result.len(), data_len + 4); // 2 invalid bytes become 2 replacement characters (3 bytes each)
    }
}