zkvmc-util 0.0.1

use std::io::Write;

/// Scans lines and prefixes lines with a given prefix. Will work even
/// when a write contains multiple lines or incomplete lines between
/// writes. It will not prefix empty lines.
#[derive(Debug)]
pub struct PrefixWriter<W: Write> {
    prefix: &'static str,
    writer: W,

    remainder: Option<String>,
}

impl<W: Write> Write for PrefixWriter<W> {
    fn write(&mut self, buf: &[u8]) -> std::io::Result<usize> {
        let input = if let Some(ref remainder) = self.remainder.take() {
            format!("{}{}", remainder, String::from_utf8_lossy(buf)).into()
        } else {
            String::from_utf8_lossy(buf)
        };

        let input_ends_with_newline = input.ends_with('\n');

        let mut lines = input.lines().peekable();
        while let Some(line) = lines.next() {
            if lines.peek().is_none() && !input_ends_with_newline {
                self.remainder = Some(line.to_owned());
                break;
            }

            if !line.is_empty() {
                self.writer.write_all(self.prefix.as_bytes())?;
            }

            self.writer.write_all(line.as_bytes())?;
            self.writer.write_all(b"\n")?;
        }

        Ok(buf.len())
    }

    fn flush(&mut self) -> std::io::Result<()> {
        if let Some(ref remainder) = self.remainder.take() {
            if !remainder.is_empty() {
                self.writer.write_all(self.prefix.as_bytes())?;
                self.writer.write_all(remainder.as_bytes())?;
            }
        }
        self.writer.flush()
    }
}

#[allow(unused)]
impl<W: Write> PrefixWriter<W> {
    /// Create a new [`PrefixWriter`] using the prefix for prefixing
    /// lines and the writer for writing the output of the prefixed
    /// lines.
    pub fn new(prefix: &'static str, writer: W) -> Self {
        Self {
            prefix,
            writer,

            remainder: None,
        }
    }

    /// Set a new prefix for [`PrefixWriter`].
    pub fn with_prefix(self, prefix: &'static str) -> Self {
        Self { prefix, ..self }
    }

    /// Set a new writer for [`PrefixWriter`].
    pub fn with_writer(self, writer: W) -> Self {
        Self { writer, ..self }
    }
}