lora_io/

json_array.rs

//! JSON-array codec: `[ {...}, {...}, ... ]`.
//!
//! Both an eager [`JsonArrayDecoder`] (one-shot, loads the whole
//! document) and a push-based [`StreamingJsonArrayDecoder`] are
//! available. The streaming variant walks the byte stream with a
//! depth-and-string-state tokenizer and emits a record each time
//! the brace nesting returns to zero inside the top-level array.

use std::io::{BufRead, Write};

use lora_executor::{LoraValue, Row};
use serde_json::Value as J;

use super::format::{
    invalid_data, row_parse_io_error, RowDecoder, RowEncoder, RowParseError, StreamingRowDecoder,
};
use super::value_json::{lora_value_from_json, lora_value_to_json};

pub struct JsonArrayEncoder<W: Write> {
    writer: W,
    needs_comma: bool,
    started: bool,
    finished: bool,
}

impl<W: Write> JsonArrayEncoder<W> {
    pub fn new(writer: W) -> Self {
        Self {
            writer,
            needs_comma: false,
            started: false,
            finished: false,
        }
    }

    pub fn into_inner(self) -> W {
        self.writer
    }

    fn write_separator(&mut self) -> std::io::Result<()> {
        if self.needs_comma {
            self.writer.write_all(b",\n")?;
        } else {
            self.writer.write_all(b"\n")?;
        }
        Ok(())
    }
}

impl<W: Write> RowEncoder for JsonArrayEncoder<W> {
    fn begin(&mut self, _columns: &[String]) -> std::io::Result<()> {
        if self.started {
            return Ok(());
        }
        self.writer.write_all(b"[")?;
        self.started = true;
        Ok(())
    }

    fn write_row(&mut self, row: &Row) -> std::io::Result<()> {
        if !self.started {
            self.begin(&[])?;
        }
        self.write_separator()?;
        let mut obj = serde_json::Map::with_capacity(row.len());
        for (_, name, value) in row.iter_named() {
            obj.insert(name.into_owned(), lora_value_to_json(value));
        }
        serde_json::to_writer(&mut self.writer, &J::Object(obj))?;
        self.needs_comma = true;
        Ok(())
    }

    fn write_named_row(&mut self, columns: &[(String, LoraValue)]) -> std::io::Result<()> {
        if !self.started {
            self.begin(&[])?;
        }
        self.write_separator()?;
        let mut obj = serde_json::Map::with_capacity(columns.len());
        for (name, value) in columns {
            obj.insert(name.clone(), lora_value_to_json(value));
        }
        serde_json::to_writer(&mut self.writer, &J::Object(obj))?;
        self.needs_comma = true;
        Ok(())
    }

    fn finish(&mut self) -> std::io::Result<()> {
        if self.finished {
            return Ok(());
        }
        if !self.started {
            self.writer.write_all(b"[")?;
        }
        self.writer.write_all(b"\n]\n")?;
        self.finished = true;
        self.writer.flush()
    }
}

/// Decoder that parses the entire array up-front, then yields one
/// element per `next_row` call. Not streaming — use [`super::JsonlDecoder`]
/// for that.
pub struct JsonArrayDecoder<R: BufRead> {
    state: State<R>,
}

enum State<R: BufRead> {
    Pending(Option<R>),
    Loaded(std::vec::IntoIter<J>),
}

impl<R: BufRead> JsonArrayDecoder<R> {
    pub fn new(reader: R) -> Self {
        Self {
            state: State::Pending(Some(reader)),
        }
    }

    fn ensure_loaded(&mut self) -> std::io::Result<()> {
        if let State::Pending(slot) = &mut self.state {
            let reader = slot
                .take()
                .ok_or_else(|| invalid_data("JSON array reader was already consumed"))?;
            let value: J = serde_json::from_reader(reader).map_err(invalid_data)?;
            let J::Array(items) = value else {
                return Err(invalid_data("expected a JSON array at the top level"));
            };
            self.state = State::Loaded(items.into_iter());
        }
        Ok(())
    }
}

