Fix bytes.find

Author: youknowone

Lib/test/string_tests.py

@@ -157,7 +157,8 @@ def test_count(self):
                     self.assertEqual(rem, 0, '%s != 0 for %s' % (rem, i))
                     self.assertEqual(r1, r2, '%s != %s for %s' % (r1, r2, i))
 
-    @unittest.skip("TODO: RUSTPYTHON test_bytes")
+    # TODO: RUSTPYTHON
+    @unittest.expectedFailure
     def test_find(self):
         self.checkequal(0, 'abcdefghiabc', 'find', 'abc')
         self.checkequal(9, 'abcdefghiabc', 'find', 'abc', 1)
@@ -215,7 +216,8 @@ def test_find(self):
                 if loc != -1:
                     self.assertEqual(i[loc:loc+len(j)], j)
 
-    @unittest.skip("TODO: RUSTPYTHON test_bytes")
+    # TODO: RUSTPYTHON
+    @unittest.expectedFailure
     def test_rfind(self):
         self.checkequal(9,  'abcdefghiabc', 'rfind', 'abc')
         self.checkequal(12, 'abcdefghiabc', 'rfind', '')
@@ -294,7 +296,6 @@ def test_index(self):
         else:
             self.checkraises(TypeError, 'hello', 'index', 42)
 
-    @unittest.skip("TODO: RUSTPYTHON test_bytes")
     def test_rindex(self):
         self.checkequal(12, 'abcdefghiabc', 'rindex', '')
         self.checkequal(3,  'abcdefghiabc', 'rindex', 'def')

vm/src/obj/objbytearray.rs

@@ -1,6 +1,9 @@
 //! Implementation of the python bytearray object.
+use bstr::ByteSlice;
 use crossbeam_utils::atomic::AtomicCell;
 use std::convert::TryFrom;
