Make re wtf8-compatible

RustPython · coolreader18 · Mar 26, 2025 · Mar 7, 2025 · Mar 21, 2025 · Mar 22, 2025
commit b36b32bfe810ab96cf945ff6ec04f85fcf46dd03
diff --git a/Cargo.lock b/Cargo.lock
diff --git a/Lib/test/test_re.py b/Lib/test/test_re.py
@@ -854,8 +854,6 @@ def test_string_boundaries(self):
        # Can match around the whitespace.
        self.assertEqual(len(re.findall(r"\B", " ")), 2)

-    # TODO: RUSTPYTHON
-    @unittest.expectedFailure
    def test_bigcharset(self):
        self.assertEqual(re.match("([\u2222\u2223])",
                                  "\u2222").group(1), "\u2222")

diff --git a/Lib/test/test_smtplib.py b/Lib/test/test_smtplib.py
@@ -1459,8 +1459,6 @@ def test_send_unicode_with_SMTPUTF8_via_low_level_API(self):
        self.assertIn('SMTPUTF8', self.serv.last_mail_options)
        self.assertEqual(self.serv.last_rcpt_options, [])

-    # TODO: RUSTPYTHON
-    @unittest.expectedFailure
    def test_send_message_uses_smtputf8_if_addrs_non_ascii(self):
        msg = EmailMessage()
        msg['From'] = "Páolo <főo@bar.com>"

diff --git a/vm/src/stdlib/sre.rs b/vm/src/stdlib/sre.rs
@@ -9,6 +9,7 @@ mod _sre {
            PyCallableIterator, PyDictRef, PyGenericAlias, PyInt, PyList, PyListRef, PyStr,
            PyStrRef, PyTuple, PyTupleRef, PyTypeRef,
        },
+        common::wtf8::{Wtf8, Wtf8Buf},
        common::{ascii, hash::PyHash},
        convert::ToPyObject,
        function::{ArgCallable, OptionalArg, PosArgs, PyComparisonValue},
@@ -66,10 +67,15 @@ mod _sre {
        }
    }

-    impl SreStr for &str {
+    impl SreStr for &Wtf8 {
        fn slice(&self, start: usize, end: usize, vm: &VirtualMachine) -> PyObjectRef {
            vm.ctx
-                .new_str(self.chars().take(end).skip(start).collect::<String>())
+                .new_str(
+                    self.code_points()
+                        .take(end)
+                        .skip(start)
+                        .collect::<Wtf8Buf>(),
+                )
                .into()
        }
    }
@@ -206,12 +212,12 @@ mod _sre {
    impl Pattern {
        fn with_str<F, R>(string: &PyObject, vm: &VirtualMachine, f: F) -> PyResult<R>
        where
-            F: FnOnce(&str) -> PyResult<R>,
+            F: FnOnce(&Wtf8) -> PyResult<R>,
        {
            let string = string.payload::<PyStr>().ok_or_else(|| {
                vm.new_type_error(format!("expected string got '{}'", string.class()))
            })?;
-            f(string.as_str())
+            f(string.as_wtf8())
        }

        fn with_bytes<F, R>(string: &PyObject, vm: &VirtualMachine, f: F) -> PyResult<R>
@@ -425,7 +431,7 @@ mod _sre {
                let is_template = if zelf.isbytes {
                    Self::with_bytes(&repl, vm, |x| Ok(x.contains(&b'\\')))?
                } else {
-                    Self::with_str(&repl, vm, |x| Ok(x.contains('\\')))?
+                    Self::with_str(&repl, vm, |x| Ok(x.contains("\\".as_ref())))?
                };

                if is_template {

diff --git a/vm/sre_engine/Cargo.toml b/vm/sre_engine/Cargo.toml
@@ -14,7 +14,12 @@ license.workspace = true
 name = "benches"
 harness = false

+[features]
+default = ["wtf8"]
+wtf8 = ["rustpython-common"]
+
 [dependencies]
+rustpython-common = { workspace = true, optional = true }
 num_enum = { workspace = true }
 bitflags = { workspace = true }
 optional = "0.5"

diff --git a/vm/sre_engine/src/string.rs b/vm/sre_engine/src/string.rs
@@ -1,3 +1,6 @@
+#[cfg(feature = "wtf8")]
+use rustpython_common::wtf8::Wtf8;
+
 #[derive(Debug, Clone, Copy)]
 pub struct StringCursor {
    pub(crate) ptr: *const u8,
@@ -148,6 +151,73 @@ impl StrDrive for &str {
    }
 }

+#[cfg(feature = "wtf8")]
+impl StrDrive for &Wtf8 {
+    #[inline]
+    fn count(&self) -> usize {
+        self.code_points().count()
+    }
+
+    #[inline]
+    fn create_cursor(&self, n: usize) -> StringCursor {
+        let mut cursor = StringCursor {
+            ptr: self.as_bytes().as_ptr(),
+            position: 0,
+        };
+        Self::skip(&mut cursor, n);
+        cursor
+    }
+
+    #[inline]
+    fn adjust_cursor(&self, cursor: &mut StringCursor, n: usize) {
+        if cursor.ptr.is_null() || cursor.position > n {
+            *cursor = Self::create_cursor(self, n);
+        } else if cursor.position < n {
+            Self::skip(cursor, n - cursor.position);
+        }
+    }
+
+    #[inline]
+    fn advance(cursor: &mut StringCursor) -> u32 {
+        cursor.position += 1;
+        unsafe { next_code_point(&mut cursor.ptr) }
+    }
+
+    #[inline]
+    fn peek(cursor: &StringCursor) -> u32 {
+        let mut ptr = cursor.ptr;
+        unsafe { next_code_point(&mut ptr) }
+    }
+
+    #[inline]
+    fn skip(cursor: &mut StringCursor, n: usize) {
+        cursor.position += n;
+        for _ in 0..n {
+            unsafe { next_code_point(&mut cursor.ptr) };
+        }
+    }
+
+    #[inline]
+    fn back_advance(cursor: &mut StringCursor) -> u32 {
+        cursor.position -= 1;
+        unsafe { next_code_point_reverse(&mut cursor.ptr) }
+    }
+
+    #[inline]
+    fn back_peek(cursor: &StringCursor) -> u32 {
+        let mut ptr = cursor.ptr;
+        unsafe { next_code_point_reverse(&mut ptr) }
+    }
+
+    #[inline]
+    fn back_skip(cursor: &mut StringCursor, n: usize) {
+        cursor.position -= n;
+        for _ in 0..n {
+            unsafe { next_code_point_reverse(&mut cursor.ptr) };
+        }
+    }
+}
+
 /// Reads the next code point out of a byte iterator (assuming a
 /// UTF-8-like encoding).
 ///