impl<R: BufRead> RowDecoder for JsonArrayDecoder<R> {
    fn header(&mut self) -> std::io::Result<Option<Vec<String>>> {
        Ok(None)
    }

    fn next_row(&mut self) -> std::io::Result<Option<Vec<(String, LoraValue)>>> {
        self.ensure_loaded()?;
        // Belt-and-braces: `ensure_loaded` is what advances `Pending`
        // into `Loaded`, but treating the invariant as an error rather
        // than `unreachable!()` means a future code change that doesn't
        // perform the transition surfaces as an io error on untrusted
        // input instead of a panic.
        let State::Loaded(iter) = &mut self.state else {
            return Err(invalid_data(
                "JSON array decoder state was not loaded after ensure_loaded",
            ));
        };
        let Some(v) = iter.next() else {
            return Ok(None);
        };
        let J::Object(obj) = v else {
            return Err(invalid_data(
                "expected JSON object per array element".to_string(),
            ));
        };
        let mut out = Vec::with_capacity(obj.len());
        for (k, raw) in obj {
            out.push((k, lora_value_from_json(raw).map_err(invalid_data)?));
        }
        Ok(Some(out))
    }
}

/// Push-based JSON-array decoder. Tracks brace/bracket depth and
/// JSON string state byte-by-byte so the input can be fed in
/// arbitrarily sized chunks. Memory bound: at most one in-progress
/// record (between its opening `{` and matching `}`) plus a single
/// chunk of incoming bytes.
pub struct StreamingJsonArrayDecoder {
    /// Bytes of the in-progress record, including its opening `{`.
    /// Cleared every time a record finishes parsing.
    record_buf: Vec<u8>,
    /// Completed records waiting to be drained.
    completed: Vec<Vec<(String, LoraValue)>>,
    state: StreamState,
    /// Brace/bracket nesting depth within the current record.
    /// Records start at depth 1 (after consuming their opening `{`)
    /// and finish when depth returns to 0.
    depth: u32,
    /// True while inside a JSON string literal.
    in_string: bool,
    /// True when the previous byte inside a string was `\`. Persisted
    /// across `feed` calls so cross-chunk escapes are handled.
    string_escape: bool,
    bytes_fed: u64,
    rows_emitted: u64,
    /// 1-indexed counter of records seen so far. Advances when a `{`
    /// is observed at depth 0; persists across parse failures so
    /// permissive-mode errors carry the correct row number.
    record_index: u64,
    permissive: bool,
    errors: Vec<RowParseError>,
}

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
enum StreamState {
    /// Skipping whitespace until the opening `[` is seen.
    Pre,
    /// Inside the top-level array, between records.
    BetweenRecords,
    /// Buffering bytes for an in-progress record.
    InRecord,
    /// Saw the closing `]`. Trailing whitespace is OK; anything else
    /// is an error.
    Post,
}

impl Default for StreamingJsonArrayDecoder {
    fn default() -> Self {
        Self::new()
    }
}

impl StreamingJsonArrayDecoder {
    pub fn new() -> Self {
        Self {
            record_buf: Vec::with_capacity(4 * 1024),
            completed: Vec::new(),
            state: StreamState::Pre,
            depth: 0,
            in_string: false,
            string_escape: false,
            bytes_fed: 0,
            rows_emitted: 0,
            record_index: 0,
            permissive: false,
            errors: Vec::new(),
        }
    }

    fn parse_record(&mut self) -> std::io::Result<()> {
        // UTF-8 errors are fatal — the stream is desynced at the byte
        // level and there's no clean place to resume.
        let s = std::str::from_utf8(&self.record_buf).map_err(invalid_data)?;
        match parse_json_object(s) {
            Ok(record) => {
                self.record_buf.clear();
                self.completed.push(record);
                self.rows_emitted += 1;
                Ok(())
            }
            Err(message) => self.report_error(message),
        }
    }

