Implement json.decoder.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,7 +220,7 @@ 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)
+        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() });
         // SAFETY: we only output valid utf8 from write_json_string
@@ -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

@@ -96,3 +96,133 @@ 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(s: &str, end: usize, strict: bool) -> Result<(String, usize), DecodeError> {
+    let mut chunks: Vec<StrOrChar> = Vec::new();
+    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 {
+            '"' => {
+                chunks.push(StrOrChar::Str(&s[chunk_start..i]));
+                let len = chunks.iter().map(|x| x.len()).sum();
+                let mut out = String::with_capacity(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));
+            }
+            '\\' => {
+                let (_, (_, c)) = match chars.next() {
+                    Some(next) => next,
+                    // unterminated string
+                    None => break,
+                };
+                let esc = match c {
+                    '"' => "\"",
+                    '\\' => "\\",
+                    '/' => "/",
+                    'b' => "\x08",
+                    'f' => "\x0c",
+                    'n' => "\n",
+                    'r' => "\r",
+                    't' => "\t",
+                    'u' => {
+                        let surrogate_err = || DecodeError::new("unpaired surrogate", char_i);
+                        chunks.push(StrOrChar::Str(&s[chunk_start..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());
+                                }
+                            }
+                        }
+                        chunks.push(StrOrChar::Char(
+                            std::char::from_u32(uni).ok_or_else(surrogate_err)?,
+                        ));
+                        continue;
+                    }
+                    _ => {
+                        return Err(DecodeError::new(
+                            format!("Invalid \\escape: {:?}", c),
+                            char_i,
+                        ))
+                    }
+                };
+                chunks.push(StrOrChar::Str(&s[chunk_start..i]));
+                chunk_start = i + 2;
+                chunks.push(StrOrChar::Str(esc));
+            }
+            '\x00'..='\x1f' if strict => {
+                return Err(DecodeError::new(
+                    format!("Invalid control character {:?} at", c),
+                    char_i,
+                ));
+            }
+            _ => {}
+        }
+    }
+    Err(DecodeError::new("Unterminated string starting at", end - 1))
+}
+
+#[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)
+}