You can not select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
145 lines
5.2 KiB
145 lines
5.2 KiB
# cython: embedsignature=True
|
|
|
|
from cpython.mem cimport PyMem_Malloc, PyMem_Free
|
|
from cpython.ref cimport Py_INCREF, Py_DECREF
|
|
from libc.string cimport memset
|
|
from libc.string cimport memcpy
|
|
|
|
cdef class PyMalloc:
|
|
cdef void _set(self, malloc_t malloc):
|
|
self.malloc = malloc
|
|
|
|
cdef PyMalloc WrapMalloc(malloc_t malloc):
|
|
cdef PyMalloc o = PyMalloc()
|
|
o._set(malloc)
|
|
return o
|
|
|
|
cdef class PyFree:
|
|
cdef void _set(self, free_t free):
|
|
self.free = free
|
|
|
|
cdef PyFree WrapFree(free_t free):
|
|
cdef PyFree o = PyFree()
|
|
o._set(free)
|
|
return o
|
|
|
|
Default_Malloc = WrapMalloc(PyMem_Malloc)
|
|
Default_Free = WrapFree(PyMem_Free)
|
|
|
|
cdef class Pool:
|
|
"""Track allocated memory addresses, and free them all when the Pool is
|
|
garbage collected. This provides an easy way to avoid memory leaks, and
|
|
removes the need for deallocation functions for complicated structs.
|
|
|
|
>>> from cymem.cymem cimport Pool
|
|
>>> cdef Pool mem = Pool()
|
|
>>> data1 = <int*>mem.alloc(10, sizeof(int))
|
|
>>> data2 = <float*>mem.alloc(12, sizeof(float))
|
|
|
|
Attributes:
|
|
size (size_t): The current size (in bytes) allocated by the pool.
|
|
addresses (dict): The currently allocated addresses and their sizes. Read-only.
|
|
pymalloc (PyMalloc): The allocator to use (default uses PyMem_Malloc).
|
|
pyfree (PyFree): The free to use (default uses PyMem_Free).
|
|
"""
|
|
|
|
def __cinit__(self, PyMalloc pymalloc=Default_Malloc,
|
|
PyFree pyfree=Default_Free):
|
|
self.size = 0
|
|
self.addresses = {}
|
|
self.refs = []
|
|
self.pymalloc = pymalloc
|
|
self.pyfree = pyfree
|
|
|
|
def __dealloc__(self):
|
|
cdef size_t addr
|
|
if self.addresses is not None:
|
|
for addr in self.addresses:
|
|
if addr != 0:
|
|
self.pyfree.free(<void*>addr)
|
|
|
|
cdef void* alloc(self, size_t number, size_t elem_size) except NULL:
|
|
"""Allocate a 0-initialized number*elem_size-byte block of memory, and
|
|
remember its address. The block will be freed when the Pool is garbage
|
|
collected.
|
|
"""
|
|
cdef void* p = self.pymalloc.malloc(number * elem_size)
|
|
if p == NULL:
|
|
raise MemoryError("Error assigning %d bytes" % (number * elem_size))
|
|
memset(p, 0, number * elem_size)
|
|
self.addresses[<size_t>p] = number * elem_size
|
|
self.size += number * elem_size
|
|
return p
|
|
|
|
cdef void* realloc(self, void* p, size_t new_size) except NULL:
|
|
"""Resizes the memory block pointed to by p to new_size bytes, returning
|
|
a non-NULL pointer to the new block. new_size must be larger than the
|
|
original.
|
|
|
|
If p is not in the Pool or new_size is 0, a MemoryError is raised.
|
|
"""
|
|
if <size_t>p not in self.addresses:
|
|
raise ValueError("Pointer %d not found in Pool %s" % (<size_t>p, self.addresses))
|
|
if new_size == 0:
|
|
raise ValueError("Realloc requires new_size > 0")
|
|
assert new_size > self.addresses[<size_t>p]
|
|
cdef void* new_ptr = self.alloc(1, new_size)
|
|
if new_ptr == NULL:
|
|
raise MemoryError("Error reallocating to %d bytes" % new_size)
|
|
memcpy(new_ptr, p, self.addresses[<size_t>p])
|
|
self.free(p)
|
|
self.addresses[<size_t>new_ptr] = new_size
|
|
return new_ptr
|
|
|
|
cdef void free(self, void* p) except *:
|
|
"""Frees the memory block pointed to by p, which must have been returned
|
|
by a previous call to Pool.alloc. You don't necessarily need to free
|
|
memory addresses manually --- you can instead let the Pool be garbage
|
|
collected, at which point all the memory will be freed.
|
|
|
|
If p is not in Pool.addresses, a KeyError is raised.
|
|
"""
|
|
self.size -= self.addresses.pop(<size_t>p)
|
|
self.pyfree.free(p)
|
|
|
|
def own_pyref(self, object py_ref):
|
|
self.refs.append(py_ref)
|
|
|
|
|
|
cdef class Address:
|
|
"""A block of number * size-bytes of 0-initialized memory, tied to a Python
|
|
ref-counted object. When the object is garbage collected, the memory is freed.
|
|
|
|
>>> from cymem.cymem cimport Address
|
|
>>> cdef Address address = Address(10, sizeof(double))
|
|
>>> d10 = <double*>address.ptr
|
|
|
|
Args:
|
|
number (size_t): The number of elements in the memory block.
|
|
elem_size (size_t): The size of each element.
|
|
|
|
Attributes:
|
|
ptr (void*): Pointer to the memory block.
|
|
addr (size_t): Read-only size_t cast of the pointer.
|
|
pymalloc (PyMalloc): The allocator to use (default uses PyMem_Malloc).
|
|
pyfree (PyFree): The free to use (default uses PyMem_Free).
|
|
"""
|
|
def __cinit__(self, size_t number, size_t elem_size,
|
|
PyMalloc pymalloc=Default_Malloc, PyFree pyfree=Default_Free):
|
|
self.ptr = NULL
|
|
self.pymalloc = pymalloc
|
|
self.pyfree = pyfree
|
|
|
|
def __init__(self, size_t number, size_t elem_size):
|
|
self.ptr = self.pymalloc.malloc(number * elem_size)
|
|
if self.ptr == NULL:
|
|
raise MemoryError("Error assigning %d bytes" % number * elem_size)
|
|
memset(self.ptr, 0, number * elem_size)
|
|
|
|
property addr:
|
|
def __get__(self):
|
|
return <size_t>self.ptr
|
|
|
|
def __dealloc__(self):
|
|
if self.ptr != NULL:
|
|
self.pyfree.free(self.ptr)
|