libpna 0.35.0

//! I/O helper types for streaming reads and writes.

mod finish;

pub(crate) use self::finish::TryIntoInner;
use crate::chunk::MAX_CHUNK_DATA_LENGTH;
use std::io;

const MIN_CHUNK_BODY_SIZE: usize = 1;

#[inline]
const fn clamp_chunk_size(size: usize) -> usize {
    if size < MIN_CHUNK_BODY_SIZE {
        MIN_CHUNK_BODY_SIZE
    } else {
        size
    }
}

pub(crate) struct FlattenWriter {
    pub(crate) inner: Vec<Vec<u8>>,
    max_chunk_size: usize,
}

impl FlattenWriter {
    #[inline]
    pub(crate) const fn new() -> Self {
        Self::with_max_chunk_size(MAX_CHUNK_DATA_LENGTH)
    }

    #[inline]
    pub(crate) const fn with_max_chunk_size(max_chunk_size: usize) -> Self {
        Self {
            inner: Vec::new(),
            max_chunk_size: clamp_chunk_size(max_chunk_size),
        }
    }

    #[inline]
    pub(crate) fn set_max_chunk_size(&mut self, size: usize) {
        self.max_chunk_size = clamp_chunk_size(size);
    }
}

impl io::Write for FlattenWriter {
    #[inline]
    fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
        if buf.is_empty() {
            return Ok(0);
        }
        let written = buf.len();
        let mut rest = buf;
        if let Some(last) = self.inner.last_mut() {
            let free = self.max_chunk_size.saturating_sub(last.len());
            let (head, tail) = rest.split_at(rest.len().min(free));
            last.extend_from_slice(head);
            rest = tail;
        }
        self.inner
            .extend(rest.chunks(self.max_chunk_size).map(<[u8]>::to_vec));
        Ok(written)
    }

    #[inline]
    fn flush(&mut self) -> io::Result<()> {
        Ok(())
    }
}

/// A reader that chains an iterator of readers.
///
/// This is similar to [`std::io::Chain`] but for an arbitrary number
/// of readers from an iterator. It reads from the current reader until it
/// is exhausted and then moves to the next one.
pub(crate) struct ChainReader<I, R> {
    inner: I,
    current: Option<R>,
}

impl<I, R> ChainReader<I, R>
where
    I: Iterator<Item = R>,
{
    #[inline]
    pub(crate) fn new<In: IntoIterator<IntoIter = I>>(into_iter: In) -> Self {
        let mut inner = into_iter.into_iter();
        Self {
            current: inner.next(),
            inner,
        }
    }
}

impl<I, R> io::Read for ChainReader<I, R>
where
    I: Iterator<Item = R>,
    R: io::Read,
{
    #[inline]
    fn read(&mut self, buf: &mut [u8]) -> io::Result<usize> {
        if buf.is_empty() {
            return Ok(0);
        }
        while let Some(c) = &mut self.current {
            match c.read(buf) {
                Ok(0) => {
                    self.current = self.inner.next();
                    continue;
                }
                other => return other,
            }
        }
        Ok(0)
    }
}

#[cfg(test)]
pub(crate) mod tests {
    use super::*;
    use std::io::prelude::*;
    #[cfg(all(target_family = "wasm", target_os = "unknown"))]
    use wasm_bindgen_test::wasm_bindgen_test as test;

    /// Test helper that simulates a source which only yields a few bytes per read call.
    pub(crate) struct PartialReader<I>
    where
        I: IntoIterator<Item = u8>,
    {
        data: Vec<u8>,
        pos: usize,
        chunk_sizes: I::IntoIter,
    }

    impl<I> PartialReader<I>
    where
        I: IntoIterator<Item = u8>,
    {
        pub(crate) fn new(data: Vec<u8>, chunk_sizes: I) -> Self {
            Self {
                data,
                pos: 0,
                chunk_sizes: chunk_sizes.into_iter(),
            }
        }

        fn next_chunk_len(&mut self, fallback: usize) -> usize {
            self.chunk_sizes.next().map_or(fallback, usize::from)
        }
    }

    impl<I> Read for PartialReader<I>
    where
        I: IntoIterator<Item = u8>,
    {
        fn read(&mut self, buf: &mut [u8]) -> io::Result<usize> {
            if self.pos >= self.data.len() {
                return Ok(0);
            }
            let requested = self.next_chunk_len(buf.len());
            let remaining = self.data.len() - self.pos;
            let written = requested.min(remaining).min(buf.len());
            let start = self.pos;
            let end = start + written;
            buf[..written].copy_from_slice(&self.data[start..end]);
            self.pos = end;
            Ok(written)
        }
    }

    #[test]
    fn chain_empty() {
        let reader = ChainReader::new(Vec::<&[u8]>::new());
        assert_eq!("", io::read_to_string(reader).unwrap());
    }

    #[test]
    fn chain_empty_in_empty() {
        let reader = ChainReader::new([&b""[..]]);
        assert_eq!("", io::read_to_string(reader).unwrap());
    }

    #[test]
    fn chain_contain_empty() {
        let reader = ChainReader::new([&b"abc"[..], &b""[..], &b"def"[..]]);
        assert_eq!("abcdef", io::read_to_string(reader).unwrap());
    }

    #[test]
    fn chain_consecutive_empty() {
        let reader = ChainReader::new([&b"abc"[..], &b""[..], &b""[..], &b"def"[..]]);
        assert_eq!("abcdef", io::read_to_string(reader).unwrap());
    }

    #[test]
    fn flatten_writer_coalesces_small_writes_into_minimal_chunks() {
        const N: usize = 1000;
        const MAX: usize = 10;
        let mut writer = FlattenWriter::with_max_chunk_size(MAX);
        for _ in 0..N {
            writer.write_all(b"a").unwrap();
        }
        assert_eq!(writer.inner.concat(), vec![b'a'; N]);
        assert!(writer.inner.iter().all(|c| c.len() <= MAX));
        assert_eq!(writer.inner.len(), N.div_ceil(MAX));
    }

    #[test]
    fn flatten_writer_tops_up_last_then_splits_remainder() {
        const MAX: usize = 10;
        let mut writer = FlattenWriter::with_max_chunk_size(MAX);
        writer.write_all(b"abcd").unwrap();
        writer.write_all(&[b'X'; 25]).unwrap();
        let mut expected = Vec::from(&b"abcd"[..]);
        expected.resize(29, b'X');
        assert_eq!(writer.inner.concat(), expected);
        assert!(writer.inner.iter().all(|c| c.len() <= MAX));
        assert_eq!(writer.inner.len(), expected.len().div_ceil(MAX));
    }

    #[test]
    fn chain_cross_boundary_small_buf() {
        let mut reader = ChainReader::new([&b"ab"[..], &b""[..], &b"cde"[..]]);
        let mut out = [0u8; 2];

        // 1st read: fills with "ab"
        assert_eq!(reader.read(&mut out).unwrap(), 2);
        assert_eq!(&out, b"ab");

        // 2nd read: next chunk gets "cd"
        assert_eq!(reader.read(&mut out).unwrap(), 2);
        assert_eq!(&out, b"cd");

        // 3rd read: remaining "e"
        assert_eq!(reader.read(&mut out).unwrap(), 1);
        assert_eq!(&out[..1], b"e");

        // 4th read: EOF
        assert_eq!(reader.read(&mut out).unwrap(), 0);
    }
}