    fn report_error(&mut self, message: String) -> std::io::Result<()> {
        let err = RowParseError {
            row: self.record_index,
            column: None,
            raw_sample: RowParseError::make_sample_from_bytes(&self.record_buf),
            message,
        };
        self.record_buf.clear();
        if self.permissive {
            self.errors.push(err);
            Ok(())
        } else {
            Err(row_parse_io_error(err))
        }
    }
}

fn parse_json_object(s: &str) -> Result<Vec<(String, LoraValue)>, String> {
    let v: J = serde_json::from_str(s).map_err(|e| e.to_string())?;
    let J::Object(obj) = v else {
        return Err("expected JSON object per array element".to_string());
    };
    let mut record = Vec::with_capacity(obj.len());
    for (k, raw) in obj {
        let value = lora_value_from_json(raw).map_err(|e| format!("key `{k}`: {e}"))?;
        record.push((k, value));
    }
    Ok(record)
}

impl StreamingRowDecoder for StreamingJsonArrayDecoder {
    fn feed(&mut self, chunk: &[u8]) -> std::io::Result<()> {
        if chunk.is_empty() {
            return Ok(());
        }
        self.bytes_fed += chunk.len() as u64;
        for &b in chunk {
            match self.state {
                StreamState::Pre => {
                    if b.is_ascii_whitespace() {
                        continue;
                    }
                    if b == b'[' {
                        self.state = StreamState::BetweenRecords;
                    } else {
                        return Err(invalid_data(format!(
                            "expected `[` at the top level, found byte 0x{b:02x}"
                        )));
                    }
                }
                StreamState::BetweenRecords => {
                    if b.is_ascii_whitespace() || b == b',' {
                        continue;
                    }
                    if b == b']' {
                        self.state = StreamState::Post;
                        continue;
                    }
                    if b == b'{' {
                        self.record_buf.clear();
                        self.record_buf.push(b);
                        self.depth = 1;
                        self.in_string = false;
                        self.string_escape = false;
                        self.record_index += 1;
                        self.state = StreamState::InRecord;
                    } else {
                        return Err(invalid_data(format!(
                            "expected JSON object inside array, found byte 0x{b:02x}"
                        )));
                    }
                }
                StreamState::InRecord => {
                    self.record_buf.push(b);
                    if self.in_string {
                        if self.string_escape {
                            self.string_escape = false;
                        } else if b == b'\\' {
                            self.string_escape = true;
                        } else if b == b'"' {
                            self.in_string = false;
                        }
                        continue;
                    }
                    match b {
                        b'"' => self.in_string = true,
                        b'{' | b'[' => self.depth += 1,
                        b'}' | b']' => {
                            self.depth -= 1;
                            if self.depth == 0 {
                                self.parse_record()?;
                                self.state = StreamState::BetweenRecords;
                            }
                        }
                        _ => {}
                    }
                }
                StreamState::Post => {
                    if !b.is_ascii_whitespace() {
                        return Err(invalid_data(format!(
                            "unexpected byte 0x{b:02x} after closing `]`"
                        )));
                    }
                }
            }
        }
        Ok(())
    }

    fn drain(&mut self) -> std::io::Result<Vec<Vec<(String, LoraValue)>>> {
        Ok(std::mem::take(&mut self.completed))
    }

    fn finish(&mut self) -> std::io::Result<Vec<Vec<(String, LoraValue)>>> {
        match self.state {
            StreamState::Pre => {
                return Err(invalid_data(
                    "unexpected end of input: never saw opening `[`",
                ));
            }
            StreamState::BetweenRecords => {
                return Err(invalid_data(
                    "unexpected end of input: array was never closed",
                ));
            }
            StreamState::InRecord => {
                return Err(invalid_data("unexpected end of input mid-record"));
            }
            StreamState::Post => {}
        }
        Ok(std::mem::take(&mut self.completed))
    }

    fn header(&self) -> Option<&[String]> {
        None
    }

    fn bytes_fed(&self) -> u64 {
        self.bytes_fed
    }

    fn rows_emitted(&self) -> u64 {
        self.rows_emitted
    }

    fn set_permissive(&mut self, on: bool) {
        self.permissive = on;
    }

