Codex Handbook
utils/stream-parser/src/utf8_stream.rs 333 lines
use std::error::Error;use std::fmt;use crate::StreamTextChunk;use crate::StreamTextParser;/// Error returned by [`Utf8StreamParser`] when streamed bytes are not valid UTF-8.#[derive(Debug, Clone, PartialEq, Eq)]pub enum Utf8StreamParserError {    /// The provided bytes contain an invalid UTF-8 sequence.    InvalidUtf8 {        /// Byte offset in the parser's buffered bytes where decoding failed.        valid_up_to: usize,        /// Length in bytes of the invalid sequence.        error_len: usize,    },    /// EOF was reached with a buffered partial UTF-8 code point.    IncompleteUtf8AtEof,}impl fmt::Display for Utf8StreamParserError {    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {        match self {            Self::InvalidUtf8 {                valid_up_to,                error_len,            } => write!(                f,                "invalid UTF-8 in streamed bytes at offset {valid_up_to} (error length {error_len})"            ),            Self::IncompleteUtf8AtEof => {                write!(f, "incomplete UTF-8 code point at end of stream")            }        }    }}impl Error for Utf8StreamParserError {}/// Wraps a [`StreamTextParser`] and accepts raw bytes, buffering partial UTF-8 code points.////// This is useful when upstream data arrives as `&[u8]` and a code point may be split across/// chunk boundaries (for example `0xC3` followed by `0xA9` for `é`).#[derive(Debug)]pub struct Utf8StreamParser<P> {    inner: P,    pending_utf8: Vec<u8>,}impl<P> Utf8StreamParser<P>where    P: StreamTextParser,{    pub fn new(inner: P) -> Self {        Self {            inner,            pending_utf8: Vec::new(),        }    }    /// Feed a raw byte chunk.    ///    /// If the chunk contains invalid UTF-8, this returns an error and rolls back the entire    /// pushed chunk so callers can decide how to recover without the inner parser seeing a partial    /// prefix from that chunk.    pub fn push_bytes(        &mut self,        chunk: &[u8],    ) -> Result<StreamTextChunk<P::Extracted>, Utf8StreamParserError> {        let old_len = self.pending_utf8.len();        self.pending_utf8.extend_from_slice(chunk);        match std::str::from_utf8(&self.pending_utf8) {            Ok(text) => {                let out = self.inner.push_str(text);                self.pending_utf8.clear();                Ok(out)            }            Err(err) => {                if let Some(error_len) = err.error_len() {                    self.pending_utf8.truncate(old_len);                    return Err(Utf8StreamParserError::InvalidUtf8 {                        valid_up_to: err.valid_up_to(),                        error_len,                    });                }                let valid_up_to = err.valid_up_to();                if valid_up_to == 0 {                    return Ok(StreamTextChunk::default());                }                let text = match std::str::from_utf8(&self.pending_utf8[..valid_up_to]) {                    Ok(text) => text,                    Err(prefix_err) => {                        self.pending_utf8.truncate(old_len);                        let error_len = prefix_err.error_len().unwrap_or(0);                        return Err(Utf8StreamParserError::InvalidUtf8 {                            valid_up_to: prefix_err.valid_up_to(),                            error_len,                        });                    }                };                let out = self.inner.push_str(text);                self.pending_utf8.drain(..valid_up_to);                Ok(out)            }        }    }    pub fn finish(&mut self) -> Result<StreamTextChunk<P::Extracted>, Utf8StreamParserError> {        if !self.pending_utf8.is_empty() {            match std::str::from_utf8(&self.pending_utf8) {                Ok(_) => {}                Err(err) => {                    if let Some(error_len) = err.error_len() {                        return Err(Utf8StreamParserError::InvalidUtf8 {                            valid_up_to: err.valid_up_to(),                            error_len,                        });                    }                    return Err(Utf8StreamParserError::IncompleteUtf8AtEof);                }            }        }        let mut out = if self.pending_utf8.is_empty() {            StreamTextChunk::default()        } else {            let text = match std::str::from_utf8(&self.pending_utf8) {                Ok(text) => text,                Err(err) => {                    let error_len = err.error_len().unwrap_or(0);                    return Err(Utf8StreamParserError::InvalidUtf8 {                        valid_up_to: err.valid_up_to(),                        error_len,                    });                }            };            let out = self.inner.push_str(text);            self.pending_utf8.clear();            out        };        let mut tail = self.inner.finish();        out.visible_text.push_str(&tail.visible_text);        out.extracted.append(&mut tail.extracted);        Ok(out)    }    /// Return the wrapped parser if no undecoded UTF-8 bytes are buffered.    ///    /// Use [`Self::finish`] first if you want to flush buffered text into the wrapped parser.    pub fn into_inner(self) -> Result<P, Utf8StreamParserError> {        if self.pending_utf8.is_empty() {            return Ok(self.inner);        }        match std::str::from_utf8(&self.pending_utf8) {            Ok(_) => Ok(self.inner),            Err(err) => {                if let Some(error_len) = err.error_len() {                    return Err(Utf8StreamParserError::InvalidUtf8 {                        valid_up_to: err.valid_up_to(),                        error_len,                    });                }                Err(Utf8StreamParserError::IncompleteUtf8AtEof)            }        }    }    /// Return the wrapped parser without validating or flushing buffered undecoded bytes.    ///    /// This may drop a partial UTF-8 code point that was buffered across chunk boundaries.    pub fn into_inner_lossy(self) -> P {        self.inner    }}#[cfg(test)]mod tests {    use super::Utf8StreamParser;    use super::Utf8StreamParserError;    use crate::CitationStreamParser;    use crate::StreamTextChunk;    use crate::StreamTextParser;    use pretty_assertions::assert_eq;    fn collect_bytes(        parser: &mut Utf8StreamParser<CitationStreamParser>,        chunks: &[&[u8]],    ) -> Result<StreamTextChunk<String>, Utf8StreamParserError> {        let mut all = StreamTextChunk::default();        for chunk in chunks {            let next = parser.push_bytes(chunk)?;            all.visible_text.push_str(&next.visible_text);            all.extracted.extend(next.extracted);        }        let tail = parser.finish()?;        all.visible_text.push_str(&tail.visible_text);        all.extracted.extend(tail.extracted);        Ok(all)    }    #[test]    fn utf8_stream_parser_handles_split_code_points_across_chunks() {        let chunks: [&[u8]; 3] = [            b"A\xC3",            b"\xA9<oai-mem-citation>\xE4",            b"\xB8\xAD</oai-mem-citation>Z",        ];        let mut parser = Utf8StreamParser::new(CitationStreamParser::new());        let out = match collect_bytes(&mut parser, &chunks) {            Ok(out) => out,            Err(err) => panic!("valid UTF-8 stream should parse: {err}"),        };        assert_eq!(out.visible_text, "AéZ");        assert_eq!(out.extracted, vec!["中".to_string()]);    }    #[test]    fn utf8_stream_parser_rolls_back_on_invalid_utf8_chunk() {        let mut parser = Utf8StreamParser::new(CitationStreamParser::new());        let first = match parser.push_bytes(&[0xC3]) {            Ok(out) => out,            Err(err) => panic!("leading byte may be buffered until next chunk: {err}"),        };        assert!(first.is_empty());        let err = match parser.push_bytes(&[0x28]) {            Ok(out) => panic!("invalid continuation byte should error, got output: {out:?}"),            Err(err) => err,        };        assert_eq!(            err,            Utf8StreamParserError::InvalidUtf8 {                valid_up_to: 0,                error_len: 1,            }        );        let second = match parser.push_bytes(&[0xA9, b'x']) {            Ok(out) => out,            Err(err) => panic!("state should still allow a valid continuation: {err}"),        };        let tail = match parser.finish() {            Ok(out) => out,            Err(err) => panic!("stream should finish: {err}"),        };        assert_eq!(second.visible_text, "éx");        assert!(second.extracted.is_empty());        assert!(tail.is_empty());    }    #[test]    fn utf8_stream_parser_rolls_back_entire_chunk_when_invalid_byte_follows_valid_prefix() {        let mut parser = Utf8StreamParser::new(CitationStreamParser::new());        let err = match parser.push_bytes(b"ok\xFF") {            Ok(out) => panic!("invalid byte should error, got output: {out:?}"),            Err(err) => err,        };        assert_eq!(            err,            Utf8StreamParserError::InvalidUtf8 {                valid_up_to: 2,                error_len: 1,            }        );        let next = match parser.push_bytes(b"!") {            Ok(out) => out,            Err(err) => panic!("parser should recover after rollback: {err}"),        };        assert_eq!(next.visible_text, "!");        assert!(next.extracted.is_empty());    }    #[test]    fn utf8_stream_parser_errors_on_incomplete_code_point_at_eof() {        let mut parser = Utf8StreamParser::new(CitationStreamParser::new());        let out = match parser.push_bytes(&[0xE2, 0x82]) {            Ok(out) => out,            Err(err) => panic!("partial code point should be buffered: {err}"),        };        assert!(out.is_empty());        let err = match parser.finish() {            Ok(out) => panic!("unfinished code point should error, got output: {out:?}"),            Err(err) => err,        };        assert_eq!(err, Utf8StreamParserError::IncompleteUtf8AtEof);    }    #[test]    fn utf8_stream_parser_into_inner_errors_when_partial_code_point_is_buffered() {        let mut parser = Utf8StreamParser::new(CitationStreamParser::new());        let out = match parser.push_bytes(&[0xC3]) {            Ok(out) => out,            Err(err) => panic!("partial code point should be buffered: {err}"),        };        assert!(out.is_empty());        let err = match parser.into_inner() {            Ok(_) => panic!("buffered partial code point should be rejected"),            Err(err) => err,        };        assert_eq!(err, Utf8StreamParserError::IncompleteUtf8AtEof);    }    #[test]    fn utf8_stream_parser_into_inner_lossy_drops_buffered_partial_code_point() {        let mut parser = Utf8StreamParser::new(CitationStreamParser::new());        let out = match parser.push_bytes(&[0xC3]) {            Ok(out) => out,            Err(err) => panic!("partial code point should be buffered: {err}"),        };        assert!(out.is_empty());        let mut inner = parser.into_inner_lossy();        let tail = inner.finish();        assert!(tail.is_empty());    }}