Rusty CPython Iterator

Who am I?

• Yuya Nishihara

  twitter: @yujauja

• Embedded software engineer

• Mercurial committer

  • DVCS implemented in Python
    • with C
    • and Rust

Rust-Python Bindings

Two major crates:

PyO3 "minimum required Rust version is 1.37.0-nightly"

rust-cpython Basics

py_class!(class IntSet |py| {
    data hash_set: RefCell<HashSet<i32>>;
});

which (basically) generates:

struct IntSet { ...: *mut ffi::PyObject }

impl IntSet {
    // data accessor
    fn hash_set<'a>(&'a self, py: Python<'a>)
       -> &'a RefCell<HashSet<i32>> { ... }
}

struct IntSet { ...: *mut ffi::PyObject }

impl IntSet {
    fn hash_set<'a>(&'a self, py: Python<'a>)
       -> &'a RefCell<HashSet<i32>> { ... }
}

• All data backed by CPython memory (i.e. reference counted.)
• Any data field must be Send + 'static,
• py: Python<'a> (!Sync nor !Send) marker guarantees GIL acquired,
• so that IntSet can be Send and Sync.

Iterator

Python iterator

s = IntSet(...)
for e in s:  # iter(s) => IntSetIter
    ...

over Rust collection:

py_class!(class IntSetIter |py| {
    data iter: RefCell<Iter<'_, i32>>;  // doesn't compile

    def __next__(&self) -> PyResult<Option<PyInt>> {
        self.iter(py).borrow().next()
    }
})

but...

The Problem

<'a> cannot live in the Python world.

hash_set.iter()
    : fn(&'a self) -> Iter<'a, i32>

Someone says

but <'static> can!!

hash_set.iter()
    : fn(&'static self) -> Iter<'static, i32>

unsafe { &*ptr }

The idea

Cheated lifetime so the iterator can (out)live in the Python world.

let hash_set: &'a RefCell<HashSet<i32>> = self.hash_set(py);
let ptr: *mut HashSet<i32> = hash_set.as_ptr();
let static_ref: &'static HashSet<i32> = unsafe { &*ptr };
let static_iter: Iter<'static, i32> = static_ref.iter();

How to make it safe?

// &'a -> &'static
unsafe { &*ptr }

Things to consider:

• Real object lifetime
• Multi-threading stuff
• Mutation of underlying object
• Safe interface on top of these three

Object lifetime

Incref the owner PyObject and keep it.

struct PyLeaked<T> {
    owner: PyObject,
    data: T,  // T: Iter<'static, i32>
    ...
}

PyLeaked {
    owner: owner.clone_ref(py),
    data: static_ref.iter(),
}

Multi-threading stuff

Easy. Just require GIL for any operation.

let leaked_iter: PyLeaked<Iter<'static, i32>>;

fn __next__<'a>(&'a self, py: Python<'a>) ... {
    let mut iter = leaked_iter.try_borrow_mut(py);
                        // not try_borrow_mut()
    iter.next()
}

Mutation #1

Invalidate references when underlying object mutated.

struct IntSet { hash_set: PySharedRefCell<HashSet<i32>> }
struct IntSetIter { leaked_iter: PyLeaked<Iter<'static, i32>> }

leaked_iter = hash_set.do_some_weird_thing();
              // remember generation
impl IntSet { fn add(...) {
    self.hash_set(py).borrow_mut(); } }  // generation += 1
impl IntSetIter { fn __next__(...) {
    self.leaked_iter(py).try_borrow_mut(py); } }
                         // compare generation values

# for e in int_set:
int_set_iter = iter(int_set)  # generation = 3
while True:
    e = next(int_set_iter)  # compare generation values
    int_set.add(1)  # generation = 3 + 1

Mutation #2

let hash_set: PySharedRefCell<HashSet<i32>>;
let leaked_iter: PyLeaked<Iter<'static, i32>>;

{
    let mut iter = leaked_iter.try_borrow_mut(py)?;
                            // borrow += 1
    hash_set.borrow_mut();  // panics if borrow > 0
    // drop(iter): borrow -= 1
}
hash_set.borrow_mut();  // okay

Safe interface

fn map<U>(self, py: Python,
          f: impl FnOnce(&'static T) -> U) -> PyLeaked<U>
fn try_borrow_mut(&mut self, py: Python)
                  -> PyResult<PyLeakedRefMut<T>>

py_class!(class IntSet |py| {
    data hash_set: PySharedRefCell<HashSet<i32>>;
});

let hash_set = self.shared_hash_set(py);
let leaked = hash_set.leak_immutable();
let leaked_iter = leaked.map(py, |s| s.iter());

def __next__(&self) -> PyResult<Option<PyInt>> {
    let mut iter = self.leaked_iter(py).try_borrow_mut(py)?;
    iter.next()
}

Safe interface is not safe

unsafe fn map<U>(self, py: Python,
                 f: impl FnOnce(&'static T) -> U) -> PyLeaked<U>
unsafe fn try_borrow_mut(&mut self, py: Python)
                         -> PyResult<PyLeakedRefMut<T>>

because the lifetime is cheated.

Many ways to obtain the underlying &'static:

let outer: &'static ...;
leaked.map(py, |s| { outer = s });
leaked_iter.try_borrow_mut(py)?.as_slice();
    // slice::Iter::as_slice(&self) -> &'a [T]
leaked_iter.try_borrow_mut(py)?.next();
    // slice::Iter::next(&mut self) -> Option<&'a T>

Lifetime is hard

Source code

$ hg clone https://www.mercurial-scm.org/repo/hg-committed/

• https://www.mercurial-scm.org/repo/hg-committed/file/tip/rust/hg-cpython/src/ref_sharing.rs
• will hopefully be upstreamed soon!

Further Reading

• Sharing references between Python and Rust, by Raphaël Gomès

  https://raphaelgomes.dev/blog/articles/2019-07-01-sharing-references-between-python-and-rust.html

• PyO3 and rust-cpython

  https://pyo3.rs/master/rust_cpython.html