From bd7aed4b126e43c5ebb8cf339034d0f26ba06215 Mon Sep 17 00:00:00 2001
From: Adam Reichold <adam.reichold@t-online.de>
Date: Sat, 8 Jul 2023 21:49:18 +0200
Subject: [PATCH] Add implementation of Iterator::size_hint for PyIterator

When the Python iterator backing `PyIterator` has a `__length_hint__` special
method, we can use this as a lower bound for Rust's `Iterator::size_hint` to
e.g. support pre-allocation of collections.

This is implemented using `PyObject_LengthHint` which is not available in the
stable ABI and hence so is `Iterator::size_hint`. This should be fine since this
is an optimization in any case and the stable ABI is expected to have slightly
worse performance overall.
---
 benches/bench_any.rs                  | 17 ++++++++++++++++-
 pyo3-ffi/src/boolobject.rs            |  4 ++--
 pyo3-macros-backend/src/attributes.rs |  2 +-
 src/impl_/pyclass.rs                  |  2 +-
 src/types/iterator.rs                 | 17 +++++++++++++++++
 src/types/mod.rs                      |  1 +
 6 files changed, 38 insertions(+), 5 deletions(-)

diff --git a/benches/bench_any.rs b/benches/bench_any.rs
index ec23cedc..765497fa 100644
--- a/benches/bench_any.rs
+++ b/benches/bench_any.rs
@@ -5,7 +5,7 @@ use pyo3::{
         PyBool, PyByteArray, PyBytes, PyDict, PyFloat, PyFrozenSet, PyInt, PyList, PyMapping,
         PySequence, PySet, PyString, PyTuple,
     },
-    PyAny, Python,
+    PyAny, PyResult, Python,
 };
 
 #[derive(PartialEq, Eq, Debug)]
@@ -71,8 +71,23 @@ fn bench_identify_object_type(b: &mut Bencher<'_>) {
     });
 }
 
+fn bench_collect_generic_iterator(b: &mut Bencher<'_>) {
+    Python::with_gil(|py| {
+        let collection = py.eval("list(range(1 << 20))", None, None).unwrap();
+
+        b.iter(|| {
+            collection
+                .iter()
+                .unwrap()
+                .collect::<PyResult<Vec<_>>>()
+                .unwrap()
+        });
+    });
+}
+
 fn criterion_benchmark(c: &mut Criterion) {
     c.bench_function("identify_object_type", bench_identify_object_type);
+    c.bench_function("collect_generic_iterator", bench_collect_generic_iterator);
 }
 
 criterion_group!(benches, criterion_benchmark);
diff --git a/pyo3-ffi/src/boolobject.rs b/pyo3-ffi/src/boolobject.rs
index 81e1308e..78972ff0 100644
--- a/pyo3-ffi/src/boolobject.rs
+++ b/pyo3-ffi/src/boolobject.rs
@@ -24,12 +24,12 @@ extern "C" {
 
 #[inline]
 pub unsafe fn Py_False() -> *mut PyObject {
-    addr_of_mut!(_Py_FalseStruct) as *mut PyLongObject as *mut PyObject
+    addr_of_mut!(_Py_FalseStruct) as *mut PyObject
 }
 
 #[inline]
 pub unsafe fn Py_True() -> *mut PyObject {
-    addr_of_mut!(_Py_TrueStruct) as *mut PyLongObject as *mut PyObject
+    addr_of_mut!(_Py_TrueStruct) as *mut PyObject
 }
 
 #[inline]
diff --git a/pyo3-macros-backend/src/attributes.rs b/pyo3-macros-backend/src/attributes.rs
index dc19c83c..e51a5a04 100644
--- a/pyo3-macros-backend/src/attributes.rs
+++ b/pyo3-macros-backend/src/attributes.rs
@@ -179,7 +179,7 @@ pub fn take_pyo3_options<T: Parse>(attrs: &mut Vec<syn::Attribute>) -> Result<Ve
     let mut out = Vec::new();
     take_attributes(attrs, |attr| {
         if let Some(options) = get_pyo3_options(attr)? {
-            out.extend(options.into_iter());
+            out.extend(options);
             Ok(true)
         } else {
             Ok(false)
diff --git a/src/impl_/pyclass.rs b/src/impl_/pyclass.rs
index 2ba722d4..07fdcc44 100644
--- a/src/impl_/pyclass.rs
+++ b/src/impl_/pyclass.rs
@@ -122,7 +122,7 @@ impl<T> Default for PyClassImplCollector<T> {
 
 impl<T> Clone for PyClassImplCollector<T> {
     fn clone(&self) -> Self {
-        Self::new()
+        *self
     }
 }
 
diff --git a/src/types/iterator.rs b/src/types/iterator.rs
index 3b45c113..71ca5d30 100644
--- a/src/types/iterator.rs
+++ b/src/types/iterator.rs
@@ -55,6 +55,12 @@ impl<'p> Iterator for &'p PyIterator {
             None => PyErr::take(py).map(Err),
         }
     }
+
+    #[cfg(not(Py_LIMITED_API))]
+    fn size_hint(&self) -> (usize, Option<usize>) {
+        let hint = unsafe { ffi::PyObject_LengthHint(self.0.as_ptr(), 0) };
+        (hint.max(0) as usize, None)
+    }
 }
 
 // PyIter_Check does not exist in the limited API until 3.8
@@ -313,4 +319,15 @@ def fibonacci(target):
             );
         });
     }
+
+    #[test]
+    #[cfg(not(Py_LIMITED_API))]
+    fn length_hint_becomes_size_hint_lower_bound() {
+        Python::with_gil(|py| {
+            let list = py.eval("[1, 2, 3]", None, None).unwrap();
+            let iter = list.iter().unwrap();
+            let hint = iter.size_hint();
+            assert_eq!(hint, (3, None));
+        });
+    }
 }
diff --git a/src/types/mod.rs b/src/types/mod.rs
index 6fbc5d2e..6cf52593 100644
--- a/src/types/mod.rs
+++ b/src/types/mod.rs
@@ -196,6 +196,7 @@ macro_rules! pyobject_native_type_info(
             const MODULE: ::std::option::Option<&'static str> = $module;
 
             #[inline]
+            #[allow(clippy::redundant_closure_call)]
             fn type_object_raw(py: $crate::Python<'_>) -> *mut $crate::ffi::PyTypeObject {
                 $typeobject(py)
             }