pyo3_polars/
alloc.rs

1use std::alloc::{GlobalAlloc, Layout, System};
2use std::ffi::c_char;
3
4use once_cell::race::OnceRef;
5use pyo3::ffi::{PyCapsule_Import, Py_IsInitialized};
6use pyo3::Python;
7
8unsafe extern "C" fn fallback_alloc(size: usize, align: usize) -> *mut u8 {
9    System.alloc(Layout::from_size_align_unchecked(size, align))
10}
11
12unsafe extern "C" fn fallback_dealloc(ptr: *mut u8, size: usize, align: usize) {
13    System.dealloc(ptr, Layout::from_size_align_unchecked(size, align))
14}
15
16unsafe extern "C" fn fallback_alloc_zeroed(size: usize, align: usize) -> *mut u8 {
17    System.alloc_zeroed(Layout::from_size_align_unchecked(size, align))
18}
19
20unsafe extern "C" fn fallback_realloc(
21    ptr: *mut u8,
22    size: usize,
23    align: usize,
24    new_size: usize,
25) -> *mut u8 {
26    System.realloc(
27        ptr,
28        Layout::from_size_align_unchecked(size, align),
29        new_size,
30    )
31}
32
33#[repr(C)]
34struct AllocatorCapsule {
35    alloc: unsafe extern "C" fn(usize, usize) -> *mut u8,
36    dealloc: unsafe extern "C" fn(*mut u8, usize, usize),
37    alloc_zeroed: unsafe extern "C" fn(usize, usize) -> *mut u8,
38    realloc: unsafe extern "C" fn(*mut u8, usize, usize, usize) -> *mut u8,
39}
40
41static FALLBACK_ALLOCATOR_CAPSULE: AllocatorCapsule = AllocatorCapsule {
42    alloc: fallback_alloc,
43    alloc_zeroed: fallback_alloc_zeroed,
44    dealloc: fallback_dealloc,
45    realloc: fallback_realloc,
46};
47
48static ALLOCATOR_CAPSULE_NAME: &[u8] = b"polars.polars._allocator\0";
49
50/// A memory allocator that relays allocations to the allocator used by Polars.
51///
52/// You can use it as the global memory allocator:
53///
54/// ```rust
55/// use pyo3_polars::PolarsAllocator;
56///
57/// #[global_allocator]
58/// static ALLOC: PolarsAllocator = PolarsAllocator::new();
59/// ```
60///
61/// If the allocator capsule (`polars.polars._allocator`) is not available,
62/// this allocator fallbacks to [`std::alloc::System`].
63pub struct PolarsAllocator(OnceRef<'static, AllocatorCapsule>);
64
65impl PolarsAllocator {
66    fn get_allocator(&self) -> &'static AllocatorCapsule {
67        // Do not allocate in this function,
68        // otherwise it will cause infinite recursion.
69        self.0.get_or_init(|| {
70            let r = (unsafe { Py_IsInitialized() } != 0)
71                .then(|| {
72                    Python::with_gil(|_| unsafe {
73                        (PyCapsule_Import(ALLOCATOR_CAPSULE_NAME.as_ptr() as *const c_char, 0)
74                            as *const AllocatorCapsule)
75                            .as_ref()
76                    })
77                })
78                .flatten();
79            #[cfg(debug_assertions)]
80            if r.is_none() {
81                // Do not use eprintln; it may alloc.
82                let msg = b"failed to get allocator capsule\n";
83                // Message length type is platform-dependent.
84                let msg_len = msg.len().try_into().unwrap();
85                unsafe { libc::write(2, msg.as_ptr() as *const libc::c_void, msg_len) };
86            }
87            r.unwrap_or(&FALLBACK_ALLOCATOR_CAPSULE)
88        })
89    }
90
91    /// Create a `PolarsAllocator`.
92    pub const fn new() -> Self {
93        PolarsAllocator(OnceRef::new())
94    }
95}
96
97impl Default for PolarsAllocator {
98    fn default() -> Self {
99        Self::new()
100    }
101}
102
103unsafe impl GlobalAlloc for PolarsAllocator {
104    #[inline]
105    unsafe fn alloc(&self, layout: Layout) -> *mut u8 {
106        (self.get_allocator().alloc)(layout.size(), layout.align())
107    }
108
109    #[inline]
110    unsafe fn dealloc(&self, ptr: *mut u8, layout: Layout) {
111        (self.get_allocator().dealloc)(ptr, layout.size(), layout.align());
112    }
113
114    #[inline]
115    unsafe fn alloc_zeroed(&self, layout: Layout) -> *mut u8 {
116        (self.get_allocator().alloc_zeroed)(layout.size(), layout.align())
117    }
118
119    #[inline]
120    unsafe fn realloc(&self, ptr: *mut u8, layout: Layout, new_size: usize) -> *mut u8 {
121        (self.get_allocator().realloc)(ptr, layout.size(), layout.align(), new_size)
122    }
123}
pyo3_polars/alloc.rs

pyo3_polars/
alloc.rs
