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

4 years ago
  1. # cython: embedsignature=True
  2. from cpython.mem cimport PyMem_Malloc, PyMem_Free
  3. from cpython.ref cimport Py_INCREF, Py_DECREF
  4. from libc.string cimport memset
  5. from libc.string cimport memcpy
  6. cdef class PyMalloc:
  7. cdef void _set(self, malloc_t malloc):
  8. self.malloc = malloc
  9. cdef PyMalloc WrapMalloc(malloc_t malloc):
  10. cdef PyMalloc o = PyMalloc()
  11. o._set(malloc)
  12. return o
  13. cdef class PyFree:
  14. cdef void _set(self, free_t free):
  15. self.free = free
  16. cdef PyFree WrapFree(free_t free):
  17. cdef PyFree o = PyFree()
  18. o._set(free)
  19. return o
  20. Default_Malloc = WrapMalloc(PyMem_Malloc)
  21. Default_Free = WrapFree(PyMem_Free)
  22. cdef class Pool:
  23. """Track allocated memory addresses, and free them all when the Pool is
  24. garbage collected. This provides an easy way to avoid memory leaks, and
  25. removes the need for deallocation functions for complicated structs.
  26. >>> from cymem.cymem cimport Pool
  27. >>> cdef Pool mem = Pool()
  28. >>> data1 = <int*>mem.alloc(10, sizeof(int))
  29. >>> data2 = <float*>mem.alloc(12, sizeof(float))
  30. Attributes:
  31. size (size_t): The current size (in bytes) allocated by the pool.
  32. addresses (dict): The currently allocated addresses and their sizes. Read-only.
  33. pymalloc (PyMalloc): The allocator to use (default uses PyMem_Malloc).
  34. pyfree (PyFree): The free to use (default uses PyMem_Free).
  35. """
  36. def __cinit__(self, PyMalloc pymalloc=Default_Malloc,
  37. PyFree pyfree=Default_Free):
  38. self.size = 0
  39. self.addresses = {}
  40. self.refs = []
  41. self.pymalloc = pymalloc
  42. self.pyfree = pyfree
  43. def __dealloc__(self):
  44. cdef size_t addr
  45. if self.addresses is not None:
  46. for addr in self.addresses:
  47. if addr != 0:
  48. self.pyfree.free(<void*>addr)
  49. cdef void* alloc(self, size_t number, size_t elem_size) except NULL:
  50. """Allocate a 0-initialized number*elem_size-byte block of memory, and
  51. remember its address. The block will be freed when the Pool is garbage
  52. collected.
  53. """
  54. cdef void* p = self.pymalloc.malloc(number * elem_size)
  55. if p == NULL:
  56. raise MemoryError("Error assigning %d bytes" % (number * elem_size))
  57. memset(p, 0, number * elem_size)
  58. self.addresses[<size_t>p] = number * elem_size
  59. self.size += number * elem_size
  60. return p
  61. cdef void* realloc(self, void* p, size_t new_size) except NULL:
  62. """Resizes the memory block pointed to by p to new_size bytes, returning
  63. a non-NULL pointer to the new block. new_size must be larger than the
  64. original.
  65. If p is not in the Pool or new_size is 0, a MemoryError is raised.
  66. """
  67. if <size_t>p not in self.addresses:
  68. raise ValueError("Pointer %d not found in Pool %s" % (<size_t>p, self.addresses))
  69. if new_size == 0:
  70. raise ValueError("Realloc requires new_size > 0")
  71. assert new_size > self.addresses[<size_t>p]
  72. cdef void* new_ptr = self.alloc(1, new_size)
  73. if new_ptr == NULL:
  74. raise MemoryError("Error reallocating to %d bytes" % new_size)
  75. memcpy(new_ptr, p, self.addresses[<size_t>p])
  76. self.free(p)
  77. self.addresses[<size_t>new_ptr] = new_size
  78. return new_ptr
  79. cdef void free(self, void* p) except *:
  80. """Frees the memory block pointed to by p, which must have been returned
  81. by a previous call to Pool.alloc. You don't necessarily need to free
  82. memory addresses manually --- you can instead let the Pool be garbage
  83. collected, at which point all the memory will be freed.
  84. If p is not in Pool.addresses, a KeyError is raised.
  85. """
  86. self.size -= self.addresses.pop(<size_t>p)
  87. self.pyfree.free(p)
  88. def own_pyref(self, object py_ref):
  89. self.refs.append(py_ref)
  90. cdef class Address:
  91. """A block of number * size-bytes of 0-initialized memory, tied to a Python
  92. ref-counted object. When the object is garbage collected, the memory is freed.
  93. >>> from cymem.cymem cimport Address
  94. >>> cdef Address address = Address(10, sizeof(double))
  95. >>> d10 = <double*>address.ptr
  96. Args:
  97. number (size_t): The number of elements in the memory block.
  98. elem_size (size_t): The size of each element.
  99. Attributes:
  100. ptr (void*): Pointer to the memory block.
  101. addr (size_t): Read-only size_t cast of the pointer.
  102. pymalloc (PyMalloc): The allocator to use (default uses PyMem_Malloc).
  103. pyfree (PyFree): The free to use (default uses PyMem_Free).
  104. """
  105. def __cinit__(self, size_t number, size_t elem_size,
  106. PyMalloc pymalloc=Default_Malloc, PyFree pyfree=Default_Free):
  107. self.ptr = NULL
  108. self.pymalloc = pymalloc
  109. self.pyfree = pyfree
  110. def __init__(self, size_t number, size_t elem_size):
  111. self.ptr = self.pymalloc.malloc(number * elem_size)
  112. if self.ptr == NULL:
  113. raise MemoryError("Error assigning %d bytes" % number * elem_size)
  114. memset(self.ptr, 0, number * elem_size)
  115. property addr:
  116. def __get__(self):
  117. return <size_t>self.ptr
  118. def __dealloc__(self):
  119. if self.ptr != NULL:
  120. self.pyfree.free(self.ptr)