Skip to main content

ferrin_schema/
partial_json.rs

1//! Repair and parse truncated JSON.
2//!
3//! Streaming models emit JSON text incrementally. [`repair`] closes open
4//! strings, arrays and objects, completes partial literals (`tru` → `true`)
5//! and drops trailing separators so that any prefix of a valid document
6//! becomes parseable. [`parse_partial`] tries a direct parse first and falls
7//! back to repair.
8//!
9//! Derived from the `fixJson` state machine of the Vercel AI SDK (Apache-2.0,
10//! Copyright 2023 Vercel, Inc.), translated from TypeScript to Rust and
11//! modified; see `NOTICE`.
12
13use std::borrow::Cow;
14
15use serde_json::Value;
16
17/// Outcome of [`parse_partial`].
18#[derive(Debug, Clone, PartialEq, Eq)]
19pub struct PartialParse {
20    /// The parsed value, if parsing (possibly after repair) succeeded.
21    pub value: Option<Value>,
22    /// How the value was obtained.
23    pub state: PartialParseState,
24}
25
26/// How a partial parse succeeded or failed.
27#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
28#[non_exhaustive]
29pub enum PartialParseState {
30    /// The input parsed as-is.
31    SuccessfulParse,
32    /// The input parsed after repair.
33    RepairedParse,
34    /// The input could not be parsed even after repair.
35    FailedParse,
36}
37
38/// Parses `text`, repairing it first if a direct parse fails.
39#[must_use]
40pub fn parse_partial(text: &str) -> PartialParse {
41    if let Ok(value) = serde_json::from_str::<Value>(text) {
42        return PartialParse {
43            value: Some(value),
44            state: PartialParseState::SuccessfulParse,
45        };
46    }
47    let repaired = repair(text);
48    match serde_json::from_str::<Value>(&repaired) {
49        Ok(value) => PartialParse {
50            value: Some(value),
51            state: PartialParseState::RepairedParse,
52        },
53        Err(_) => PartialParse {
54            value: None,
55            state: PartialParseState::FailedParse,
56        },
57    }
58}
59
60#[derive(Debug, Clone, Copy, PartialEq, Eq)]
61enum State {
62    Root,
63    Finish,
64    InsideString,
65    InsideStringEscape,
66    InsideStringUnicodeEscape,
67    InsideLiteral,
68    InsideNumber,
69    InsideObjectStart,
70    InsideObjectKey,
71    InsideObjectKeyEscape,
72    InsideObjectAfterKey,
73    InsideObjectBeforeValue,
74    InsideObjectAfterValue,
75    InsideObjectAfterComma,
76    InsideArrayStart,
77    InsideArrayAfterValue,
78    InsideArrayAfterComma,
79}
80
81struct Repairer<'a> {
82    input: &'a str,
83    stack: Vec<State>,
84    /// Byte index one past the last character that belongs to valid output.
85    last_valid_end: usize,
86    literal_start: usize,
87    unicode_escape_digits: u8,
88    unicode_escape_value: u16,
89    pending_high_surrogate: bool,
90}
91
92impl Repairer<'_> {
93    fn top(&self) -> State {
94        // The stack always holds at least `Root`/`Finish`.
95        self.stack.last().copied().unwrap_or(State::Finish)
96    }
97
98    fn replace_top(&mut self, state: State) {
99        self.stack.pop();
100        self.stack.push(state);
101    }
102
103    fn process_value_start(&mut self, ch: char, start: usize, end: usize, swap: State) {
104        match ch {
105            '"' => {
106                self.last_valid_end = end;
107                self.replace_top(swap);
108                self.stack.push(State::InsideString);
109            }
110            'f' | 't' | 'n' => {
111                self.last_valid_end = end;
112                self.literal_start = start;
113                self.replace_top(swap);
114                self.stack.push(State::InsideLiteral);
115            }
116            '-' => {
117                self.replace_top(swap);
118                self.stack.push(State::InsideNumber);
119            }
120            '0'..='9' => {
121                self.last_valid_end = end;
122                self.replace_top(swap);
123                self.stack.push(State::InsideNumber);
124            }
125            '{' => {
126                self.last_valid_end = end;
127                self.replace_top(swap);
128                self.stack.push(State::InsideObjectStart);
129            }
130            '[' => {
131                self.last_valid_end = end;
132                self.replace_top(swap);
133                self.stack.push(State::InsideArrayStart);
134            }
135            _ => {}
136        }
137    }
138
139    fn process_after_object_value(&mut self, ch: char, end: usize) {
140        match ch {
141            ',' => self.replace_top(State::InsideObjectAfterComma),
142            '}' => {
143                self.last_valid_end = end;
144                self.stack.pop();
145            }
146            _ => {}
147        }
148    }
149
150    fn process_after_array_value(&mut self, ch: char, end: usize) {
151        match ch {
152            ',' => self.replace_top(State::InsideArrayAfterComma),
153            ']' => {
154                self.last_valid_end = end;
155                self.stack.pop();
156            }
157            _ => {}
158        }
159    }
160
161    fn step(&mut self, ch: char, start: usize, end: usize) {
162        match self.top() {
163            State::Root => self.process_value_start(ch, start, end, State::Finish),
164            State::Finish => {}
165            State::InsideObjectStart => match ch {
166                '"' => self.replace_top(State::InsideObjectKey),
167                '}' => {
168                    self.last_valid_end = end;
169                    self.stack.pop();
170                }
171                _ => {}
172            },
173            State::InsideObjectAfterComma => {
174                if ch == '"' {
175                    self.replace_top(State::InsideObjectKey);
176                }
177            }
178            State::InsideObjectKey => match ch {
179                '"' => self.replace_top(State::InsideObjectAfterKey),
180                '\\' => self.stack.push(State::InsideObjectKeyEscape),
181                _ => {}
182            },
183            State::InsideObjectKeyEscape => {
184                self.stack.pop();
185            }
186            State::InsideObjectAfterKey => {
187                if ch == ':' {
188                    self.replace_top(State::InsideObjectBeforeValue);
189                }
190            }
191            State::InsideObjectBeforeValue => {
192                self.process_value_start(ch, start, end, State::InsideObjectAfterValue);
193            }
194            State::InsideObjectAfterValue => self.process_after_object_value(ch, end),
195            State::InsideString => match ch {
196                '"' => {
197                    self.stack.pop();
198                    self.last_valid_end = end;
199                }
200                '\\' => self.stack.push(State::InsideStringEscape),
201                _ => self.last_valid_end = end,
202            },
203            State::InsideArrayStart => {
204                if ch == ']' {
205                    self.last_valid_end = end;
206                    self.stack.pop();
207                } else {
208                    // Whitespace before the first element is valid output; a
209                    // lone `-` is not (it would leave `[-]`).
210                    if ch != '-' {
211                        self.last_valid_end = end;
212                    }
213                    self.process_value_start(ch, start, end, State::InsideArrayAfterValue);
214                }
215            }
216            State::InsideArrayAfterValue => match ch {
217                ',' => self.replace_top(State::InsideArrayAfterComma),
218                ']' => {
219                    self.last_valid_end = end;
220                    self.stack.pop();
221                }
222                _ => self.last_valid_end = end,
223            },
224            State::InsideArrayAfterComma => {
225                self.process_value_start(ch, start, end, State::InsideArrayAfterValue);
226            }
227            State::InsideStringEscape => {
228                self.stack.pop();
229                if ch == 'u' {
230                    self.unicode_escape_digits = 0;
231                    self.unicode_escape_value = 0;
232                    self.stack.push(State::InsideStringUnicodeEscape);
233                } else {
234                    self.last_valid_end = end;
235                }
236            }
237            State::InsideStringUnicodeEscape => {
238                if let Some(digit) = ch.to_digit(16) {
239                    self.unicode_escape_value = (self.unicode_escape_value << 4) | digit as u16;
240                    self.unicode_escape_digits += 1;
241                    if self.unicode_escape_digits == 4 {
242                        self.stack.pop();
243                        if (0xD800..=0xDBFF).contains(&self.unicode_escape_value) {
244                            // A high surrogate alone is not a Unicode scalar.
245                            self.pending_high_surrogate = true;
246                        } else if !self.pending_high_surrogate
247                            || (0xDC00..=0xDFFF).contains(&self.unicode_escape_value)
248                        {
249                            self.pending_high_surrogate = false;
250                            self.last_valid_end = end;
251                        }
252                    }
253                }
254            }
255            State::InsideNumber => match ch {
256                '0'..='9' => self.last_valid_end = end,
257                'e' | 'E' | '-' | '+' | '.' => {}
258                ',' => {
259                    self.stack.pop();
260                    if self.top() == State::InsideArrayAfterValue {
261                        self.process_after_array_value(ch, end);
262                    }
263                    if self.top() == State::InsideObjectAfterValue {
264                        self.process_after_object_value(ch, end);
265                    }
266                }
267                '}' => {
268                    self.stack.pop();
269                    if self.top() == State::InsideObjectAfterValue {
270                        self.process_after_object_value(ch, end);
271                    }
272                }
273                ']' => {
274                    self.stack.pop();
275                    if self.top() == State::InsideArrayAfterValue {
276                        self.process_after_array_value(ch, end);
277                    }
278                }
279                _ => {
280                    self.stack.pop();
281                }
282            },
283            State::InsideLiteral => {
284                let partial = &self.input[self.literal_start..end];
285                if !"false".starts_with(partial)
286                    && !"true".starts_with(partial)
287                    && !"null".starts_with(partial)
288                {
289                    self.stack.pop();
290                    if self.top() == State::InsideObjectAfterValue {
291                        self.process_after_object_value(ch, end);
292                    } else if self.top() == State::InsideArrayAfterValue {
293                        self.process_after_array_value(ch, end);
294                    }
295                } else {
296                    self.last_valid_end = end;
297                }
298            }
299        }
300    }
301
302    fn finish(self) -> String {
303        let mut result = String::with_capacity(self.last_valid_end + self.stack.len());
304        result.push_str(&self.input[..self.last_valid_end]);
305        for state in self.stack.iter().rev() {
306            match state {
307                State::InsideString => result.push('"'),
308                State::InsideObjectKey
309                | State::InsideObjectAfterKey
310                | State::InsideObjectAfterComma
311                | State::InsideObjectStart
312                | State::InsideObjectBeforeValue
313                | State::InsideObjectAfterValue => result.push('}'),
314                State::InsideArrayStart
315                | State::InsideArrayAfterComma
316                | State::InsideArrayAfterValue => result.push(']'),
317                State::InsideLiteral => {
318                    let partial = &self.input[self.literal_start..];
319                    for literal in ["true", "false", "null"] {
320                        if let Some(rest) = literal.strip_prefix(partial) {
321                            result.push_str(rest);
322                            break;
323                        }
324                    }
325                }
326                State::Root
327                | State::Finish
328                | State::InsideStringEscape
329                | State::InsideObjectKeyEscape
330                | State::InsideStringUnicodeEscape
331                | State::InsideNumber => {}
332            }
333        }
334        result
335    }
336}
337
338/// Repairs truncated JSON so that it parses.
339///
340/// Returns the input unchanged (borrowed) when no repair was needed.
341#[must_use]
342pub fn repair(input: &str) -> Cow<'_, str> {
343    let mut repairer = Repairer {
344        input,
345        stack: vec![State::Root],
346        last_valid_end: 0,
347        literal_start: 0,
348        unicode_escape_digits: 0,
349        unicode_escape_value: 0,
350        pending_high_surrogate: false,
351    };
352    for (start, ch) in input.char_indices() {
353        let end = start + ch.len_utf8();
354        repairer.step(ch, start, end);
355    }
356    let result = repairer.finish();
357    if result == input {
358        Cow::Borrowed(input)
359    } else {
360        Cow::Owned(result)
361    }
362}