flat_drop/
lib.rs

1//! In this crate, we define the [FlatDrop] type.
2//! `FlatDrop<K>` behaves just like a `K`, but with a custom `Drop` implementation
3//! that avoids blowing up the stack when dropping large objects.
4//! Instead of recursively dropping subobjects, we perform a depth-first search
5//! and iteratively drop subobjects.
6//!
7//! To use this crate, you can replace recursive [Box]es and [Arc]s in your types
8//! with `FlatDrop<Box<T>>` or `FlatDrop<Arc<T>>`. You'll need to implement the
9//! [Recursive] trait for your type, which performs one step of the iterative
10//! dropping procedure.
11//!
12//! This crate uses `unsafe` internally, but the external API is safe.
13//!
14//! # Example
15//!
16//! ```
17//! use flat_drop::{FlatDrop, Recursive};
18//!
19//! /// Peano natural numbers.
20//! enum Natural {
21//!     Zero,
22//!     Succ(FlatDrop<Box<Natural>>),
23//! }
24//!
25//! impl Recursive for Natural {
26//!     type Container = Box<Natural>;
27//!
28//!     fn destruct(self) -> impl Iterator<Item = Self::Container> {
29//!         match self {
30//!             Natural::Zero => None,
31//!             Natural::Succ(pred) => Some(pred.into_inner()),
32//!         }
33//!         .into_iter()
34//!     }
35//! }
36//!
37//! impl Natural {
38//!     pub fn from_usize(value: usize) -> Self {
39//!         (0..value).fold(Self::Zero, |nat, _| {
40//!             Self::Succ(FlatDrop::new(Box::new(nat)))
41//!         })
42//!     }
43//! }
44//!
45//! // Create a new thread with a 4kb stack and allocate a number far bigger than 4 * 1024.
46//! const STACK_SIZE: usize = 4 * 1024;
47//!
48//! fn task() {
49//!     let nat = Natural::from_usize(STACK_SIZE * 100);
50//!     drop(std::hint::black_box(nat));
51//! }
52//!
53//! std::thread::Builder::new()
54//!     .stack_size(STACK_SIZE)
55//!     .spawn(task)
56//!     .unwrap()
57//!     .join()
58//!     .unwrap();
59//! ```
60
61use std::{
62    mem::ManuallyDrop,
63    ops::{Deref, DerefMut},
64    rc::Rc,
65    sync::Arc,
66};
67
68/// The [Recursive::destruct] function decomposes an object into some component parts.
69/// Usually, [Recursive::Container] is something like `Box<Self>` or `Arc<Self>`.
70pub trait Recursive {
71    type Container;
72
73    fn destruct(self) -> impl Iterator<Item = Self::Container>;
74}
75
76/// A trait for a smart pointer that contains (at most) a single value.
77pub trait IntoOptionInner {
78    type Inner;
79
80    /// A (potentially) fallible operation to convert the container into its internal value.
81    /// This should never drop any data.
82    ///
83    /// If `Self == Box`, this will always return `Some(*self)`.
84    /// If `Self == Arc`, this is `Arc::into_inner`.
85    fn into_option_inner(self) -> Option<Self::Inner>;
86}
87
88/// If `K` is a container of a recursive type, such as `Box<T>` where `T: Recursive`,
89/// `FlatDrop<K>` behaves just like `K`, but with a custom `Drop` implementation.
90/// In this implementation, we gather the recursive parts of the object iteratively
91/// and drop them without recursion, avoiding stack overflows when dropping
92/// large recursive objects.
93///
94/// # Safety
95///
96/// We keep the invariant that the inner object is always initialised, but will
97/// be dropped (exactly once) in the `drop` implementation.
98#[derive(Clone, Debug, Default, PartialEq, Eq, PartialOrd, Ord, Hash)]
99#[repr(transparent)]
100pub struct FlatDrop<K>(ManuallyDrop<K>)
101where
102    K: IntoOptionInner,
103    K::Inner: Recursive<Container = K>;
104
105impl<K> Drop for FlatDrop<K>
106where
107    K: IntoOptionInner,
108    K::Inner: Recursive<Container = K>,
109{
110    fn drop(&mut self) {
111        // Move out of the inner `ManuallyDrop`.
112        // Safety: the inner value has not yet been dropped, and will not be used again.
113        let value = unsafe { ManuallyDrop::take(&mut self.0) };
114
115        // Construct a sequence of containers to drop.
116        let mut to_drop = vec![value];
117
118        // Iteratively decompose each container from this list.
119        // This avoids creating excessive stack frames when destroying large objects.
120        while let Some(container) = to_drop.pop() {
121            if let Some(value) = container.into_option_inner() {
122                to_drop.extend(value.destruct());
123            }
124        }
125
126        // The drop glue will be a no-op since the field is `ManuallyDrop`.
127    }
128}
129
130// Now that we've defined the core parts of the library, we'll make some API.
131
132impl<T> IntoOptionInner for Box<T> {
133    type Inner = T;
134
135    fn into_option_inner(self) -> Option<Self::Inner> {
136        Some(*self)
137    }
138}
139
140impl<T> IntoOptionInner for Rc<T> {
141    type Inner = T;
142
143    fn into_option_inner(self) -> Option<Self::Inner> {
144        Rc::into_inner(self)
145    }
146}
147
148impl<T> IntoOptionInner for Arc<T> {
149    type Inner = T;
150
151    fn into_option_inner(self) -> Option<Self::Inner> {
152        Arc::into_inner(self)
153    }
154}
155
156impl<K> FlatDrop<K>
157where
158    K: IntoOptionInner,
159    K::Inner: Recursive<Container = K>,
160{
161    pub const fn new(container: K) -> Self {
162        Self(ManuallyDrop::new(container))
163    }
164
165    pub fn into_inner(mut self) -> K {
166        // Safety: This value is always initialised.
167        // Once we take it, we need to be careful to not call `drop` on `self`.
168        let value = unsafe { ManuallyDrop::take(&mut self.0) };
169        // This doesn't leak, because `self` is contained purely on the stack.
170        std::mem::forget(self);
171        value
172    }
173}
174
175impl<K, T> AsRef<T> for FlatDrop<K>
176where
177    T: ?Sized,
178    K: IntoOptionInner,
179    K::Inner: Recursive<Container = K>,
180    K: AsRef<T>,
181{
182    fn as_ref(&self) -> &T {
183        (**self).as_ref()
184    }
185}
186
187impl<K, T> AsMut<T> for FlatDrop<K>
188where
189    T: ?Sized,
190    K: IntoOptionInner,
191    K::Inner: Recursive<Container = K>,
192    K: AsMut<T>,
193{
194    fn as_mut(&mut self) -> &mut T {
195        (**self).as_mut()
196    }
197}
198
199impl<K> Deref for FlatDrop<K>
200where
201    K: IntoOptionInner,
202    K::Inner: Recursive<Container = K>,
203{
204    type Target = K;
205
206    fn deref(&self) -> &K {
207        self.0.deref()
208    }
209}
210
211impl<K> DerefMut for FlatDrop<K>
212where
213    K: IntoOptionInner,
214    K::Inner: Recursive<Container = K>,
215{
216    fn deref_mut(&mut self) -> &mut K {
217        self.0.deref_mut()
218    }
219}
220
221impl<K> From<K> for FlatDrop<K>
222where
223    K: IntoOptionInner,
224    K::Inner: Recursive<Container = K>,
225{
226    fn from(value: K) -> Self {
227        Self::new(value)
228    }
229}
230
231impl<T> FlatDrop<Box<T>>
232where
233    T: Recursive<Container = Box<T>>,
234{
235    pub fn new_boxed(value: T) -> Self {
236        Self::new(Box::new(value))
237    }
238}
239
240impl<T> FlatDrop<Rc<T>>
241where
242    T: Recursive<Container = Rc<T>>,
243{
244    pub fn new_rc(value: T) -> Self {
245        Self::new(Rc::new(value))
246    }
247}
248
249impl<T> FlatDrop<Arc<T>>
250where
251    T: Recursive<Container = Arc<T>>,
252{
253    pub fn new_arc(value: T) -> Self {
254        Self::new(Arc::new(value))
255    }
256}
257
258#[cfg(feature = "serde")]
259impl<K> serde::Serialize for FlatDrop<K>
260where
261    K: IntoOptionInner,
262    K::Inner: Recursive<Container = K>,
263    K: serde::Serialize,
264{
265    fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
266    where
267        S: serde::Serializer,
268    {
269        <K as serde::Serialize>::serialize(self, serializer)
270    }
271}
272
273#[cfg(feature = "serde")]
274impl<'de, K> serde::Deserialize<'de> for FlatDrop<K>
275where
276    K: IntoOptionInner,
277    K::Inner: Recursive<Container = K>,
278    K: serde::Deserialize<'de>,
279{
280    fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
281    where
282        D: serde::Deserializer<'de>,
283    {
284        <K as serde::Deserialize>::deserialize(deserializer).map(Self::new)
285    }
286}
287
288#[cfg(test)]
289mod tests {
290    use crate::{FlatDrop, Recursive};
291
292    /// Peano natural numbers.
293    enum Natural {
294        Zero,
295        Succ(FlatDrop<Box<Natural>>),
296    }
297
298    impl Recursive for Natural {
299        type Container = Box<Natural>;
300
301        fn destruct(self) -> impl Iterator<Item = Self::Container> {
302            match self {
303                Natural::Zero => None,
304                Natural::Succ(pred) => Some(pred.into_inner()),
305            }
306            .into_iter()
307        }
308    }
309
310    impl Natural {
311        pub fn from_usize(value: usize) -> Self {
312            (0..value).fold(Self::Zero, |nat, _| {
313                Self::Succ(FlatDrop::new(Box::new(nat)))
314            })
315        }
316    }
317
318    #[test]
319    fn test_large_natural() {
320        // Create a new thread with a 4kb stack and allocate a number far bigger than 4 * 1024.
321        const STACK_SIZE: usize = 4 * 1024;
322
323        fn task() {
324            let nat = Natural::from_usize(STACK_SIZE * 100);
325            println!("Dropping...");
326            drop(std::hint::black_box(nat));
327            println!("Dropped.");
328        }
329
330        std::thread::Builder::new()
331            .stack_size(STACK_SIZE)
332            .spawn(task)
333            .unwrap()
334            .join()
335            .unwrap();
336    }
337}
flat_drop/lib.rs

flat_drop/
lib.rs
