Merge pull request RustPython#3785 from youknowone/exact-size-operator

PyExactSizeIterator

Author: youknowone

stdlib/src/array.rs

@@ -23,7 +23,7 @@ mod array {
                 BufferDescriptor, BufferMethods, BufferResizeGuard, PyBuffer, PyIterReturn,
                 PyMappingMethods,
             },
-            sequence::{SequenceMutOp, SequenceOp},
+            sequence::{SequenceExt, SequenceMutExt},
             sliceable::{
                 SaturatedSlice, SequenceIndex, SliceableSequenceMutOp, SliceableSequenceOp,
             },

vm/src/builtins/list.rs

@@ -6,9 +6,10 @@ use crate::{
     class::PyClassImpl,
     convert::ToPyObject,
     function::{FuncArgs, OptionalArg, PyComparisonValue},
+    iter::PyExactSizeIterator,
     protocol::{PyIterReturn, PyMappingMethods, PySequence, PySequenceMethods},
     recursion::ReprGuard,
-    sequence::{MutObjectSequenceOp, ObjectSequenceOp, SequenceMutOp, SequenceOp},
+    sequence::{MutObjectSequenceOp, SequenceExt, SequenceMutExt},
     sliceable::{saturate_index, SequenceIndex, SliceableSequenceMutOp, SliceableSequenceOp},
     types::{
         AsMapping, AsSequence, Comparable, Constructor, Hashable, Initializer, IterNext,
@@ -491,7 +492,9 @@ impl Comparable for PyList {
         let other = class_or_notimplemented!(Self, other);
         let a = &*zelf.borrow_vec();
         let b = &*other.borrow_vec();
-        a.cmp(vm, b, op).map(PyComparisonValue::Implemented)
+        a.iter()
+            .richcompare(b.iter(), op, vm)
+            .map(PyComparisonValue::Implemented)
     }
 }
 

vm/src/builtins/str.rs

@@ -11,7 +11,7 @@ use crate::{
     function::{ArgIterable, FuncArgs, OptionalArg, OptionalOption, PyComparisonValue},
     intern::PyInterned,
     protocol::{PyIterReturn, PyMappingMethods, PySequenceMethods},
-    sequence::SequenceOp,
+    sequence::SequenceExt,
     sliceable::{SequenceIndex, SliceableSequenceOp},
     types::{
         AsMapping, AsSequence, Comparable, Constructor, Hashable, IterNext, IterNextIterable,

vm/src/builtins/tuple.rs

@@ -4,9 +4,10 @@ use crate::{
     class::PyClassImpl,
     convert::{ToPyObject, TransmuteFromObject},
     function::{OptionalArg, PyArithmeticValue, PyComparisonValue},
+    iter::PyExactSizeIterator,
     protocol::{PyIterReturn, PyMappingMethods, PySequenceMethods},
     recursion::ReprGuard,
-    sequence::{ObjectSequenceOp, SequenceOp},
+    sequence::SequenceExt,
     sliceable::{SequenceIndex, SliceableSequenceOp},
     stdlib::sys,
     types::{
@@ -400,7 +401,9 @@ impl Comparable for PyTuple {
             return Ok(res.into());
         }
         let other = class_or_notimplemented!(Self, other);
-        zelf.cmp(vm, other, op).map(PyComparisonValue::Implemented)
+        zelf.iter()
+            .richcompare(other.iter(), op, vm)
+            .map(PyComparisonValue::Implemented)
     }
 }
 

vm/src/bytesinner.rs

@@ -7,7 +7,7 @@ use crate::{
     function::{ArgIterable, Either, OptionalArg, OptionalOption, PyComparisonValue},
     identifier,
     protocol::PyBuffer,
-    sequence::{SequenceMutOp, SequenceOp},
+    sequence::{SequenceExt, SequenceMutExt},
     types::PyComparisonOp,
     AsObject, PyObject, PyObjectRef, PyPayload, PyResult, TryFromBorrowedObject, VirtualMachine,
 };

vm/src/iter.rs

@@ -0,0 +1,50 @@
+use crate::{types::PyComparisonOp, vm::VirtualMachine, PyObjectRef, PyResult};
+use itertools::Itertools;
+
+pub trait PyExactSizeIterator<'a>: ExactSizeIterator<Item = &'a PyObjectRef> + Sized {
+    fn eq(self, other: impl PyExactSizeIterator<'a>, vm: &VirtualMachine) -> PyResult<bool> {
+        let lhs = self;
+        let rhs = other;
+        if lhs.len() != rhs.len() {
+            return Ok(false);
+        }
+        for (a, b) in lhs.zip_eq(rhs) {
+            if !vm.identical_or_equal(a, b)? {
+                return Ok(false);
+            }
+        }
+        Ok(true)
+    }
+
+    fn richcompare(
+        self,
+        other: impl PyExactSizeIterator<'a>,
+        op: PyComparisonOp,
+        vm: &VirtualMachine,
+    ) -> PyResult<bool> {
+        let less = match op {
+            PyComparisonOp::Eq => return PyExactSizeIterator::eq(self, other, vm),
+            PyComparisonOp::Ne => return PyExactSizeIterator::eq(self, other, vm).map(|eq| !eq),
+            PyComparisonOp::Lt | PyComparisonOp::Le => true,
+            PyComparisonOp::Gt | PyComparisonOp::Ge => false,
+        };
+
+        let lhs = self;
+        let rhs = other;
+        let lhs_len = lhs.len();
+        let rhs_len = rhs.len();
+        for (a, b) in lhs.zip(rhs) {
+            let ret = if less {
+                vm.bool_seq_lt(a, b)?
+            } else {
+                vm.bool_seq_gt(a, b)?
+            };
+            if let Some(v) = ret {
+                return Ok(v);
+            }
+        }
+        Ok(op.eval_ord(lhs_len.cmp(&rhs_len)))
+    }
+}
+
+impl<'a, T> PyExactSizeIterator<'a> for T where T: ExactSizeIterator<Item = &'a PyObjectRef> + Sized {}

vm/src/lib.rs

@@ -57,6 +57,7 @@ mod frozen;
 pub mod function;
 pub mod import;
 mod intern;
+pub mod iter;
 pub mod object;
 pub mod prelude;
 pub mod protocol;

vm/src/sequence.rs

@@ -4,70 +4,8 @@ use crate::{
     vm::VirtualMachine,
     AsObject, PyObject, PyObjectRef, PyResult,
 };
-use itertools::Itertools;
 use optional::Optioned;
-use std::{collections::VecDeque, ops::Range};
-
-pub trait ObjectSequenceOp<'a> {
-    type Iter: ExactSizeIterator<Item = &'a PyObjectRef>;
-
-    fn iter(&'a self) -> Self::Iter;
-
-    fn eq(&'a self, vm: &VirtualMachine, other: &'a Self) -> PyResult<bool> {
-        let lhs = self.iter();
-        let rhs = other.iter();
-        if lhs.len() != rhs.len() {
-            return Ok(false);
-        }
-        for (a, b) in lhs.zip_eq(rhs) {
-            if !vm.identical_or_equal(a, b)? {
-                return Ok(false);
-            }
-        }
-        Ok(true)
-    }
-
-    fn cmp(&'a self, vm: &VirtualMachine, other: &'a Self, op: PyComparisonOp) -> PyResult<bool> {
-        let less = match op {
-            PyComparisonOp::Eq => return self.eq(vm, other),
-            PyComparisonOp::Ne => return self.eq(vm, other).map(|eq| !eq),
-            PyComparisonOp::Lt | PyComparisonOp::Le => true,
-            PyComparisonOp::Gt | PyComparisonOp::Ge => false,
-        };
-
-        let lhs = self.iter();
-        let rhs = other.iter();
-        let lhs_len = lhs.len();
-        let rhs_len = rhs.len();
-        for (a, b) in lhs.zip(rhs) {
-            let ret = if less {
-                vm.bool_seq_lt(a, b)?
-            } else {
-                vm.bool_seq_gt(a, b)?
-            };
-            if let Some(v) = ret {
-                return Ok(v);
-            }
-        }
-        Ok(op.eval_ord(lhs_len.cmp(&rhs_len)))
-    }
-}
-
-impl<'a> ObjectSequenceOp<'a> for [PyObjectRef] {
-    type Iter = core::slice::Iter<'a, PyObjectRef>;
-
-    fn iter(&'a self) -> Self::Iter {
-        self.iter()
-    }
-}
-
-impl<'a> ObjectSequenceOp<'a> for VecDeque<PyObjectRef> {
-    type Iter = std::collections::vec_deque::Iter<'a, PyObjectRef>;
-
-    fn iter(&'a self) -> Self::Iter {
-        self.iter()
-    }
-}
+use std::ops::Range;
 
 pub trait MutObjectSequenceOp<'a> {
     type Guard;
@@ -242,7 +180,7 @@ pub trait MutObjectSequenceOp<'a> {
     }
 }
 
-pub trait SequenceOp<T: Clone>
+pub trait SequenceExt<T: Clone>
 where
     Self: AsRef<[T]>,
 {
@@ -256,9 +194,9 @@ where
     }
 }
 