+use std::mem::size_of;
+use std::str::FromStr;
 use std::sync::{RwLock, RwLockReadGuard, RwLockWriteGuard};
 
 use super::objbyteinner::{
@@ -21,8 +24,6 @@ use crate::pyobject::{
     PyValue, ThreadSafe, TryFromObject, TypeProtocol,
 };
 use crate::vm::VirtualMachine;
-use std::mem::size_of;
-use std::str::FromStr;
 
 /// "bytearray(iterable_of_ints) -> bytearray\n\
 ///  bytearray(string, encoding[, errors]) -> bytearray\n\
@@ -327,25 +328,25 @@ impl PyByteArray {
 
     #[pymethod(name = "find")]
     fn find(&self, options: ByteInnerFindOptions, vm: &VirtualMachine) -> PyResult<isize> {
-        let index = self.borrow_value().find(options, false, vm)?;
+        let index = self.borrow_value().find(options, |h, n| h.find(n), vm)?;
         Ok(index.map_or(-1, |v| v as isize))
     }
 
     #[pymethod(name = "index")]
     fn index(&self, options: ByteInnerFindOptions, vm: &VirtualMachine) -> PyResult<usize> {
-        let index = self.borrow_value().find(options, false, vm)?;
+        let index = self.borrow_value().find(options, |h, n| h.find(n), vm)?;
         index.ok_or_else(|| vm.new_value_error("substring not found".to_owned()))
     }
 
     #[pymethod(name = "rfind")]
     fn rfind(&self, options: ByteInnerFindOptions, vm: &VirtualMachine) -> PyResult<isize> {
-        let index = self.borrow_value().find(options, true, vm)?;
+        let index = self.borrow_value().find(options, |h, n| h.rfind(n), vm)?;
         Ok(index.map_or(-1, |v| v as isize))
     }
 
     #[pymethod(name = "rindex")]
     fn rindex(&self, options: ByteInnerFindOptions, vm: &VirtualMachine) -> PyResult<usize> {
-        let index = self.borrow_value().find(options, true, vm)?;
+        let index = self.borrow_value().find(options, |h, n| h.rfind(n), vm)?;
         index.ok_or_else(|| vm.new_value_error("substring not found".to_owned()))
     }
 

vm/src/obj/objbyteinner.rs

@@ -823,35 +823,17 @@ impl PyByteInner {
     }
 
     #[inline]
-    pub fn find(
+    pub fn find<F>(
         &self,
         options: ByteInnerFindOptions,
-        reverse: bool,
+        find: F,
         vm: &VirtualMachine,
-    ) -> PyResult<Option<usize>> {
+    ) -> PyResult<Option<usize>>
+    where
+        F: Fn(&[u8], &[u8]) -> Option<usize>,
+    {
         let (needle, range) = options.get_value(self.elements.len(), vm)?;
-        if !range.is_normal() {
-            return Ok(None);
-        }
-        if needle.is_empty() {
-            return Ok(Some(if reverse { range.end } else { range.start }));
-        }
-        let haystack = &self.elements[range.clone()];
-        let windows = haystack.windows(needle.len());
-        if reverse {
-            for (i, w) in windows.rev().enumerate() {
-                if w == needle.as_slice() {
-                    return Ok(Some(range.end - i - needle.len()));
-                }
-            }
-        } else {
-            for (i, w) in windows.enumerate() {
-                if w == needle.as_slice() {
-                    return Ok(Some(range.start + i));
-                }
-            }
-        }
-        Ok(None)
+        Ok(self.elements.py_find(&needle, range, find))
     }
 
     pub fn maketrans(from: PyByteInner, to: PyByteInner, vm: &VirtualMachine) -> PyResult {

vm/src/obj/objbytes.rs

@@ -1,3 +1,4 @@
+use bstr::ByteSlice;
 use crossbeam_utils::atomic::AtomicCell;
 use std::mem::size_of;
 use std::ops::Deref;
@@ -300,25 +301,25 @@ impl PyBytes {
 
     #[pymethod(name = "find")]
     fn find(&self, options: ByteInnerFindOptions, vm: &VirtualMachine) -> PyResult<isize> {
-        let index = self.inner.find(options, false, vm)?;
+        let index = self.inner.find(options, |h, n| h.find(n), vm)?;
         Ok(index.map_or(-1, |v| v as isize))
     }
 
     #[pymethod(name = "index")]
     fn index(&self, options: ByteInnerFindOptions, vm: &VirtualMachine) -> PyResult<usize> {
-        let index = self.inner.find(options, false, vm)?;
+        let index = self.inner.find(options, |h, n| h.find(n), vm)?;
         index.ok_or_else(|| vm.new_value_error("substring not found".to_owned()))
     }
 
     #[pymethod(name = "rfind")]
     fn rfind(&self, options: ByteInnerFindOptions, vm: &VirtualMachine) -> PyResult<isize> {
-        let index = self.inner.find(options, true, vm)?;
+        let index = self.inner.find(options, |h, n| h.rfind(n), vm)?;
         Ok(index.map_or(-1, |v| v as isize))
     }
 
     #[pymethod(name = "rindex")]
     fn rindex(&self, options: ByteInnerFindOptions, vm: &VirtualMachine) -> PyResult<usize> {
-        let index = self.inner.find(options, true, vm)?;
+        let index = self.inner.find(options, |h, n| h.rfind(n), vm)?;
         index.ok_or_else(|| vm.new_value_error("substring not found".to_owned()))
     }
 

vm/src/obj/objstr.rs

@@ -811,6 +811,7 @@ impl PyString {
         Ok(joined)
     }
 
+    #[inline]
     fn _find<F>(
         &self,
         sub: PyStringRef,
@@ -822,12 +823,7 @@ impl PyString {
         F: Fn(&str, &str) -> Option<usize>,
     {
         let range = adjust_indices(start, end, self.value.len());
-        if range.is_normal() {
-            if let Some(index) = find(&self.value[range.clone()], &sub.value) {
-                return Some(range.start + index);
-            }
-        }
-        None
+        self.value.py_find(&sub.value, range, find)
     }
 
     #[pymethod]

vm/src/obj/pystr.rs

@@ -186,6 +186,7 @@ pub trait PyCommonString<E> {
         }
     }
 
+    #[inline]
     fn py_strip<'a, S, FC, FD>(
         &'a self,
         chars: OptionalOption<S>,
@@ -203,4 +204,18 @@ pub trait PyCommonString<E> {
             None => func_default(self),
         }
     }
+
+    #[inline]
+    fn py_find<F>(&self, needle: &Self, range: std::ops::Range<usize>, find: F) -> Option<usize>
+    where
+        F: Fn(&Self, &Self) -> Option<usize>,
+    {
+        if range.is_normal() {
+            let start = range.start;
+            if let Some(index) = find(self.get_slice(range), &needle) {
+                return Some(start + index);
+            }
+        }
+        None
+    }
 }