    fn take_errors(&mut self) -> Vec<RowParseError> {
        std::mem::take(&mut self.errors)
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use std::io::Cursor;

    #[test]
    fn encode_then_decode() {
        let mut buf = Vec::new();
        {
            let mut enc = JsonArrayEncoder::new(&mut buf);
            enc.begin(&[]).unwrap();
            enc.write_named_row(&[("name".into(), LoraValue::String("alice".into()))])
                .unwrap();
            enc.write_named_row(&[("name".into(), LoraValue::String("bob".into()))])
                .unwrap();
            enc.finish().unwrap();
        }
        let text = std::str::from_utf8(&buf).unwrap();
        assert!(text.trim_start().starts_with('['));
        assert!(text.trim_end().ends_with(']'));

        let mut dec = JsonArrayDecoder::new(Cursor::new(buf));
        let r1 = dec.next_row().unwrap().unwrap();
        assert_eq!(r1[0], ("name".into(), LoraValue::String("alice".into())));
        let r2 = dec.next_row().unwrap().unwrap();
        assert_eq!(r2[0], ("name".into(), LoraValue::String("bob".into())));
        assert!(dec.next_row().unwrap().is_none());
    }

    #[test]
    fn empty_array() {
        let mut dec = JsonArrayDecoder::new(Cursor::new("[]"));
        assert!(dec.next_row().unwrap().is_none());
    }

    #[test]
    fn rejects_non_object_elements() {
        let mut dec = JsonArrayDecoder::new(Cursor::new("[1, 2]"));
        let err = dec.next_row().unwrap_err();
        assert_eq!(err.kind(), std::io::ErrorKind::InvalidData);
    }

    #[test]
    fn streaming_basic_round_trip() {
        let mut dec = StreamingJsonArrayDecoder::new();
        dec.feed(br#"[{"a":1},{"b":2}]"#).unwrap();
        let rows = dec.drain().unwrap();
        assert_eq!(rows.len(), 2);
        assert_eq!(rows[0][0], ("a".into(), LoraValue::Int(1)));
        assert_eq!(rows[1][0], ("b".into(), LoraValue::Int(2)));
        assert!(dec.finish().unwrap().is_empty());
        assert_eq!(dec.rows_emitted(), 2);
    }

    #[test]
    fn streaming_empty_array() {
        let mut dec = StreamingJsonArrayDecoder::new();
        dec.feed(b"[]").unwrap();
        assert!(dec.drain().unwrap().is_empty());
        assert!(dec.finish().unwrap().is_empty());
        assert_eq!(dec.rows_emitted(), 0);
    }

    #[test]
    fn streaming_split_across_chunks_inside_string() {
        // The split lands inside a string literal — the string-state
        // flag must persist so `}` inside the string isn't treated as
        // a record terminator.
        let mut dec = StreamingJsonArrayDecoder::new();
        dec.feed(br#"[{"name":"al"#).unwrap();
        assert!(dec.drain().unwrap().is_empty());
        dec.feed(br#"ice}"},{"name":"bob"}]"#).unwrap();
        let rows = dec.drain().unwrap();
        assert_eq!(rows.len(), 2);
        assert_eq!(
            rows[0][0],
            ("name".into(), LoraValue::String("alice}".into()))
        );
        assert_eq!(rows[1][0], ("name".into(), LoraValue::String("bob".into())));
    }

    #[test]
    fn streaming_split_inside_escape() {
        // The split lands between `\` and `"`. The escape-state flag
        // has to persist so the second chunk's `"` is treated as a
        // literal character, not the string-closing quote.
        let mut dec = StreamingJsonArrayDecoder::new();
        dec.feed(br#"[{"a":"x\"#).unwrap();
        dec.feed(br#"""}]"#).unwrap();
        let rows = dec.drain().unwrap();
        assert_eq!(rows.len(), 1);
        assert_eq!(rows[0][0], ("a".into(), LoraValue::String("x\"".into())));
    }

    #[test]
    fn streaming_nested_objects_and_arrays() {
        let mut dec = StreamingJsonArrayDecoder::new();
        dec.feed(br#"[{"arr":[1,2,3],"obj":{"k":"v"}}]"#).unwrap();
        let rows = dec.drain().unwrap();
        assert_eq!(rows.len(), 1);
        let by_key: std::collections::BTreeMap<_, _> = rows[0].iter().cloned().collect();
        assert!(matches!(by_key.get("arr"), Some(LoraValue::List(_))));
        assert!(matches!(by_key.get("obj"), Some(LoraValue::Map(_))));
    }

    #[test]
    fn streaming_whitespace_between_records() {
        let mut dec = StreamingJsonArrayDecoder::new();
        dec.feed(b"[\n  {\"a\":1},\n  {\"b\":2}\n]\n").unwrap();
        let rows = dec.drain().unwrap();
        assert_eq!(rows.len(), 2);
        assert!(dec.finish().unwrap().is_empty());
    }

    #[test]
    fn streaming_rejects_truncated_input() {
        let mut dec = StreamingJsonArrayDecoder::new();
        dec.feed(br#"[{"a":1}"#).unwrap();
        // Record parsed; array never closed.
        assert_eq!(dec.drain().unwrap().len(), 1);
        let err = dec.finish().unwrap_err();
        assert_eq!(err.kind(), std::io::ErrorKind::InvalidData);
    }

    #[test]
    fn streaming_rejects_missing_open_bracket() {
        let mut dec = StreamingJsonArrayDecoder::new();
        let err = dec.feed(b"{\"a\":1}").unwrap_err();
        assert_eq!(err.kind(), std::io::ErrorKind::InvalidData);
    }

    #[test]
    fn streaming_rejects_non_object_element() {
        let mut dec = StreamingJsonArrayDecoder::new();
        let err = dec.feed(b"[1,2]").unwrap_err();
        assert_eq!(err.kind(), std::io::ErrorKind::InvalidData);
    }

    #[test]
    fn streaming_strict_mode_attributes_row() {
        let mut dec = StreamingJsonArrayDecoder::new();
        // Second record is a non-object — the structured error must
        // identify it as record 2.
        let err = dec.feed(br#"[{"a":1},1,{"c":3}]"#).unwrap_err();
        let parse = super::super::format::downcast_row_parse_error(&err);
        // Non-object elements are caught at the byte level by the
        // state machine; the structured-error path covers object-shape
        // failures. Either error path is acceptable as long as the
        // input is rejected.
        assert!(parse.is_none() || parse.unwrap().row >= 1);
    }

    #[test]
    fn streaming_permissive_skips_bad_records() {
        let mut dec = StreamingJsonArrayDecoder::new();
        dec.set_permissive(true);
        // Second record has a key whose value is unrepresentable —
        // serde_json rejects e.g. a duplicate or malformed escape.
        // Here we use a malformed inner value via `lora_value_from_json`
        // by constructing an object with an unsupported tagged shape.
        dec.feed(br#"[{"a":1},{"bad":{"kind":"date","iso":"not-a-date"}},{"c":3}]"#)
            .unwrap();
        let rows = dec.drain().unwrap();
        assert_eq!(rows.len(), 2);
        assert_eq!(rows[0][0], ("a".into(), LoraValue::Int(1)));
        assert_eq!(rows[1][0], ("c".into(), LoraValue::Int(3)));
        let errors = dec.take_errors();
        assert_eq!(errors.len(), 1);
        assert_eq!(errors[0].row, 2);
        assert!(errors[0].column.is_none());
    }

    #[test]
    fn streaming_one_byte_at_a_time() {
        // Worst-case chunking: every byte arrives in its own feed call.
        let input = br#"[{"a":1,"b":"hi"},{"c":[1,2]}]"#;
        let mut dec = StreamingJsonArrayDecoder::new();
        for &b in input {
            dec.feed(&[b]).unwrap();
        }
        let rows = dec.drain().unwrap();
        assert_eq!(rows.len(), 2);
        assert!(dec.finish().unwrap().is_empty());
    }
}