-impl<T: Clone> SequenceOp<T> for [T] {}
+impl<T: Clone> SequenceExt<T> for [T] {}
 
-pub trait SequenceMutOp<T: Clone>
+pub trait SequenceMutExt<T: Clone>
 where
     Self: AsRef<[T]>,
 {
@@ -282,7 +220,7 @@ where
     }
 }
 
-impl<T: Clone> SequenceMutOp<T> for Vec<T> {
+impl<T: Clone> SequenceMutExt<T> for Vec<T> {
     fn as_vec_mut(&mut self) -> &mut Vec<T> {
         self
     }

vm/src/stdlib/collections.rs

@@ -9,9 +9,10 @@ mod _collections {
         },
         common::lock::{PyMutex, PyRwLock, PyRwLockReadGuard, PyRwLockWriteGuard},
         function::{FuncArgs, KwArgs, OptionalArg, PyComparisonValue},
+        iter::PyExactSizeIterator,
         protocol::{PyIterReturn, PySequenceMethods},
         recursion::ReprGuard,
-        sequence::{MutObjectSequenceOp, ObjectSequenceOp},
+        sequence::MutObjectSequenceOp,
         sliceable,
         sliceable::saturate_index,
         types::{
@@ -565,7 +566,9 @@ mod _collections {
             let other = class_or_notimplemented!(Self, other);
             let lhs = zelf.borrow_deque();
             let rhs = other.borrow_deque();
-            lhs.cmp(vm, &rhs, op).map(PyComparisonValue::Implemented)
+            lhs.iter()
+                .richcompare(rhs.iter(), op, vm)
+                .map(PyComparisonValue::Implemented)
         }
     }