replace cached_py_string and pystring_fast_new with safe alternatives (#201)

davidhewitt · web-flow · commit 9abadb78d82e · 2025-05-18T13:29:51.000-04:00
diff --git a/crates/jiter/src/lib.rs b/crates/jiter/src/lib.rs
@@ -172,7 +172,9 @@ pub use value::{JsonArray, JsonObject, JsonValue};
 #[cfg(feature = "python")]
 pub use py_lossless_float::{FloatMode, LosslessFloat};
 #[cfg(feature = "python")]
-pub use py_string_cache::{cache_clear, cache_usage, cached_py_string, pystring_fast_new, StringCacheMode};
+pub use py_string_cache::{
+    cache_clear, cache_usage, cached_py_string, cached_py_string_ascii, pystring_ascii_new, StringCacheMode,
+};
 #[cfg(feature = "python")]
 pub use python::{map_json_error, PythonParse};
 
diff --git a/crates/jiter/src/py_string_cache.rs b/crates/jiter/src/py_string_cache.rs
@@ -5,6 +5,8 @@ use pyo3::exceptions::{PyTypeError, PyValueError};
 use pyo3::prelude::*;
 use pyo3::types::{PyBool, PyString};
 
+use crate::string_decoder::StringOutput;
+
 #[derive(Debug, Clone, Copy)]
 pub enum StringCacheMode {
     All,
@@ -50,38 +52,40 @@ impl From<bool> for StringCacheMode {
 }
 
 pub trait StringMaybeCache {
-    fn get_key<'py>(py: Python<'py>, json_str: &str, ascii_only: bool) -> Bound<'py, PyString>;
+    fn get_key<'py>(py: Python<'py>, string_output: StringOutput<'_, '_>) -> Bound<'py, PyString>;
 
-    fn get_value<'py>(py: Python<'py>, json_str: &str, ascii_only: bool) -> Bound<'py, PyString> {
-        Self::get_key(py, json_str, ascii_only)
+    fn get_value<'py>(py: Python<'py>, string_output: StringOutput<'_, '_>) -> Bound<'py, PyString> {
+        Self::get_key(py, string_output)
     }
 }
 
 pub struct StringCacheAll;
 
 impl StringMaybeCache for StringCacheAll {
-    fn get_key<'py>(py: Python<'py>, json_str: &str, ascii_only: bool) -> Bound<'py, PyString> {
-        cached_py_string(py, json_str, ascii_only)
+    fn get_key<'py>(py: Python<'py>, string_output: StringOutput<'_, '_>) -> Bound<'py, PyString> {
+        // Safety: string_output carries the safety information
+        unsafe { cached_py_string_maybe_ascii(py, string_output.as_str(), string_output.ascii_only()) }
     }
 }
 
 pub struct StringCacheKeys;
 
 impl StringMaybeCache for StringCacheKeys {
-    fn get_key<'py>(py: Python<'py>, json_str: &str, ascii_only: bool) -> Bound<'py, PyString> {
-        cached_py_string(py, json_str, ascii_only)
+    fn get_key<'py>(py: Python<'py>, string_output: StringOutput<'_, '_>) -> Bound<'py, PyString> {
+        // Safety: string_output carries the safety information
+        unsafe { cached_py_string_maybe_ascii(py, string_output.as_str(), string_output.ascii_only()) }
     }
 
-    fn get_value<'py>(py: Python<'py>, json_str: &str, ascii_only: bool) -> Bound<'py, PyString> {
-        pystring_fast_new(py, json_str, ascii_only)
+    fn get_value<'py>(py: Python<'py>, string_output: StringOutput<'_, '_>) -> Bound<'py, PyString> {
+        unsafe { pystring_fast_new_maybe_ascii(py, string_output.as_str(), string_output.ascii_only()) }
     }
 }
 
 pub struct StringNoCache;
 
 impl StringMaybeCache for StringNoCache {
-    fn get_key<'py>(py: Python<'py>, json_str: &str, ascii_only: bool) -> Bound<'py, PyString> {
-        pystring_fast_new(py, json_str, ascii_only)
+    fn get_key<'py>(py: Python<'py>, string_output: StringOutput<'_, '_>) -> Bound<'py, PyString> {
+        unsafe { pystring_fast_new_maybe_ascii(py, string_output.as_str(), string_output.ascii_only()) }
     }
 }
 
@@ -108,12 +112,33 @@ pub fn cache_clear() {
     get_string_cache().clear();
 }
 
