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
90struct Parser<'a> {
91    s: &'a [u8],
92    i: usize,
93}
94
95impl Parser<'_> {
96    fn ws(&mut self) {
97        while self.s.get(self.i).is_some_and(u8::is_ascii_whitespace) {
98            self.i += 1;
99        }
100    }
101
102    fn eat(&mut self, c: u8) -> Result<()> {
103        self.ws();
104        if self.s.get(self.i) == Some(&c) {
105            self.i += 1;
106            Ok(())
107        } else {
108            Err(invalid(&format!("expected `{}`", c as char), self.i))
109        }
110    }
111
112    fn value(&mut self) -> Result<Json> {
113        self.ws();
114        match self.s.get(self.i) {
115            Some(b'{') => {
116                self.i += 1;
117                let mut items = Vec::new();
118                self.ws();
119                if self.s.get(self.i) == Some(&b'}') {
120                    self.i += 1;
121                    return Ok(Json::Object(items));
122                }
123                loop {
124                    self.ws();
125                    let key = self.string()?;
126                    self.eat(b':')?;
127                    items.push((key, self.value()?));
128                    self.ws();
129                    match self.s.get(self.i) {
130                        Some(b',') => self.i += 1,
131                        Some(b'}') => {
132                            self.i += 1;
133                            return Ok(Json::Object(items));
134                        }
135                        _ => return Err(invalid("expected `,` or `}`", self.i)),
136                    }
137                }
138            }
139            Some(b'[') => {
140                self.i += 1;
141                let mut items = Vec::new();
142                self.ws();
143                if self.s.get(self.i) == Some(&b']') {
144                    self.i += 1;
145                    return Ok(Json::Array(items));
146                }
147                loop {
148                    items.push(self.value()?);
149                    self.ws();
150                    match self.s.get(self.i) {
151                        Some(b',') => self.i += 1,
152                        Some(b']') => {
153                            self.i += 1;
154                            return Ok(Json::Array(items));
155                        }
156                        _ => return Err(invalid("expected `,` or `]`", self.i)),
157                    }
158                }
159            }
160            Some(b'"') => Ok(Json::Str(self.string()?)),
161            Some(b't') if self.s[self.i..].starts_with(b"true") => {
162                self.i += 4;
163                Ok(Json::Bool(true))
164            }
165            Some(b'f') if self.s[self.i..].starts_with(b"false") => {
166                self.i += 5;
167                Ok(Json::Bool(false))
168            }
169            Some(b'n') if self.s[self.i..].starts_with(b"null") => {
170                self.i += 4;
171                Ok(Json::Null)
172            }
173            Some(c) if *c == b'-' || c.is_ascii_digit() => {
174                let start = self.i;
175                self.i += 1;
176                while self
177                    .s
178                    .get(self.i)
179                    .is_some_and(|c| c.is_ascii_digit() || b".eE+-".contains(c))
180                {
181                    self.i += 1;
182                }
183                let text = std::str::from_utf8(&self.s[start..self.i]).unwrap_or("");
184                text.parse()
185                    .map(Json::Int)
186                    .or_else(|_| text.parse().map(Json::Float))
187                    .map_err(|_| invalid("number expected", start))
188            }
189            _ => Err(invalid("value expected", self.i)),
190        }
191    }
192
193    fn string(&mut self) -> Result<String> {
194        if self.s.get(self.i) != Some(&b'"') {
195            return Err(invalid("string expected", self.i));
196        }
197        self.i += 1;
198        let mut out = String::new();
199        loop {
200            let start = self.i;
201            while self.s.get(self.i).is_some_and(|&c| c != b'"' && c != b'\\') {
202                self.i += 1;
203            }
204            out.push_str(
205                std::str::from_utf8(&self.s[start..self.i])
206                    .map_err(|_| invalid("invalid UTF-8", start))?,
207            );
208            match self.s.get(self.i) {
209                Some(b'"') => {
210                    self.i += 1;
211                    return Ok(out);
212                }
213                Some(b'\\') => {
214                    let esc = *self
215                        .s
216                        .get(self.i + 1)
217                        .ok_or_else(|| invalid("bad escape", self.i))?;
218                    self.i += 2;
219                    match esc {
220                        b'"' => out.push('"'),
221                        b'\\' => out.push('\\'),
222                        b'/' => out.push('/'),
223                        b'b' => out.push('\u{8}'),
224                        b'f' => out.push('\u{c}'),
225                        b'n' => out.push('\n'),
226                        b'r' => out.push('\r'),
227                        b't' => out.push('\t'),
228                        b'u' => {
229                            let mut code = self.hex4()?;
230                            if (0xD800..0xDC00).contains(&code)
231                                && self.s[self.i..].starts_with(b"\\u")
232                            {
233                                self.i += 2;
234                                let low = self.hex4()?;
235                                code = 0x10000
236                                    + ((code - 0xD800) << 10)
237                                    + (low.wrapping_sub(0xDC00) & 0x3FF);
238                            }
239                            out.push(char::from_u32(code).unwrap_or('\u{FFFD}'));
240                        }
241                        _ => return Err(invalid("bad escape", self.i - 1)),
242                    }
243                }
244                _ => return Err(invalid("unterminated string", start)),
245            }
246        }
247    }
248
249    fn hex4(&mut self) -> Result<u32> {
250        let digits = self
251            .s
252            .get(self.i..self.i + 4)
253            .ok_or_else(|| invalid("bad \\u escape", self.i))?;
254        self.i += 4;
255        std::str::from_utf8(digits)
256            .ok()
257            .and_then(|d| u32::from_str_radix(d, 16).ok())
258            .ok_or_else(|| invalid("bad \\u escape", self.i - 4))
259    }
260}
261
262/// Parses one JSON value (trailing whitespace allowed).
263pub fn parse(bytes: &[u8]) -> Result<Json> {
264    let mut p = Parser { s: bytes, i: 0 };
265    let v = p.value()?;
266    p.ws();
267    if p.i != bytes.len() {
268        return Err(invalid("trailing data", p.i));
269    }
270    Ok(v)
271}
272
273/// `s` as a JSON string literal.
274pub(crate) fn quote(s: &str) -> String {
275    let mut out = String::with_capacity(s.len() + 2);
276    out.push('"');
277    for c in s.chars() {
278        match c {
279            '"' => out.push_str("\\\""),
280            '\\' => out.push_str("\\\\"),
281            '\n' => out.push_str("\\n"),
282            '\r' => out.push_str("\\r"),
283            '\t' => out.push_str("\\t"),
284            c if (c as u32) < 0x20 => out.push_str(&format!("\\u{:04x}", c as u32)),
285            c => out.push(c),
286        }
287    }
288    out.push('"');
289    out
290}
291
292/// A string map as a JSON object.
293pub(crate) fn object(map: &BTreeMap<String, String>) -> String {
294    let items: Vec<String> = map
295        .iter()
296        .map(|(k, v)| format!("{}:{}", quote(k), quote(v)))
297        .collect();
298    format!("{{{}}}", items.join(","))
299}
nuvec_io/json.rs

nuvec_io/
json.rs
