rustpython_vm/
iter.rs

1use crate::{PyObject, PyObjectRef, PyResult, types::PyComparisonOp, vm::VirtualMachine};
2use itertools::Itertools;
3
4pub trait PyExactSizeIterator<'a>: ExactSizeIterator<Item = &'a PyObject> + Sized {
5    fn eq(self, other: impl PyExactSizeIterator<'a>, vm: &VirtualMachine) -> PyResult<bool> {
6        let lhs = self;
7        let rhs = other;
8        if lhs.len() != rhs.len() {
9            return Ok(false);
10        }
11        for (a, b) in lhs.zip_eq(rhs) {
12            if !vm.identical_or_equal(a, b)? {
13                return Ok(false);
14            }
15        }
16        Ok(true)
17    }
18
19    fn richcompare(
20        self,
21        other: impl PyExactSizeIterator<'a>,
22        op: PyComparisonOp,
23        vm: &VirtualMachine,
24    ) -> PyResult<bool> {
25        let less = match op {
26            PyComparisonOp::Eq => return PyExactSizeIterator::eq(self, other, vm),
27            PyComparisonOp::Ne => return PyExactSizeIterator::eq(self, other, vm).map(|eq| !eq),
28            PyComparisonOp::Lt | PyComparisonOp::Le => true,
29            PyComparisonOp::Gt | PyComparisonOp::Ge => false,
30        };
31
32        let lhs = self;
33        let rhs = other;
34        let lhs_len = lhs.len();
35        let rhs_len = rhs.len();
36        for (a, b) in lhs.zip(rhs) {
37            if vm.bool_eq(a, b)? {
38                continue;
39            }
40            let ret = if less {
41                vm.bool_seq_lt(a, b)?
42            } else {
43                vm.bool_seq_gt(a, b)?
44            };
45            if let Some(v) = ret {
46                return Ok(v);
47            }
48        }
49        Ok(op.eval_ord(lhs_len.cmp(&rhs_len)))
50    }
51}
52
53impl<'a, T> PyExactSizeIterator<'a> for T where T: ExactSizeIterator<Item = &'a PyObject> + Sized {}
54
55/// Compare two sequences whose item-comparison callbacks may mutate the
56/// containers. `get(i)` must snapshot `(len, item_at_i)` and drop any
57/// container lock before returning.
58///
59/// After an unequal item, sizes are read again: if a callback shrank a
60/// sequence so `i` is now past the end, the result is a length comparison
61/// rather than the item inequality.
62pub fn richcompare_mutating_seqs(
63    mut get_a: impl FnMut(usize) -> (usize, Option<PyObjectRef>),
64    mut get_b: impl FnMut(usize) -> (usize, Option<PyObjectRef>),
65    op: PyComparisonOp,
66    vm: &VirtualMachine,
67) -> PyResult<bool> {
68    if matches!(op, PyComparisonOp::Eq | PyComparisonOp::Ne) {
69        let (a_len, _) = get_a(0);
70        let (b_len, _) = get_b(0);
71        if a_len != b_len {
72            return Ok(op == PyComparisonOp::Ne);
73        }
74    }
75
76    let mut i = 0usize;
77    loop {
78        let (a_len, a_item) = get_a(i);
79        let (b_len, b_item) = get_b(i);
80        let (Some(a_item), Some(b_item)) = (a_item, b_item) else {
81            return Ok(op.eval_ord(a_len.cmp(&b_len)));
82        };
83
84        if vm.bool_eq(&a_item, &b_item)? {
85            i += 1;
86            continue;
87        }
88
89        let (a_len, _) = get_a(i);
90        let (b_len, _) = get_b(i);
91        if i >= a_len || i >= b_len {
92            return Ok(op.eval_ord(a_len.cmp(&b_len)));
93        }
94
95        return match op {
96            PyComparisonOp::Eq => Ok(false),
97            PyComparisonOp::Ne => Ok(true),
98            _ => a_item.rich_compare_bool(&b_item, op, vm),
99        };
100    }
101}
rustpython_vm/iter.rs

rustpython_vm/
iter.rs
