map_box/
lib.rs

1use std::alloc::{handle_alloc_error, Layout};
2use std::mem::MaybeUninit;
3use std::{alloc, ptr};
4
5pub trait Map<T1> {
6    type Target<T2>;
7
8    /// Returns a box, with function f applied to the value inside.
9    /// This function will re-use the allocation when possible.
10    fn map_box<T2>(self, f: impl FnMut(T1) -> T2) -> Self::Target<T2>;
11}
12
13impl<T1> Map<T1> for Box<T1> {
14    type Target<T2> = Box<T2>;
15
16    /// Returns a box, with function f applied to the value inside.
17    /// This function will re-use the allocation when possible.
18    fn map_box<T2>(self, mut f: impl FnMut(T1) -> T2) -> Self::Target<T2> {
19        // Get layouts of the types
20        let from_layout = Layout::new::<T1>();
21        let to_layout = Layout::new::<T2>();
22
23        // If `T1` or `T2` is a ZST, we let `Box` handle the logic for us. This is optimal.
24        // If the alignment requirements of T1 and T2 are different, we also call this.
25        // This is because `dealloc` requires the alignment to be identical, and there's no way to realloc with a different Layout
26        if from_layout.size() == 0
27            || to_layout.size() == 0
28            || from_layout.align() != to_layout.align()
29        {
30            return Box::new(f(*self));
31        }
32
33        // Safety: Read T1, safe since `from_ptr` was created from a Box<T1>
34        let from_ptr = Box::into_raw(self);
35        let v = unsafe { ptr::read(from_ptr) };
36
37        // Apply `f`. Create a temporary Box so the allocation of the Box is deallocated on panic of `f`
38        let tmp_box: Box<MaybeUninit<T1>> =
39            unsafe { Box::from_raw(from_ptr as *mut MaybeUninit<T1>) };
40        let v = f(v);
41        Box::into_raw(tmp_box);
42
43        // Generate a `to_ptr` from `from_ptr` that can fit a `T2`
44        let to_ptr = if to_layout.size() != from_layout.size() {
45            // We need to re-allocate, because the size of the Box is incorrect
46            let to_ptr = unsafe {
47                // Safety: from_layout was used to allocate the Box and to_layout is non-zero
48                alloc::realloc(from_ptr as *mut u8, from_layout, to_layout.size()) as *mut T2
49            };
50            // Realloc will return null in case of an alloc error, we need to check for this
51            if to_ptr.is_null() {
52                // We still need to dealloc the old box
53                // Safety: from_ptr is still valid if `realloc` returns null
54                unsafe {
55                    alloc::dealloc(from_ptr as *mut u8, from_layout);
56                }
57                handle_alloc_error(to_layout)
58            }
59            to_ptr
60        } else {
61            // Size and alignment are correct, so we can re-use the allocation
62            from_ptr as *mut T2
63        };
64
65        unsafe {
66            // Safety: The logic above guarantees that to_ptr can fit a `T2`
67            ptr::write(to_ptr, v);
68            Box::from_raw(to_ptr)
69        }
70    }
71}
72
73#[cfg(test)]
74mod tests {
75    use crate::Map;
76
77    #[test]
78    fn same_type() {
79        let b = Box::new(21u64);
80        let b = b.map_box(|v| v * 2);
81        assert_eq!(*b, 42);
82    }
83
84    #[test]
85    fn same_size() {
86        let b = Box::new(42u64);
87        let b = b.map_box(|v| v as i64);
88        assert_eq!(*b, 42);
89    }
90
91    #[test]
92    fn down_size() {
93        let b = Box::new(42u64);
94        let b = b.map_box(|v| v as u32);
95        assert_eq!(*b, 42);
96    }
97
98    #[test]
99    fn up_size() {
100        let b = Box::new(42u32);
101        let b = b.map_box(|v| v as u64);
102        assert_eq!(*b, 42);
103    }
104
105    #[test]
106    fn zst_in() {
107        let b = Box::new(());
108        let b = b.map_box(|_| 42u64);
109        assert_eq!(*b, 42);
110    }
111
112    #[test]
113    fn zst_out() {
114        let b = Box::new(42u64);
115        let b = b.map_box(|_| ());
116        assert_eq!(*b, ());
117    }
118
119    #[test]
120    fn zst_both() {
121        let b = Box::new(());
122        let b: Box<[u64; 0]> = b.map_box(|_| []);
123        assert_eq!(*b, []);
124    }
125
126    #[test]
127    #[should_panic]
128    fn panic() {
129        let b = Box::new(42u64);
130        b.map_box(|_| panic!());
131    }
132
133    #[test]
134    fn align_up_high() {
135        #[repr(align(16))]
136        struct Test1(u64);
137        #[repr(align(128))]
138        struct Test2(u64);
139
140        let b = Box::new(Test1(42));
141        let b = b.map_box(|Test1(v)| Test2(v));
142        assert_eq!(b.0, 42);
143    }
144
145    #[test]
146    fn align_down_high() {
147        #[repr(align(128))]
148        struct Test1(u64);
149        #[repr(align(16))]
150        struct Test2(u64);
151
152        let b = Box::new(Test1(42));
153        let b = b.map_box(|Test1(v)| Test2(v));
154        assert_eq!(b.0, 42);
155    }
156
157    #[test]
158    fn align_up_low() {
159        #[repr(align(1))]
160        struct Test1(u64);
161        #[repr(align(8))]
162        struct Test2(u64);
163
164        let b = Box::new(Test1(42));
165        let b = b.map_box(|Test1(v)| Test2(v));
166        assert_eq!(b.0, 42);
167    }
168
169    #[test]
170    fn align_down_low() {
171        #[repr(align(8))]
172        struct Test1(u64);
173        #[repr(align(1))]
174        struct Test2(u64);
175
176        let b = Box::new(Test1(42));
177        let b = b.map_box(|Test1(v)| Test2(v));
178        assert_eq!(b.0, 42);
179    }
180}
map_box/lib.rs

map_box/
lib.rs
