inplace_box/
impl_fn_traits.rs

1//! Implementation of `Fn*` traits for `InplaceBox`.
2//!
3//! The `FnOnce` implementation uses a dummy allocator to temporarily create a
4//! `Box` without actual memory allocation, allowing safe delegation to
5//! `Box<F>`'s `FnOnce`.
6//!
7//! ## Why the intermediate `Box` is needed
8//!
9//! When `F` is a trait object like `dyn FnOnce<Args>`, we can't call
10//! `F::call_once()` directly because `FnOnce::call_once` takes `self` by value,
11//! requiring the compiler to know the size at compile time. But trait objects
12//! are unsized types - their concrete size is only known at runtime.
13//!
14//! The `FnOnce` trait for trait objects is implemented for `Box<dyn
15//! FnOnce<Args>>`, which handles the complex dynamic dispatch and destruction
16//! logic. Rather than reimplementing this, we create a temporary `Box` with a
17//! dummy allocator to reuse the existing, well-tested implementation.
18
19use alloc::boxed::Box;
20use core::alloc::AllocError;
21use core::alloc::Allocator;
22use core::alloc::Layout;
23use core::marker::Tuple;
24use core::ptr::NonNull;
25
26use crate::InplaceBox;
27
28impl<Args: Tuple, F: FnOnce<Args> + ?Sized, const SIZE: usize> FnOnce<Args>
29    for InplaceBox<F, SIZE>
30{
31    type Output = <F as FnOnce<Args>>::Output;
32
33    #[inline]
34    extern "rust-call" fn call_once(mut self, args: Args) -> Self::Output {
35        // SAFETY: Create a temporary `Box` with a dummy allocator that doesn't
36        // actually deallocate. The original `InplaceBox` is forgotten to
37        // prevent double-drop.
38        let b = unsafe {
39            Box::from_raw_in(self.as_mut(), InplaceBoxFnOnceDummyAllocator)
40        };
41        core::mem::forget(self); // the inner object is destroyed by `call_once` below
42        <Box<_, _> as FnOnce<Args>>::call_once(b, args)
43    }
44}
45
46/// Dummy allocator for `InplaceBox::call_once` that never allocates or
47/// deallocates.
48struct InplaceBoxFnOnceDummyAllocator;
49
50// SAFETY: This allocator is unsafe. It is only to be used with `FnOnce` for
51// `InplaceBox`.
52unsafe impl Allocator for InplaceBoxFnOnceDummyAllocator {
53    #[inline]
54    fn allocate(&self, _layout: Layout) -> Result<NonNull<[u8]>, AllocError> {
55        Err(AllocError) // in fact, never called
56    }
57
58    #[inline]
59    unsafe fn deallocate(&self, _ptr: NonNull<u8>, _layout: Layout) {
60        // No-op: memory owned by InplaceBox
61    }
62}
63
64impl<Args: Tuple, F: FnMut<Args> + ?Sized, const SIZE: usize> FnMut<Args>
65    for InplaceBox<F, SIZE>
66{
67    extern "rust-call" fn call_mut(&mut self, args: Args) -> Self::Output {
68        <F as FnMut<Args>>::call_mut(self, args)
69    }
70}
71
72impl<Args: Tuple, F: Fn<Args> + ?Sized, const SIZE: usize> Fn<Args>
73    for InplaceBox<F, SIZE>
74{
75    extern "rust-call" fn call(&self, args: Args) -> Self::Output {
76        <F as Fn<Args>>::call(self, args)
77    }
78}
79
80#[cfg(test)]
81mod tests {
82    use super::*;
83
84    #[test]
85    fn fn_once() {
86        let mut counter = 1;
87        let adder: InplaceBox<dyn FnOnce<(usize,), Output = usize>, 32> =
88            InplaceBox::new(|count| {
89                let res = counter;
90                counter = res + count;
91                res
92            });
93        // call the function once
94        let prev_value = adder(4);
95        // previous count was 1
96        assert_eq!(prev_value, 1);
97        // we added 4, so now it's 5
98        assert_eq!(5, counter);
99
100        // let prev_value2 = adder(1); -- impossible - `FnOnce` call consumes
101        // the box
102    }
103
104    #[test]
105    fn fn_once_drop_or_call() {
106        struct Guard<'a>(&'a mut bool);
107        impl Drop for Guard<'_> {
108            fn drop(&mut self) {
109                *self.0 = true;
110            }
111        }
112
113        // first part - ensure that the closure is dropped, if the `FnOnce` is
114        // not called
115        let mut called = false;
116        let mut dropped = false;
117        {
118            let called = &mut called;
119            let guard = Guard(&mut dropped);
120            let b: InplaceBox<dyn FnOnce(), 32> = InplaceBox::new(move || {
121                *called = true;
122                core::mem::forget(guard);
123            });
124
125            drop(b); // drop w/o calling
126        }
127        assert!(!called);
128        assert!(dropped);
129
130        // second part - ensure that the closure is not dropped twice, the
131        // `FnOnce` call via `InplaceBox` drops it
132        called = false;
133        dropped = false;
134        {
135            let called = &mut called;
136            let guard = Guard(&mut dropped);
137            let b: InplaceBox<dyn FnOnce<(), Output = ()>, 32> =
138                InplaceBox::new(move || {
139                    *called = true;
140                    core::mem::forget(guard);
141                });
142
143            b(); // call it now
144        }
145        assert!(called);
146        assert!(!dropped);
147    }
148}