Is there a way to tell the Rust compiler to call drop on partially-initialized array elements when handling a panic?

Question

I\'m working on a custom type where I have the following requirements:

1. A collection of elements that avoid heap allocation. I am using arrays instead of a Ve

Answer 1

Is there a way to tell the Rust compiler to call drop on partially-initialized array elements when handling a panic?

No, but you can call drop yourself. You need to run code when a panic occurs.

Brute-force solution

This uses the building blocks of catch_unwind, resume_unwind, and AssertUnwindSafe to notice that a panic occurred and run some cleanup code:

fn default() -> Self {
    use std::panic::{self, AssertUnwindSafe};

    let mut temp = NoDrop::new(Self {
        data: unsafe { std::mem::uninitialized() },
    });

    let mut valid = 0;

    let panicked = {
        let mut temp = AssertUnwindSafe(&mut temp);
        let mut valid = AssertUnwindSafe(&mut valid);

        std::panic::catch_unwind(move || unsafe {
            for index in 0..CAPACITY {
                std::ptr::write(&mut temp.data[index], T::default());
                **valid += 1;
            }
        })
    };

    if let Err(e) = panicked {
        for i in 0..valid {
            unsafe { std::ptr::read(&temp.data[i]) };
        }

        panic::resume_unwind(e);
    }

    temp.into_inner()
}

Slightly nicer

Once you recognize that a type's Drop implementation is run when a panic occurs, you can use that to your advantage by creating a drop bomb — a type that cleans up when dropped but in the success path it is not dropped:

extern crate nodrop;

use nodrop::NoDrop;

use std::{mem, ptr};

const CAPACITY: usize = 5;
type Data<T> = [T; CAPACITY];

struct Temp<T> {
    data: NoDrop<Data<T>>,
    valid: usize,
}

impl<T> Temp<T> {
    unsafe fn new() -> Self {
        Self {
            data: NoDrop::new(mem::uninitialized()),
            valid: 0,
        }
    }

    unsafe fn push(&mut self, v: T) {
        if self.valid < CAPACITY {
            ptr::write(&mut self.data[self.valid], v);
            self.valid += 1;
        }
    }

    unsafe fn into_inner(mut self) -> Data<T> {
        let data = mem::replace(&mut self.data, mem::uninitialized());
        mem::forget(self);
        data.into_inner()
    }
}

impl<T> Drop for Temp<T> {
    fn drop(&mut self) {
        unsafe {
            for i in 0..self.valid {
                ptr::read(&self.data[i]);
            }
        }
    }
}

struct Composite<T>(Data<T>);

impl<T> Default for Composite<T>
where
    T: Default,
{
    fn default() -> Self {
        unsafe {
            let mut tmp = Temp::new();

            for _ in 0..CAPACITY {
                tmp.push(T::default());
            }

            Composite(tmp.into_inner())
        }
    }
}

impl<T> From<Data<T>> for Composite<T> {
    fn from(value: Data<T>) -> Self {
        Composite(value)
    }
}

struct Dummy;

impl Drop for Dummy {
    fn drop(&mut self) {
        println!("dropping");
    }
}

impl Default for Dummy {
    fn default() -> Self {
        use std::sync::atomic::{AtomicUsize, Ordering, ATOMIC_USIZE_INIT};

        static COUNT: AtomicUsize = ATOMIC_USIZE_INIT;

        let count = COUNT.fetch_add(1, Ordering::SeqCst);
        if count < 3 {
            println!("default");
            Dummy {}
        } else {
            panic!("oh noes!");
        }
    }
}

pub fn main() {
    let _v1: Composite<Dummy> = Composite::default();
}

Note that I've made some unrelated cleanups:

1. Using an atomic variable instead of unsafe static mutable variables.
2. Don't use return as the last statement of a block.
3. Converted Composite into a newtype, as data isn't a wonderful variable name.
4. Imported the mem and ptr modules for easier access.
5. Created the Data<T> type alias to avoid retyping that detail.

An elegant solution

The choice of push in the second solution is no accident. Temp is a poor implementation of a variable-sized stack-allocated vector. There's a good implementation called arrayvec which we can use instead:

extern crate arrayvec;

use arrayvec::ArrayVec;

const CAPACITY: usize = 5;
type Data<T> = [T; CAPACITY];

struct Composite<T>(Data<T>);

impl<T> Default for Composite<T>
where
    T: Default,
{
    fn default() -> Self {
        let tmp: ArrayVec<_> = (0..CAPACITY).map(|_| T::default()).collect();

        match tmp.into_inner() {
            Ok(data) => Composite(data),
            Err(_) => panic!("Didn't insert enough values"),
        }
    }
}

Would you be surprised to learn that nodrop was created in a large part to be used for arrayvec? The same author created both!