-pub fn cached_py_string<'py>(py: Python<'py>, s: &str, ascii_only: bool) -> Bound<'py, PyString> {
+/// Create a cached Python `str` from a string slice
+#[inline]
+pub fn cached_py_string<'py>(py: Python<'py>, s: &str) -> Bound<'py, PyString> {
+    // SAFETY: not setting ascii-only
+    unsafe { cached_py_string_maybe_ascii(py, s, false) }
+}
+
+/// Create a cached Python `str` from a string slice.
+///
+/// # Safety
+///
+/// Caller must pass ascii-only string.
+#[inline]
+pub unsafe fn cached_py_string_ascii<'py>(py: Python<'py>, s: &str) -> Bound<'py, PyString> {
+    // SAFETY: caller upholds invariant
+    unsafe { cached_py_string_maybe_ascii(py, s, true) }
+}
+
+/// # Safety
+///
+/// Caller must match the ascii_only flag to the string passed in.
+unsafe fn cached_py_string_maybe_ascii<'py>(py: Python<'py>, s: &str, ascii_only: bool) -> Bound<'py, PyString> {
     // from tests, 0 and 1 character strings are faster not cached
     if (2..64).contains(&s.len()) {
         get_string_cache().get_or_insert(py, s, ascii_only)
     } else {
-        pystring_fast_new(py, s, ascii_only)
+        pystring_fast_new_maybe_ascii(py, s, ascii_only)
     }
 }
 
@@ -146,13 +171,18 @@ impl Default for PyStringCache {
 impl PyStringCache {
     /// Lookup the cache for an entry with the given string. If it exists, return it.
     /// If it is not set or has a different string, insert it and return it.
-    fn get_or_insert<'py>(&mut self, py: Python<'py>, s: &str, ascii_only: bool) -> Bound<'py, PyString> {
+    ///
+    /// # Safety
+    ///
+    /// `ascii_only` must only be set to `true` if the string is guaranteed to be ASCII only.
+    unsafe fn get_or_insert<'py>(&mut self, py: Python<'py>, s: &str, ascii_only: bool) -> Bound<'py, PyString> {
         let hash = self.hash_builder.hash_one(s);
 
         let hash_index = hash as usize % CAPACITY;
 
         let set_entry = |entry: &mut Entry| {
-            let py_str = pystring_fast_new(py, s, ascii_only);
+            // SAFETY: caller upholds invariant
+            let py_str = unsafe { pystring_fast_new_maybe_ascii(py, s, ascii_only) };
             if let Some((_, old_py_str)) = entry.replace((hash, py_str.clone().unbind())) {
                 // micro-optimization: bind the old entry before dropping it so that PyO3 can
                 // fast-path the drop (Bound::drop is faster than Py::drop)
@@ -199,8 +229,14 @@ impl PyStringCache {
     }
 }
 
-pub fn pystring_fast_new<'py>(py: Python<'py>, s: &str, ascii_only: bool) -> Bound<'py, PyString> {
+/// Creatate a new Python `str` from a string slice, with a fast path for ASCII strings
+///
+/// # Safety
+///
+/// `ascii_only` must only be set to `true` if the string is guaranteed to be ASCII only.
+unsafe fn pystring_fast_new_maybe_ascii<'py>(py: Python<'py>, s: &str, ascii_only: bool) -> Bound<'py, PyString> {
     if ascii_only {
+        // SAFETY: caller upholds invariant
         unsafe { pystring_ascii_new(py, s) }
     } else {
         PyString::new(py, s)
@@ -209,22 +245,24 @@ pub fn pystring_fast_new<'py>(py: Python<'py>, s: &str, ascii_only: bool) -> Bou
 
 /// Faster creation of PyString from an ASCII string, inspired by
 /// https://github.com/ijl/orjson/blob/3.10.0/src/str/create.rs#L41
