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; CAPACITY];

struct Temp {
    data: NoDrop>,
    valid: usize,
}

impl Temp {
    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 {
        let data = mem::replace(&mut self.data, mem::uninitialized());
        mem::forget(self);
        data.into_inner()
    }
}

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

struct Composite(Data);

impl Default for Composite
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 From> for Composite {
    fn from(value: Data) -> 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 = 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 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; CAPACITY];

struct Composite(Data);

impl Default for Composite
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!