gg_alloc/
lib.rs

1use std::alloc::{GlobalAlloc, Layout};
2
3pub struct GgAlloc<A> {
4    alloc: A,
5}
6
7impl<A> GgAlloc<A> {
8    pub const fn new(alloc: A) -> Self {
9        Self { alloc }
10    }
11}
12
13fn pointer_above_2g(ptr: *mut u8) -> bool {
14    //(ptr as usize) > (isize::MAX as usize)
15    (ptr as u32) > (i32::MAX as u32)
16}
17
18fn alloc_fully_below_2g(ptr: *mut u8, layout: Layout) -> bool {
19    !pointer_above_2g(ptr) && !pointer_above_2g(unsafe { ptr.add(layout.size() - 1) })
20}
21
22unsafe impl<A> GlobalAlloc for GgAlloc<A>
23where
24    A: GlobalAlloc,
25{
26    unsafe fn alloc(&self, layout: Layout) -> *mut u8 {
27        let mut ret = self.alloc.alloc(layout);
28
29        loop {
30            if ret.is_null() {
31                // Failed to allocate, return null pointer
32                break;
33            }
34
35            if pointer_above_2g(ret) {
36                // Allocated above 2G, return this pointer
37                break;
38            }
39
40            // Here the allocation was successful but below 2G, so fill memory a bit
41            // in chunks of 128 MB
42            let mut size = 1 << 27;
43            loop {
44                let test_layout = Layout::from_size_align(size, 1).unwrap();
45                let fill_ptr = self.alloc.alloc(test_layout);
46                if !fill_ptr.is_null() && alloc_fully_below_2g(fill_ptr, test_layout) {
47                    // leak allocation and continue
48                } else {
49                    // Free and try again with half size
50                    if !fill_ptr.is_null() {
51                        self.alloc.dealloc(fill_ptr, test_layout);
52                    }
53                    size /= 2;
54
55                    if size < 1 {
56                        // Exit fill loop if size is small
57                        break;
58                    }
59                }
60            }
61
62            if !alloc_fully_below_2g(ret, layout) {
63                // Free and allocate something with half size
64                self.alloc.dealloc(ret, layout);
65            }
66
67            // fill memory a bit
68            // in chunks of same size as the requested allocation
69            loop {
70                let test_layout = layout;
71                let fill_ptr = self.alloc.alloc(test_layout);
72                if !fill_ptr.is_null() && alloc_fully_below_2g(fill_ptr, test_layout) {
73                    // leak allocation and continue
74                } else {
75                    // this allocation is either above 2GB or right at the 2GB edge, we are done
76                    if pointer_above_2g(fill_ptr) {
77                        // above 2GB, we can use it as the return pointer
78                        return fill_ptr;
79                    } else {
80                        // right at the edge, leak this allocation and try again in next iteration
81                        // of the outer loop
82                        break;
83                    }
84                }
85            }
86
87            ret = self.alloc.alloc(layout);
88        }
89
90        ret
91    }
92    unsafe fn dealloc(&self, ptr: *mut u8, layout: Layout) {
93        self.alloc.dealloc(ptr, layout);
94    }
95}
96
97#[cfg(test)]
98mod tests {
99    use super::*;
100    use std::alloc::System;
101
102    #[global_allocator]
103    static A: GgAlloc<System> = GgAlloc::new(System);
104
105    fn assert_above_2g(ptr: *mut u8) {
106        assert!(pointer_above_2g(ptr), "{:p}", ptr,);
107    }
108
109    #[test]
110    fn alloc_string() {
111        let s = format!("allocating a string!");
112        assert_above_2g(s.as_str() as *const _ as *mut _);
113    }
114
115    #[test]
116    fn alloc_one_byte() {
117        let s = format!("1");
118        assert!(
119            pointer_above_2g(s.as_str() as *const _ as *mut _),
120            "{:p}",
121            s.as_str()
122        );
123    }
124
125    #[test]
126    fn alloc_many_bytes() {
127        let s = format!("1");
128        assert_above_2g(s.as_str() as *const _ as *mut _);
129        let s = format!("12");
130        assert_above_2g(s.as_str() as *const _ as *mut _);
131        let s = format!("123");
132        assert_above_2g(s.as_str() as *const _ as *mut _);
133        let s = format!("1234");
134        assert_above_2g(s.as_str() as *const _ as *mut _);
135        let s = format!("12345");
136        assert_above_2g(s.as_str() as *const _ as *mut _);
137        let s = format!("123456");
138        assert_above_2g(s.as_str() as *const _ as *mut _);
139        let s = format!("1234567");
140        assert_above_2g(s.as_str() as *const _ as *mut _);
141        let s = format!("12345678");
142        assert_above_2g(s.as_str() as *const _ as *mut _);
143        let s = format!("123456789");
144        assert_above_2g(s.as_str() as *const _ as *mut _);
145    }
146}