add free list support

pyo3cls/src/py_class.rs

@@ -88,6 +88,33 @@ fn impl_class(cls: &syn::Ident, base: &syn::Ident,
         None
     };
 
+    let extra = {
+        if let Some(freelist) = params.get("freelist") {
+            Some(quote! {
+                impl _pyo3::freelist::PyObjectWithFreeList for #cls {
+                    #[inline]
+                    fn get_free_list() -> &'static mut _pyo3::freelist::FreeList<*mut ffi::PyObject> {
+                        static mut FREELIST: *mut _pyo3::freelist::FreeList<*mut ffi::PyObject> = 0 as *mut _;
+                        unsafe {
+                            if FREELIST.is_null() {
+                                FREELIST = Box::into_raw(Box::new(
+                                    _pyo3::freelist::FreeList::with_capacity(#freelist)));
+
+                                <#cls as _pyo3::typeob::PyTypeObject>::init_type(
+                                    _pyo3::Python::assume_gil_acquired());
+                            }
+                            std::mem::transmute(FREELIST)
+                        }
+                    }
+                }
+
+                #extra
+            })
+        } else {
+            extra
+        }
+    };
+
     quote! {
         impl _pyo3::typeob::PyTypeInfo for #cls {
             type Type = #cls;

src/freelist.rs

@@ -0,0 +1,97 @@
+// Copyright (c) 2017-present PyO3 Project and Contributors
+use std;
+
+use ffi;
+use err::PyResult;
+use python::Python;
+use typeob::{PyTypeInfo, PyObjectAlloc};
+
+/// Implementing this trait for custom class adds free allocation list to class.
+/// The performance improvement applies to types that are often created and deleted in a row,
+/// so that they can benefit from a freelist.
+pub trait PyObjectWithFreeList: PyTypeInfo {
+
+    fn get_free_list() -> &'static mut FreeList<*mut ffi::PyObject>;
+
+}
+
+pub enum Slot<T> {
+    Empty,
+    Filled(T),
+}
+
+pub struct FreeList<T> {
+    entries: Vec<Slot<T>>,
+    split: usize,
+    capacity: usize,
+}
+
+impl<T> FreeList<T> {
+
+    /// Create new `FreeList` instance with specified capacity
+    pub fn with_capacity(capacity: usize) -> FreeList<T> {
+        let entries = (0..capacity).map(|_| Slot::Empty).collect::<Vec<_>>();
+
+        FreeList {
+            entries: entries,
+            split: 0,
+            capacity: capacity,
+        }
+    }
+
+    /// Pop first non empty item
+    pub fn pop(&mut self) -> Option<T> {
+        let idx = self.split;
+        if idx == 0 {
+            None
+        } else {
+            match std::mem::replace(&mut self.entries[idx-1], Slot::Empty) {
+                Slot::Filled(v) => {
+                    self.split = idx - 1;
+                    Some(v)
+                }
+                _ => panic!("FreeList is corrupt")
+            }
+        }
+    }
+
+    /// Insert a value into the list
+    pub fn insert(&mut self, val: T) -> Option<T> {
+        let next = self.split + 1;
+        if next < self.capacity {
+            self.entries[self.split] = Slot::Filled(val);
+            self.split = next;
+            None
+        } else {
+            Some(val)
+        }
+    }
+}
+
+
+impl<T> PyObjectAlloc<T> for T where T: PyObjectWithFreeList {
+
+    unsafe fn alloc(_py: Python, value: T) -> PyResult<*mut ffi::PyObject> {
+        let obj = if let Some(obj) = <T as PyObjectWithFreeList>::get_free_list().pop() {
+            ffi::PyObject_Init(obj, <T as PyTypeInfo>::type_object());
+            obj
+        } else {
+            ffi::PyType_GenericAlloc(<T as PyTypeInfo>::type_object(), 0)
+        };
+
+        let offset = <T as PyTypeInfo>::offset();
+        let ptr = (obj as *mut u8).offset(offset) as *mut T;
+        std::ptr::write(ptr, value);
+
+        Ok(obj)
+    }
+
+    unsafe fn dealloc(_py: Python, obj: *mut ffi::PyObject) {
+        let ptr = (obj as *mut u8).offset(<T as PyTypeInfo>::offset()) as *mut T;
+        std::ptr::drop_in_place(ptr);
+
+        if let Some(obj) = <T as PyObjectWithFreeList>::get_free_list().insert(obj) {
+            ffi::PyObject_Free(obj as *mut ::c_void);
+        }
+    }
+}

src/lib.rs

@@ -121,6 +121,7 @@ pub mod typeob;
 pub mod argparse;
 pub mod function;
 pub mod buffer;
+pub mod freelist;
 
 // re-export for simplicity
 pub use std::os::raw::*;

src/objects/tuple.rs

@@ -196,7 +196,6 @@ tuple_conversion!(9, (ref0, 0, A), (ref1, 1, B), (ref2, 2, C), (ref3, 3, D),
 ///
 /// # Example
 /// ```
-/// use pyo3::ObjectProtocol;
 /// let gil = pyo3::Python::acquire_gil();
 /// let py = gil.python();
 /// let os = py.import("os").unwrap();

src/python.rs

@@ -233,7 +233,7 @@ impl<'p> Python<'p> {
     #[inline]
     pub fn init<T, F>(self, f: F) -> PyResult<T::Target>
         where F: FnOnce(PyToken) -> T,
-              T: ToInstancePtr<T> + PyTypeInfo + PyObjectAlloc<Type=T>
+              T: ToInstancePtr<T> + PyTypeInfo + PyObjectAlloc<T>
     {
         ::token::init(self, f)
     }

src/token.rs

@@ -21,12 +21,12 @@ impl PyToken {
 #[inline]
 pub fn init<'p, T, F>(py: Python<'p>, f: F) -> PyResult<T::Target>
     where F: FnOnce(PyToken) -> T,
-          T: ToInstancePtr<T> + PyTypeInfo + PyObjectAlloc<Type=T>
+          T: ToInstancePtr<T> + PyTypeInfo + PyObjectAlloc<T>
 {
     let ob = f(PyToken(PhantomData));
 
     let ob = unsafe {
-        let ob = try!(<T as PyObjectAlloc>::alloc(py, ob));
+        let ob = try!(<T as PyObjectAlloc<T>>::alloc(py, ob));
         T::from_owned_ptr(ob)
     };
     Ok(ob)

src/typeob.rs

@@ -59,12 +59,11 @@ impl<'a, T: ?Sized> PyTypeInfo for &'a T where T: PyTypeInfo {
     }
 }
 
-pub trait PyObjectAlloc {
+pub trait PyObjectAlloc<T> {
-    type Type;
 
     /// Allocates a new object (usually by calling ty->tp_alloc),
     /// and initializes it using value.
-    unsafe fn alloc(py: Python, value: Self::Type) -> PyResult<*mut ffi::PyObject>;
+    unsafe fn alloc(py: Python, value: T) -> PyResult<*mut ffi::PyObject>;
 
     /// Calls the rust destructor for the object and frees the memory
     /// (usually by calling ptr->ob_type->tp_free).
@@ -73,12 +72,11 @@ pub trait PyObjectAlloc {
 }
 
 /// A Python object allocator that is usable as a base type for #[class]
-impl<T> PyObjectAlloc for T where T : PyTypeInfo {
+impl<T> PyObjectAlloc<T> for T where T : PyTypeInfo {
-    type Type = T::Type;
 
     /// Allocates a new object (usually by calling ty->tp_alloc),
     /// and initializes it using value.
-    default unsafe fn alloc(py: Python, value: Self::Type) -> PyResult<*mut ffi::PyObject> {
+    default unsafe fn alloc(py: Python, value: T) -> PyResult<*mut ffi::PyObject> {
         // TODO: remove this
         <T as PyTypeObject>::init_type(py);
 
@@ -86,7 +84,7 @@ impl<T> PyObjectAlloc for T where T : PyTypeInfo {
             <Self as PyTypeInfo>::type_object(), 0);
 
         let offset = <Self as PyTypeInfo>::offset();
-        let ptr = (obj as *mut u8).offset(offset) as *mut Self::Type;
+        let ptr = (obj as *mut u8).offset(offset) as *mut T;
         std::ptr::write(ptr, value);
 
         Ok(obj)
@@ -96,7 +94,7 @@ impl<T> PyObjectAlloc for T where T : PyTypeInfo {
     /// (usually by calling ptr->ob_type->tp_free).
     /// This function is used as tp_dealloc implementation.
     default unsafe fn dealloc(_py: Python, obj: *mut ffi::PyObject) {
-        let ptr = (obj as *mut u8).offset(<Self as PyTypeInfo>::offset()) as *mut Self::Type;
+        let ptr = (obj as *mut u8).offset(<Self as PyTypeInfo>::offset()) as *mut T;
         std::ptr::drop_in_place(ptr);
 
         let ty = ffi::Py_TYPE(obj);
@@ -124,7 +122,7 @@ pub trait PyTypeObject {
 
 }
 
-impl<T> PyTypeObject for T where T: PyObjectAlloc + PyTypeInfo {
+impl<T> PyTypeObject for T where T: PyObjectAlloc<T> + PyTypeInfo {
 
     #[inline]
     default fn init_type(py: Python) {
@@ -151,7 +149,7 @@ impl<T> PyTypeObject for T where T: PyObjectAlloc + PyTypeInfo {
 
 pub fn initialize_type<T>(py: Python, module_name: Option<&str>, type_name: &str,
                           type_object: &mut ffi::PyTypeObject) -> PyResult<PyType>
-    where T: PyObjectAlloc + PyTypeInfo
+    where T: PyObjectAlloc<T> + PyTypeInfo
 {
     // type name
     let name = match module_name {
@@ -247,13 +245,13 @@ pub fn initialize_type<T>(py: Python, module_name: Option<&str>, type_name: &str
 
 
 unsafe extern "C" fn tp_dealloc_callback<T>(obj: *mut ffi::PyObject)
-    where T: PyTypeInfo
+    where T: PyObjectAlloc<T>
 {
     debug!("DEALLOC: {:?} - {:?}", obj,
            CStr::from_ptr((*(*obj).ob_type).tp_name).to_string_lossy());
     let guard = AbortOnDrop("Cannot unwind out of tp_dealloc");
     let py = Python::assume_gil_acquired();
-    let r = <T as PyObjectAlloc>::dealloc(py, obj);
+    let r = <T as PyObjectAlloc<T>>::dealloc(py, obj);
     mem::forget(guard);
     r
 }

tests/test_class.rs

@@ -159,6 +159,26 @@ fn new_with_two_args() {
     assert_eq!(obj._data2, 20);
 }
 
+#[py::class(freelist=10)]
+struct ClassWithFreelist{token: PyToken}
+
+#[py::ptr(ClassWithFreelist)]
+struct ClassWithFreelistPtr(PyPtr);
+
+#[test]
+fn class_with_freelist() {
+    let gil = Python::acquire_gil();
+    let py = gil.python();
+
+    let inst = py.init(|t| ClassWithFreelist{token: t}).unwrap();
+    let inst2 = py.init(|t| ClassWithFreelist{token: t}).unwrap();
+    let ptr = inst.as_ptr();
+    drop(inst);
+
+    let inst3 = py.init(|t| ClassWithFreelist{token: t}).unwrap();
+    assert_eq!(ptr, inst3.as_ptr());
+}
+
 struct TestDropCall {
     drop_called: Arc<AtomicBool>
 }