-#[cfg(not(any(PyPy, GraalPy)))]
-unsafe fn pystring_ascii_new<'py>(py: Python<'py>, s: &str) -> Bound<'py, PyString> {
-    // disabled on everything except tier-1 platforms because of a crash in the built wheels from CI,
-    // see https://github.com/pydantic/jiter/pull/175
-
-    let ptr = pyo3::ffi::PyUnicode_New(s.len() as isize, 127);
-    // see https://github.com/pydantic/jiter/pull/72#discussion_r1545485907
-    debug_assert_eq!(pyo3::ffi::PyUnicode_KIND(ptr), pyo3::ffi::PyUnicode_1BYTE_KIND);
-    let data_ptr = pyo3::ffi::PyUnicode_DATA(ptr).cast();
-    core::ptr::copy_nonoverlapping(s.as_ptr(), data_ptr, s.len());
-    core::ptr::write(data_ptr.add(s.len()), 0);
-    Bound::from_owned_ptr(py, ptr).downcast_into_unchecked()
-}
-
-// unoptimized version (albeit not that much slower) on other platforms
-#[cfg(any(PyPy, GraalPy))]
-unsafe fn pystring_ascii_new<'py>(py: Python<'py>, s: &str) -> Bound<'py, PyString> {
-    PyString::new(py, s)
+///
+/// # Safety
+///
+/// `s` must be ASCII only
+pub unsafe fn pystring_ascii_new<'py>(py: Python<'py>, s: &str) -> Bound<'py, PyString> {
+    #[cfg(not(any(PyPy, GraalPy, Py_LIMITED_API)))]
+    {
+        let ptr = pyo3::ffi::PyUnicode_New(s.len() as isize, 127);
+        // see https://github.com/pydantic/jiter/pull/72#discussion_r1545485907
+        debug_assert_eq!(pyo3::ffi::PyUnicode_KIND(ptr), pyo3::ffi::PyUnicode_1BYTE_KIND);
+        let data_ptr = pyo3::ffi::PyUnicode_DATA(ptr).cast();
+        core::ptr::copy_nonoverlapping(s.as_ptr(), data_ptr, s.len());
+        core::ptr::write(data_ptr.add(s.len()), 0);
+        Bound::from_owned_ptr(py, ptr).downcast_into_unchecked()
+    }
+
+    #[cfg(any(PyPy, GraalPy, Py_LIMITED_API))]
+    {
+        PyString::new(py, s)
+    }
 }
diff --git a/crates/jiter/src/python.rs b/crates/jiter/src/python.rs
@@ -138,7 +138,7 @@ impl<StringCache: StringMaybeCache, KeyCheck: MaybeKeyCheck, ParseNumber: MaybeP
                 let s = self
                     .parser
                     .consume_string::<StringDecoder>(&mut self.tape, self.partial_mode.allow_trailing_str())?;
