1
  2
  3
  4
  5
  6
  7
  8
  9
 10
 11
 12
 13
 14
 15
 16
 17
 18
 19
 20
 21
 22
 23
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110

// Copyright (c) 2017-present PyO3 Project and Contributors

//! Free allocation list

use crate::ffi;
use crate::type_object::{pytype_drop, PyObjectAlloc, PyTypeInfo};
use crate::Python;
use std::mem;
use std::os::raw::c_void;

/// 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,
            split: 0,
            capacity,
        }
    }

    /// Pop first non empty item
    pub fn pop(&mut self) -> Option<T> {
        let idx = self.split;
        if idx == 0 {
            None
        } else {
            match 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 for T
where
    T: PyObjectWithFreeList,
{
    unsafe fn alloc(_py: Python) -> *mut ffi::PyObject {
        if let Some(obj) = <Self as PyObjectWithFreeList>::get_free_list().pop() {
            ffi::PyObject_Init(obj, <Self as PyTypeInfo>::type_object());
            obj
        } else {
            ffi::PyType_GenericAlloc(<Self as PyTypeInfo>::type_object(), 0)
        }
    }

    unsafe fn dealloc(py: Python, obj: *mut ffi::PyObject) {
        pytype_drop::<Self>(py, obj);

        if ffi::PyObject_CallFinalizerFromDealloc(obj) < 0 {
            return;
        }

        if let Some(obj) = <Self as PyObjectWithFreeList>::get_free_list().insert(obj) {
            match Self::type_object().tp_free {
                Some(free) => free(obj as *mut c_void),
                None => {
                    let ty = ffi::Py_TYPE(obj);
                    if ffi::PyType_IS_GC(ty) != 0 {
                        ffi::PyObject_GC_Del(obj as *mut c_void);
                    } else {
                        ffi::PyObject_Free(obj as *mut c_void);
                    }

                    // For heap types, PyType_GenericAlloc calls INCREF on the type objects,
                    // so we need to call DECREF here:
                    if ffi::PyType_HasFeature(ty, ffi::Py_TPFLAGS_HEAPTYPE) != 0 {
                        ffi::Py_DECREF(ty as *mut ffi::PyObject);
                    }
                }
            }
        }
    }
}