/
pyalloc.rs
141 lines (113 loc) · 4.34 KB
/
pyalloc.rs
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
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
// This Source Code Form is subject to the terms of the Mozilla Public
// License, v. 2.0. If a copy of the MPL was not distributed with this
// file, You can obtain one at https://mozilla.org/MPL/2.0/.
use libc::{c_void, size_t};
use std::alloc;
use std::collections::HashMap;
use std::ptr::null_mut;
const MIN_ALIGN: usize = 16;
type RawAllocatorState = HashMap<*mut u8, alloc::Layout>;
/// Holds state for the raw memory allocator.
///
/// Ideally we wouldn't need to track state. But Rust's dealloc() API
/// requires passing in a Layout that matches the allocation. This means
/// we need to track the Layout for each allocation. This data structure
/// facilitates that.
///
/// TODO HashMap isn't thread safe and the Python raw allocator doesn't
/// hold the GIL. So we need a thread safe map or a mutex guarding access.
pub struct RawAllocator {
pub allocator: pyffi::PyMemAllocatorEx,
_state: Box<RawAllocatorState>,
}
extern "C" fn raw_malloc(ctx: *mut c_void, size: size_t) -> *mut c_void {
// PyMem_RawMalloc()'s docs say: Requesting zero bytes returns a distinct
// non-NULL pointer if possible, as if PyMem_RawMalloc(1) had been called
// instead.
let size = match size {
0 => 1,
val => val,
};
unsafe {
let state = ctx as *mut RawAllocatorState;
let layout = alloc::Layout::from_size_align_unchecked(size, MIN_ALIGN);
let res = alloc::alloc(layout);
(*state).insert(res, layout);
//println!("allocated {} bytes to {:?}", size, res);
res as *mut c_void
}
}
extern "C" fn raw_calloc(ctx: *mut c_void, nelem: size_t, elsize: size_t) -> *mut c_void {
// PyMem_RawCalloc()'s docs say: Requesting zero elements or elements of
// size zero bytes returns a distinct non-NULL pointer if possible, as if
// PyMem_RawCalloc(1, 1) had been called instead.
let size = match nelem * elsize {
0 => 1,
val => val,
};
unsafe {
let state = ctx as *mut RawAllocatorState;
let layout = alloc::Layout::from_size_align_unchecked(size, MIN_ALIGN);
let res = alloc::alloc_zeroed(layout);
(*state).insert(res, layout);
//println!("zero allocated {} bytes to {:?}", size, res);
res as *mut c_void
}
}
extern "C" fn raw_realloc(ctx: *mut c_void, ptr: *mut c_void, new_size: size_t) -> *mut c_void {
//println!("reallocating {:?} to {} bytes", ptr as *mut u8, new_size);
// PyMem_RawRealloc()'s docs say: If p is NULL, the call is equivalent to
// PyMem_RawMalloc(n); else if n is equal to zero, the memory block is
// resized but is not freed, and the returned pointer is non-NULL.
if ptr == null_mut() {
return raw_malloc(ctx, new_size);
}
let new_size = match new_size {
0 => 1,
val => val,
};
unsafe {
let state = ctx as *mut RawAllocatorState;
let layout = alloc::Layout::from_size_align_unchecked(new_size, MIN_ALIGN);
let key = ptr as *mut u8;
let old_layout = (*state)
.remove(&key)
.expect("original memory address not tracked");
let res = alloc::realloc(ptr as *mut u8, old_layout, new_size);
(*state).insert(res, layout);
res as *mut c_void
}
}
extern "C" fn raw_free(ctx: *mut c_void, ptr: *mut c_void) -> () {
if ptr == null_mut() {
return;
}
//println!("freeing {:?}", ptr as *mut u8);
unsafe {
let state = ctx as *mut RawAllocatorState;
let key = ptr as *mut u8;
let layout = (*state)
.get(&key)
.expect(format!("could not find allocated memory record: {:?}", key).as_str());
alloc::dealloc(key, *layout);
(*state).remove(&key);
}
}
pub fn make_raw_memory_allocator() -> RawAllocator {
// We need to allocate the HashMap on the heap so the pointer doesn't refer
// to the stack. We rebox and add the Box to our struct so lifetimes are
// managed.
let alloc = Box::new(HashMap::<*mut u8, alloc::Layout>::new());
let state = Box::into_raw(alloc);
let allocator = pyffi::PyMemAllocatorEx {
ctx: state as *mut c_void,
malloc: Some(raw_malloc),
calloc: Some(raw_calloc),
realloc: Some(raw_realloc),
free: Some(raw_free),
};
RawAllocator {
allocator,
_state: unsafe { Box::from_raw(state) },
}
}