nuvec_io/
json.rs

1//! A minimal JSON parser (RFC 8259) for safetensors headers and model files (configs,
2//! tokenizers), so that no serialization framework is needed.
3
4use std::collections::BTreeMap;
5
6use nuvec_core::{Error, Result};
7
8/// A parsed JSON value.
9#[derive(Debug, Clone, PartialEq)]
10pub enum Json {
11    /// `null`.
12    Null,
13    /// `true` or `false`.
14    Bool(bool),
15    /// A number without fraction or exponent that fits an `i64`.
16    Int(i64),
17    /// Any other number.
18    Float(f64),
19    /// A string.
20    Str(String),
21    /// An array.
22    Array(Vec<Json>),
23    /// An object, keys in file order.
24    Object(Vec<(String, Json)>),
25}
26
27impl Json {
28    /// The value of `key` in an object (the first, if repeated).
29    pub fn get(&self, key: &str) -> Option<&Json> {
30        match self {
31            Self::Object(items) => items.iter().find(|(k, _)| k == key).map(|(_, v)| v),
32            _ => None,
33        }
34    }
35
36    /// The string, if this is one.
37    pub fn as_str(&self) -> Option<&str> {
38        match self {
39            Self::Str(s) => Some(s),
40            _ => None,
41        }
42    }
43
44    /// The integer, if this is one.
45    pub fn as_i64(&self) -> Option<i64> {
46        match self {
47            Self::Int(i) => Some(*i),
48            _ => None,
49        }
50    }
51
52    /// The number (integers converted), if this is one.
53    pub fn as_f64(&self) -> Option<f64> {
54        match self {
55            Self::Int(i) => Some(*i as f64),
56            Self::Float(f) => Some(*f),
57            _ => None,
58        }
59    }
60
61    /// The boolean, if this is one.
62    pub fn as_bool(&self) -> Option<bool> {
63        match self {
64            Self::Bool(b) => Some(*b),
65            _ => None,
66        }
67    }
68
69    /// The elements, if this is an array.
70    pub fn as_array(&self) -> Option<&[Json]> {
71        match self {
72            Self::Array(items) => Some(items),
73            _ => None,
74        }
75    }
76
77    /// The entries, if this is an object.
78    pub fn as_object(&self) -> Option<&[(String, Json)]> {
79        match self {
80            Self::Object(items) => Some(items),
81            _ => None,
82        }
83    }
84}
85
86fn invalid(msg: &str, at: usize) -> Error {
87    Error::InvalidArgument(format!("JSON: {msg} at byte {at}"))
88}
89
90/// The deepest nesting of arrays and objects accepted (serde_json's default): deeper
91/// input is an error instead of a stack overflow.
92const MAX_DEPTH: usize = 128;
93
94struct Parser<'a> {
95    s: &'a [u8],
96    i: usize,
97    /// Arrays and objects currently open.
98    depth: usize,
99}
100
101impl Parser<'_> {
102    fn ws(&mut self) {
103        while self.s.get(self.i).is_some_and(u8::is_ascii_whitespace) {
104            self.i += 1;
105        }
106    }
107
108    fn eat(&mut self, c: u8) -> Result<()> {
109        self.ws();
110        if self.s.get(self.i) == Some(&c) {
111            self.i += 1;
112            Ok(())
113        } else {
114            Err(invalid(&format!("expected `{}`", c as char), self.i))
115        }
116    }
117
118    /// Enters the array or object at the current byte.
119    fn open(&mut self) -> Result<()> {
120        if self.depth == MAX_DEPTH {
121            return Err(invalid("nesting too deep", self.i));
122        }
123        self.depth += 1;
124        self.i += 1;
125        Ok(())
126    }
127
128    fn value(&mut self) -> Result<Json> {
129        self.ws();
130        match self.s.get(self.i) {
131            Some(b'{') => {
132                self.open()?;
133                let mut items = Vec::new();
134                self.ws();
135                if self.s.get(self.i) == Some(&b'}') {
136                    self.i += 1;
137                    self.depth -= 1;
138                    return Ok(Json::Object(items));
139                }
140                loop {
141                    self.ws();
142                    let key = self.string()?;
143                    self.eat(b':')?;
144                    items.push((key, self.value()?));
145                    self.ws();
146                    match self.s.get(self.i) {
147                        Some(b',') => self.i += 1,
148                        Some(b'}') => {
149                            self.i += 1;
150                            self.depth -= 1;
151                            return Ok(Json::Object(items));
152                        }
153                        _ => return Err(invalid("expected `,` or `}`", self.i)),
154                    }
155                }
156            }
157            Some(b'[') => {
158                self.open()?;
159                let mut items = Vec::new();
160                self.ws();
161                if self.s.get(self.i) == Some(&b']') {
162                    self.i += 1;
163                    self.depth -= 1;
164                    return Ok(Json::Array(items));
165                }
166                loop {
167                    items.push(self.value()?);
168                    self.ws();
169                    match self.s.get(self.i) {
170                        Some(b',') => self.i += 1,
171                        Some(b']') => {
172                            self.i += 1;
173                            self.depth -= 1;
174                            return Ok(Json::Array(items));
175                        }
176                        _ => return Err(invalid("expected `,` or `]`", self.i)),
177                    }
178                }
179            }
180            Some(b'"') => Ok(Json::Str(self.string()?)),
181            Some(b't') if self.s[self.i..].starts_with(b"true") => {
182                self.i += 4;
183                Ok(Json::Bool(true))
184            }
185            Some(b'f') if self.s[self.i..].starts_with(b"false") => {
186                self.i += 5;
187                Ok(Json::Bool(false))
188            }
189            Some(b'n') if self.s[self.i..].starts_with(b"null") => {
190                self.i += 4;
191                Ok(Json::Null)
192            }
193            Some(c) if *c == b'-' || c.is_ascii_digit() => self.number(),
194            _ => Err(invalid("value expected", self.i)),
195        }
196    }
197
198    /// The digits at the current byte (at least one); returns how many.
199    fn digits(&mut self, start: usize) -> Result<usize> {
200        let from = self.i;
201        while self.s.get(self.i).is_some_and(u8::is_ascii_digit) {
202            self.i += 1;
203        }
204        match self.i - from {
205            0 => Err(invalid("number expected", start)),
206            n => Ok(n),
207        }
208    }
209
210    /// A number in RFC 8259's grammar: `-? (0 | [1-9][0-9]*) (. [0-9]+)? ([eE] [+-]? [0-9]+)?`.
211    fn number(&mut self) -> Result<Json> {
212        let start = self.i;
213        if self.s.get(self.i) == Some(&b'-') {
214            self.i += 1;
215        }
216        let first = self.i;
217        if self.digits(start)? > 1 && self.s[first] == b'0' {
218            return Err(invalid("leading zero", start));
219        }
220        let mut integer = true;
221        if self.s.get(self.i) == Some(&b'.') {
222            self.i += 1;
223            self.digits(start)?;
224            integer = false;
225        }
226        if matches!(self.s.get(self.i), Some(b'e' | b'E')) {
227            self.i += 1;
228            if matches!(self.s.get(self.i), Some(b'+' | b'-')) {
229                self.i += 1;
230            }
231            self.digits(start)?;
232            integer = false;
233        }
234        // ASCII by construction.
235        let text = std::str::from_utf8(&self.s[start..self.i]).unwrap_or("");
236        match text.parse() {
237            Ok(n) if integer => Ok(Json::Int(n)),
238            _ => text
239                .parse()
240                .map(Json::Float)
241                .map_err(|_| invalid("number expected", start)),
242        }
243    }
244
245    fn string(&mut self) -> Result<String> {
246        if self.s.get(self.i) != Some(&b'"') {
247            return Err(invalid("string expected", self.i));
248        }
249        self.i += 1;
250        let mut out = String::new();
251        loop {
252            let start = self.i;
253            while self.s.get(self.i).is_some_and(|&c| c != b'"' && c != b'\\') {
254                self.i += 1;
255            }
256            out.push_str(
257                std::str::from_utf8(&self.s[start..self.i])
258                    .map_err(|_| invalid("invalid UTF-8", start))?,
259            );
260            match self.s.get(self.i) {
261                Some(b'"') => {
262                    self.i += 1;
263                    return Ok(out);
264                }
265                Some(b'\\') => {
266                    let esc = *self
267                        .s
268                        .get(self.i + 1)
269                        .ok_or_else(|| invalid("bad escape", self.i))?;
270                    self.i += 2;
271                    match esc {
272                        b'"' => out.push('"'),
273                        b'\\' => out.push('\\'),
274                        b'/' => out.push('/'),
275                        b'b' => out.push('\u{8}'),
276                        b'f' => out.push('\u{c}'),
277                        b'n' => out.push('\n'),
278                        b'r' => out.push('\r'),
279                        b't' => out.push('\t'),
280                        b'u' => {
281                            let at = self.i - 2;
282                            let mut code = self.hex4()?;
283                            // A high surrogate must be followed by an escaped low one; a
284                            // lone surrogate of either kind is not a character.
285                            if (0xD800..0xDC00).contains(&code) {
286                                if !self.s[self.i..].starts_with(b"\\u") {
287                                    return Err(invalid("lone surrogate", at));
288                                }
289                                self.i += 2;
290                                let low = self.hex4()?;
291                                if !(0xDC00..0xE000).contains(&low) {
292                                    return Err(invalid("invalid surrogate pair", at));
293                                }
294                                code = 0x10000 + ((code - 0xD800) << 10) + (low - 0xDC00);
295                            }
296                            out.push(
297                                char::from_u32(code)
298                                    .ok_or_else(|| invalid("lone surrogate", at))?,
299                            );
300                        }
301                        _ => return Err(invalid("bad escape", self.i - 1)),
302                    }
303                }
304                _ => return Err(invalid("unterminated string", start)),
305            }
306        }
307    }
308
309    fn hex4(&mut self) -> Result<u32> {
310        let digits = self
311            .s
312            .get(self.i..self.i + 4)
313            .ok_or_else(|| invalid("bad \\u escape", self.i))?;
314        // Exactly four hex digits: `from_str_radix` alone would accept a sign.
315        if !digits.iter().all(u8::is_ascii_hexdigit) {
316            return Err(invalid("bad \\u escape", self.i));
317        }
318        self.i += 4;
319        std::str::from_utf8(digits)
320            .ok()
321            .and_then(|d| u32::from_str_radix(d, 16).ok())
322            .ok_or_else(|| invalid("bad \\u escape", self.i - 4))
323    }
324}
325
326/// Parses one JSON value (trailing whitespace allowed).
327///
328/// Strict RFC 8259: lone surrogates, leading zeros and malformed numbers are errors, and
329/// arrays and objects nest at most 128 deep (deeper input is an error, not a stack
330/// overflow).
331pub fn parse(bytes: &[u8]) -> Result<Json> {
332    let mut p = Parser {
333        s: bytes,
334        i: 0,
335        depth: 0,
336    };
337    let v = p.value()?;
338    p.ws();
339    if p.i != bytes.len() {
340        return Err(invalid("trailing data", p.i));
341    }
342    Ok(v)
343}
344
345/// `s` as a JSON string literal.
346pub(crate) fn quote(s: &str) -> String {
347    let mut out = String::with_capacity(s.len() + 2);
348    out.push('"');
349    for c in s.chars() {
350        match c {
351            '"' => out.push_str("\\\""),
352            '\\' => out.push_str("\\\\"),
353            '\n' => out.push_str("\\n"),
354            '\r' => out.push_str("\\r"),
355            '\t' => out.push_str("\\t"),
356            c if (c as u32) < 0x20 => out.push_str(&format!("\\u{:04x}", c as u32)),
357            c => out.push(c),
358        }
359    }
360    out.push('"');
361    out
362}
363
364/// A string map as a JSON object.
365pub(crate) fn object(map: &BTreeMap<String, String>) -> String {
366    let items: Vec<String> = map
367        .iter()
368        .map(|(k, v)| format!("{}:{}", quote(k), quote(v)))
369        .collect();
370    format!("{{{}}}", items.join(","))
371}
nuvec_io/json.rs

nuvec_io/
json.rs
