Merge pull request #2265 from RustPython/coolreader18/json-scanstring

Implement _json.scanstring in Rust

Lib/test/test_json/test_speedups.py

@@ -8,8 +8,6 @@ class BadBool:
 
 
 class TestSpeedups(CTest):
-    # TODO: RUSTPYTHON
-    @unittest.expectedFailure
     def test_scanstring(self):
         self.assertEqual(self.json.decoder.scanstring.__module__, "_json")
         self.assertIs(self.json.decoder.scanstring, self.json.decoder.c_scanstring)

vm/src/stdlib/json.rs

@@ -3,11 +3,15 @@ mod machinery;
 
 #[pymodule]
 mod _json {
-    use crate::function::PyFuncArgs;
+    use super::*;
+    use crate::exceptions::PyBaseExceptionRef;
+    use crate::function::{OptionalArg, PyFuncArgs};
     use crate::obj::objiter;
     use crate::obj::objstr::PyStrRef;
     use crate::obj::{objbool, objtype::PyTypeRef};
-    use crate::pyobject::{BorrowValue, IdProtocol, PyObjectRef, PyRef, PyResult, PyValue};
+    use crate::pyobject::{
+        BorrowValue, IdProtocol, IntoPyObject, PyObjectRef, PyRef, PyResult, PyValue, TryFromObject,
+    };
     use crate::slots::Callable;
     use crate::VirtualMachine;
 
@@ -83,16 +87,8 @@ mod _json {
             let next_idx = idx + c.len_utf8();
             match c {
                 '"' => {
-                    // TODO: parse the string in rust
-                    let parse_str = vm.get_attribute(self.ctx.clone(), "parse_string")?;
-                    return vm.invoke(
-                        &parse_str,
-                        vec![
-                            pystr.into_object(),
-                            vm.ctx.new_int(next_idx),
-                            vm.ctx.new_bool(self.strict),
-                        ],
-                    );
+                    return scanstring(pystr, next_idx, OptionalArg::Present(self.strict), vm)
+                        .map(|x| x.into_pyobject(vm))
                 }
                 '{' => {
                     // TODO: parse the object in rust
@@ -224,11 +220,11 @@ mod _json {
 
     fn encode_string(s: &str, ascii_only: bool) -> String {
         let mut buf = Vec::<u8>::with_capacity(s.len() + 2);
-        super::machinery::write_json_string(s, ascii_only, &mut buf)
-            // writing to a vec can't fail
+        machinery::write_json_string(s, ascii_only, &mut buf)
+            // SAFETY: writing to a vec can't fail
             .unwrap_or_else(|_| unsafe { std::hint::unreachable_unchecked() });
-        // TODO: verify that the implementation is correct enough to use `from_utf8_unchecked`
-        String::from_utf8(buf).expect("invalid utf-8 in json output")
+        // SAFETY: we only output valid utf8 from write_json_string
+        unsafe { String::from_utf8_unchecked(buf) }
     }
 
     #[pyfunction]
@@ -240,4 +236,37 @@ mod _json {
     fn encode_basestring_ascii(s: PyStrRef) -> String {
         encode_string(s.borrow_value(), true)
     }
+
+    fn py_decode_error(
+        e: machinery::DecodeError,
+        s: PyStrRef,
+        vm: &VirtualMachine,
+    ) -> PyBaseExceptionRef {
+        let get_error = || -> PyResult<_> {
+            let cls = vm.try_class("json", "JSONDecodeError")?;
+            let exc = vm.invoke(
+                cls.as_object(),
+                vec![
+                    vm.ctx.new_str(e.msg),
+                    s.into_object(),
+                    vm.ctx.new_int(e.pos),
+                ],
+            )?;
+            PyBaseExceptionRef::try_from_object(vm, exc)
+        };
+        match get_error() {
+            Ok(x) | Err(x) => x,
+        }
+    }
+
+    #[pyfunction]
+    fn scanstring(
+        s: PyStrRef,
+        end: usize,
+        strict: OptionalArg<bool>,
+        vm: &VirtualMachine,
+    ) -> PyResult<(String, usize)> {
+        machinery::scanstring(s.borrow_value(), end, strict.unwrap_or(true))
+            .map_err(|e| py_decode_error(e, s, vm))
+    }
 }

vm/src/stdlib/json/machinery.rs

@@ -28,11 +28,11 @@
 
 use std::io;
 
-static ESCAPE_CHARS: [&[u8]; 0x20] = [
-    b"\\u0000", b"\\u0001", b"\\u0002", b"\\u0003", b"\\u0004", b"\\u0005", b"\\u0006", b"\\u0007",
-    b"\\b", b"\\t", b"\\n", b"\\u000b", b"\\f", b"\\r", b"\\u000e", b"\\u000f", b"\\u0010",
-    b"\\u0011", b"\\u0012", b"\\u0013", b"\\u0014", b"\\u0015", b"\\u0016", b"\\u0017", b"\\u0018",
-    b"\\u0019", b"\\u001a", b"\\u001b", b"\\u001c", b"\\u001d", b"\\u001e", b"\\u001f",
+static ESCAPE_CHARS: [&str; 0x20] = [
+    "\\u0000", "\\u0001", "\\u0002", "\\u0003", "\\u0004", "\\u0005", "\\u0006", "\\u0007", "\\b",
+    "\\t", "\\n", "\\u000", "\\f", "\\r", "\\u000e", "\\u000f", "\\u0010", "\\u0011", "\\u0012",
+    "\\u0013", "\\u0014", "\\u0015", "\\u0016", "\\u0017", "\\u0018", "\\u0019", "\\u001a",
+    "\\u001", "\\u001c", "\\u001d", "\\u001e", "\\u001f",
 ];
 
 // This bitset represents which bytes can be copied as-is to a JSON string (0)
@@ -48,16 +48,16 @@ static NEEDS_ESCAPING_BITSET: [u64; 4] = [
 ];
 
 #[inline(always)]
-fn json_escaped_char(c: u8) -> Option<&'static [u8]> {
+fn json_escaped_char(c: u8) -> Option<&'static str> {
     let bitset_value = NEEDS_ESCAPING_BITSET[(c / 64) as usize] & (1 << (c % 64));
     if bitset_value == 0 {
         None
     } else {
         Some(match c {
             x if x < 0x20 => ESCAPE_CHARS[c as usize],
-            b'\\' => &b"\\\\"[..],
-            b'\"' => &b"\\\""[..],
-            0x7F => &b"\\u007f"[..],
+            b'\\' => "\\\\",
+            b'\"' => "\\\"",
+            0x7F => "\\u007f",
             _ => unreachable!(),
         })
     }
@@ -72,7 +72,7 @@ pub fn write_json_string<W: io::Write>(s: &str, ascii_only: bool, w: &mut W) ->
             if c.is_ascii() {
                 if let Some(escaped) = json_escaped_char(c as u8) {
                     w.write_all(&bytes[write_start_idx..idx])?;
-                    w.write_all(escaped)?;
+                    w.write_all(escaped.as_bytes())?;
                     write_start_idx = idx + 1;
                 }
             } else {
@@ -88,7 +88,7 @@ pub fn write_json_string<W: io::Write>(s: &str, ascii_only: bool, w: &mut W) ->
         for (idx, c) in s.bytes().enumerate() {
             if let Some(escaped) = json_escaped_char(c) {
                 w.write_all(&bytes[write_start_idx..idx])?;
-                w.write_all(escaped)?;
+                w.write_all(escaped.as_bytes())?;
                 write_start_idx = idx + 1;
             }
         }
@@ -96,3 +96,137 @@ pub fn write_json_string<W: io::Write>(s: &str, ascii_only: bool, w: &mut W) ->
     w.write_all(&bytes[write_start_idx..])?;
     w.write_all(b"\"")
 }
+
+#[derive(Debug)]
+pub struct DecodeError {
+    pub msg: String,
+    pub pos: usize,
+}
+impl DecodeError {
+    fn new(msg: impl Into<String>, pos: usize) -> Self {
+        Self {
+            msg: msg.into(),
+            pos,
+        }
+    }
+}
+
+enum StrOrChar<'a> {
+    Str(&'a str),
+    Char(char),
+}
+impl StrOrChar<'_> {
+    fn len(&self) -> usize {
+        match self {
+            StrOrChar::Str(s) => s.len(),
+            StrOrChar::Char(c) => c.len_utf8(),
+        }
+    }
+}
+pub fn scanstring<'a>(
+    s: &'a str,
+    end: usize,
+    strict: bool,
+) -> Result<(String, usize), DecodeError> {
+    let mut chunks: Vec<StrOrChar<'a>> = Vec::new();
+    let mut output_len = 0usize;
+    let mut push_chunk = |chunk: StrOrChar<'a>| {
+        output_len += chunk.len();
+        chunks.push(chunk);
+    };
+    let unterminated_err = || DecodeError::new("Unterminated string starting at", end - 1);
+    let mut chunk_start = end;
+    let mut chars = s.char_indices().enumerate().skip(end).peekable();
+    while let Some((char_i, (i, c))) = chars.next() {
+        match c {
+            '"' => {
+                push_chunk(StrOrChar::Str(&s[chunk_start..i]));
+                let mut out = String::with_capacity(output_len);
+                for x in chunks {
+                    match x {
+                        StrOrChar::Str(s) => out.push_str(s),
+                        StrOrChar::Char(c) => out.push(c),
+                    }
+                }
+                return Ok((out, char_i + 1));
+            }
+            '\\' => {
+                push_chunk(StrOrChar::Str(&s[chunk_start..i]));
+                let (_, (_, c)) = chars.next().ok_or_else(unterminated_err)?;
+                let esc = match c {
+                    '"' => "\"",
+                    '\\' => "\\",
+                    '/' => "/",
+                    'b' => "\x08",
+                    'f' => "\x0c",
+                    'n' => "\n",
+                    'r' => "\r",
+                    't' => "\t",
+                    'u' => {
+                        let surrogate_err = || DecodeError::new("unpaired surrogate", char_i);
+                        let mut uni = decode_unicode(&mut chars, char_i)?;
+                        chunk_start = char_i + 6;
+                        if (0xd800..=0xdbff).contains(&uni) {
+                            // uni is a surrogate -- try to find its pair
+                            if let Some(&(pos2, (_, '\\'))) = chars.peek() {
+                                // ok, the next char starts an escape
+                                chars.next();
+                                if let Some((_, (_, 'u'))) = chars.peek() {
+                                    // ok, it's a unicode escape
+                                    chars.next();
+                                    let uni2 = decode_unicode(&mut chars, pos2)?;
+                                    chunk_start = pos2 + 6;
+                                    if (0xdc00..=0xdfff).contains(&uni2) {
+                                        // ok, we found what we were looking for -- \uXXXX\uXXXX, both surrogates
+                                        uni = 0x10000 + (((uni - 0xd800) << 10) | (uni2 - 0xdc00));
+                                    } else {
+                                        // if we don't find a matching surrogate, error -- until str
+                                        // isn't utf8 internally, we can't parse surrogates
+                                        return Err(surrogate_err());
+                                    }
+                                } else {
+                                    return Err(surrogate_err());
+                                }
+                            }
+                        }
+                        push_chunk(StrOrChar::Char(
+                            std::char::from_u32(uni).ok_or_else(surrogate_err)?,
+                        ));
+                        continue;
+                    }
+                    _ => {
+                        return Err(DecodeError::new(
+                            format!("Invalid \\escape: {:?}", c),
+                            char_i,
+                        ))
+                    }
+                };
+                chunk_start = i + 2;
+                push_chunk(StrOrChar::Str(esc));
+            }
+            '\x00'..='\x1f' if strict => {
+                return Err(DecodeError::new(
+                    format!("Invalid control character {:?} at", c),
+                    char_i,
+                ));
+            }
+            _ => {}
+        }
+    }
+    Err(unterminated_err())
+}
+
+#[inline]
+fn decode_unicode<I>(it: &mut I, pos: usize) -> Result<u32, DecodeError>
+where
+    I: Iterator<Item = (usize, (usize, char))>,
+{
+    let err = || DecodeError::new("Invalid \\uXXXX escape", pos);
+    let mut uni = 0;
+    for x in (0..4).rev() {
+        let (_, (_, c)) = it.next().ok_or_else(err)?;
+        let d = c.to_digit(16).ok_or_else(err)?;
+        uni += d * 16u32.pow(x);
+    }
+    Ok(uni)
+}

vm/src/stdlib/serde_json.rs

@@ -3,6 +3,7 @@ pub(crate) use _serde_json::make_module;
 #[pymodule]
 mod _serde_json {
     use crate::common::borrow::BorrowValue;
+    use crate::exceptions::PyBaseExceptionRef;
     use crate::obj::objstr::PyStrRef;
     use crate::py_serde;
     use crate::pyobject::{PyResult, TryFromObject};
@@ -17,22 +18,30 @@ mod _serde_json {
             Ok(res)
         })();
 
-        res.or_else(|err| {
-            let decode_error = vm.try_class("json", "JSONDecodeError")?;
-            let from_serde = vm.get_attribute(decode_error.into_object(), "_from_serde")?;
-            let mut err_msg = err.to_string();
-            let pos = err_msg.rfind(" at line ").unwrap();
-            err_msg.truncate(pos);
-            let decode_error = vm.invoke(
-                &from_serde,
-                vec![
-                    vm.ctx.new_str(err_msg),
-                    s.into_object(),
-                    vm.ctx.new_int(err.line()),
-                    vm.ctx.new_int(err.column()),
-                ],
-            )?;
-            TryFromObject::try_from_object(vm, decode_error)
+        res.map_err(|err| match json_exception(err, s, vm) {
+            Ok(x) | Err(x) => x,
         })
     }
+
+    fn json_exception(
+        err: serde_json::Error,
+        s: PyStrRef,
+        vm: &VirtualMachine,
+    ) -> PyResult<PyBaseExceptionRef> {
+        let decode_error = vm.try_class("json", "JSONDecodeError")?;
+        let from_serde = vm.get_attribute(decode_error.into_object(), "_from_serde")?;
+        let mut err_msg = err.to_string();
+        let pos = err_msg.rfind(" at line ").unwrap();
+        err_msg.truncate(pos);
+        let decode_error = vm.invoke(
+            &from_serde,
+            vec![
+                vm.ctx.new_str(err_msg),
+                s.into_object(),
+                vm.ctx.new_int(err.line()),
+                vm.ctx.new_int(err.column()),
+            ],
+        )?;
+        PyBaseExceptionRef::try_from_object(vm, decode_error)
+    }
 }