-                Ok(StringCache::get_value(py, s.as_str(), s.ascii_only()).into_any())
+                Ok(StringCache::get_value(py, s).into_any())
             }
             Peek::Array => {
                 let peek_first = match self.parser.array_first() {
@@ -198,14 +198,14 @@ impl<StringCache: StringMaybeCache, KeyCheck: MaybeKeyCheck, ParseNumber: MaybeP
         if let Some(first_key) = self.parser.object_first::<StringDecoder>(&mut self.tape)? {
             let first_key_s = first_key.as_str();
             check_keys.check(first_key_s, self.parser.index)?;
-            let first_key = StringCache::get_key(py, first_key_s, first_key.ascii_only());
+            let first_key = StringCache::get_key(py, first_key);
             let peek = self.parser.peek()?;
             let first_value = self.check_take_value(py, peek)?;
             set_item(first_key, first_value);
             while let Some(key) = self.parser.object_step::<StringDecoder>(&mut self.tape)? {
                 let key_s = key.as_str();
                 check_keys.check(key_s, self.parser.index)?;
-                let key = StringCache::get_key(py, key_s, key.ascii_only());
+                let key = StringCache::get_key(py, key);
                 let peek = self.parser.peek()?;
                 let value = self.check_take_value(py, peek)?;
                 set_item(key, value);
diff --git a/crates/jiter/src/string_decoder.rs b/crates/jiter/src/string_decoder.rs
@@ -1,4 +1,3 @@
-use std::borrow::Cow;
 use std::ops::Range;
 use std::str::{from_utf8, from_utf8_unchecked};
 
@@ -26,48 +25,87 @@ where
 pub struct StringDecoder;
 
 #[derive(Debug)]
-pub enum StringOutput<'t, 'j>
+pub enum StringOutputType<'t, 'j>
 where
     'j: 't,
 {
-    Tape(&'t str, bool),
-    Data(&'j str, bool),
+    Tape(&'t str),
+    Data(&'j str),
 }
 
-impl From<StringOutput<'_, '_>> for String {
-    fn from(val: StringOutput) -> Self {
-        match val {
-            StringOutput::Tape(s, _) => s.to_owned(),
-            StringOutput::Data(s, _) => s.to_owned(),
-        }
+/// This submodule is used to create a safety boundary where the `ascii_only`
+/// flag can be used to carry soundness information about the string.
+mod string_output {
+    use std::borrow::Cow;
+
+    use super::StringOutputType;
+
+    #[derive(Debug)]
+    pub struct StringOutput<'t, 'j>
+    where
+        'j: 't,
+    {
+        data: StringOutputType<'t, 'j>,
+        // SAFETY: this is used as an invariant to determine if the string is ascii only
+        // so this should not be set except when known
+        ascii_only: bool,
     }
-}
 
-impl<'j> From<StringOutput<'_, 'j>> for Cow<'j, str> {
-    fn from(val: StringOutput<'_, 'j>) -> Self {
-        match val {
-            StringOutput::Tape(s, _) => s.to_owned().into(),
-            StringOutput::Data(s, _) => s.into(),
+    impl From<StringOutput<'_, '_>> for String {
+        fn from(val: StringOutput) -> Self {
+            match val.data {
+                StringOutputType::Tape(s) | StringOutputType::Data(s) => s.to_owned(),
+            }
         }
     }
-}
 
-impl<'t> StringOutput<'t, '_> {
-    pub fn as_str(&self) -> &'t str {
-        match self {
-            Self::Tape(s, _) => s,
-            Self::Data(s, _) => s,
+    impl<'j> From<StringOutput<'_, 'j>> for Cow<'j, str> {
+        fn from(val: StringOutput<'_, 'j>) -> Self {
+            match val.data {
+                StringOutputType::Tape(s) => s.to_owned().into(),
+                StringOutputType::Data(s) => s.into(),
+            }
         }
     }
 
-    pub fn ascii_only(&self) -> bool {
-        match self {
-            Self::Tape(_, ascii_only) => *ascii_only,
-            Self::Data(_, ascii_only) => *ascii_only,
+    impl<'t, 'j> StringOutput<'t, 'j>
+    where
+        'j: 't,
+    {
+        /// # Safety
+        ///
+        /// `accii_only` must only be set to true if the string is ascii only
+        pub unsafe fn tape(data: &'t str, ascii_only: bool) -> Self {
+            StringOutput {
+                data: StringOutputType::Tape(data),
+                ascii_only,
+            }
+        }
+
+        /// # Safety
+        ///
+        /// `accii_only` must only be set to true if the string is ascii only
+        pub unsafe fn data(data: &'j str, ascii_only: bool) -> Self {
+            StringOutput {
+                data: StringOutputType::Data(data),
+                ascii_only,
+            }
+        }
+
+        pub fn as_str(&self) -> &'t str {
+            match self.data {
+                StringOutputType::Tape(s) | StringOutputType::Data(s) => s,
+            }
+        }
+
+        pub fn ascii_only(&self) -> bool {
+            self.ascii_only
         }
     }
 }
 
+pub use string_output::StringOutput;
+
 impl<'t, 'j> AbstractStringDecoder<'t, 'j> for StringDecoder
 where
     'j: 't,
@@ -85,7 +123,7 @@ where
         match decode_chunk(data, start, true, allow_partial)? {
             (StringChunk::StringEnd, ascii_only, index) => {
                 let s = to_str(&data[start..index], ascii_only, start)?;
-                Ok((StringOutput::Data(s, ascii_only), index + 1))
+                Ok((unsafe { StringOutput::data(s, ascii_only) }, index + 1))
             }
             (StringChunk::Backslash, ascii_only, index) => {
                 decode_to_tape(data, index, tape, start, ascii_only, allow_partial)
@@ -134,7 +172,7 @@ fn decode_to_tape<'t, 'j>(
                 tape.extend_from_slice(&data[index..new_index]);
                 index = new_index + 1;
                 let s = to_str(tape, ascii_only, start)?;
-                return Ok((StringOutput::Tape(s, ascii_only), index));
+                return Ok((unsafe { StringOutput::tape(s, ascii_only) }, index));
             }
             (StringChunk::Backslash, ascii_only_new, index_new) => {
                 ascii_only = ascii_only_new;
diff --git a/crates/jiter/tests/python.rs b/crates/jiter/tests/python.rs
@@ -1,7 +1,7 @@
 use pyo3::prelude::*;
 use pyo3::types::PyString;
 
-use jiter::{pystring_fast_new, JsonValue, PythonParse, StringCacheMode};
+use jiter::{pystring_ascii_new, JsonValue, PythonParse, StringCacheMode};
 
 #[cfg(feature = "num-bigint")]
 #[test]
@@ -71,19 +71,10 @@ fn test_cache_into() {
 }
 
 #[test]
-fn test_pystring_fast_new_non_ascii() {
-    let json = "£100 💩";
-    Python::with_gil(|py| {
-        let s = pystring_fast_new(py, json, false);
-        assert_eq!(s.to_string(), "£100 💩");
-    });
-}
-
-#[test]
-fn test_pystring_fast_new_ascii() {
+fn test_pystring_ascii_new() {
     let json = "100abc";
     Python::with_gil(|py| {
-        let s = pystring_fast_new(py, json, true);
+        let s = unsafe { pystring_ascii_new(py, json) };
         assert_eq!(s.to_string(), "100abc");
     });
 }
