smalloca 0.1.0

#![cfg_attr(not(std), no_std)]

extern "C" {
    /// This is the external smalloca function which essentially is just a _builtin_alloca.
    fn _smalloca(size: usize, dto: *mut u8, wrapper: extern "C" fn(*mut u8, *mut u8));
}

#[repr(C)]
/// We can not build a closure and pass it as a C function. Instead we use the wrapper functions,
/// which adhere to the C ABI and pass all data via a Data Transfer Object. This is essentially
/// the long way around a closure.
struct SmallocaDto<T, F, R>
where
    F: FnOnce(&mut [T]) -> R,
{
    size: usize,
    len: usize,
    ret: Option<R>,
    prototype: Option<T>,
    function: Option<F>,
}

/// Initializes the slice with clones of the prototype.
fn clone_init<T>(smalloca: &mut [T], prototype: T)
where
    T: Clone,
{
    // Initialize the data
    // The data that is inside smalloca is uninitialized at this point. If we override it with
    // initialised data, the destructor will be executed on the uninitialized data. We have to
    // prevent that by calling mem::forget;
    if smalloca.len() > 0 {
        for i in 0..smalloca.len() - 1 {
            let uninit = core::mem::replace::<T>(&mut smalloca[i], prototype.clone());
            core::mem::forget(uninit);
        }

        let uninit = core::mem::replace::<T>(&mut smalloca[smalloca.len() - 1], prototype);
        core::mem::forget(uninit);
    };
}

fn default_init<T>(smalloca: &mut [T])
where
    T: Default,
{
    for i in 0..smalloca.len() {
        let uninit = core::mem::replace::<T>(&mut smalloca[i], T::default());
        core::mem::forget(uninit);
    }
}

/// Zeroes out the slice. This has to happen in order to run the destructors.
fn zero_out<T>(smalloca: &mut [T]) {
    // Now we need to manually drop the the elements in smalloca
    // FIXME: Why doesn't this happen automatically, when smalloca gets out of scope?
    for i in 0..smalloca.len() {
        unsafe { core::mem::replace::<T>(&mut smalloca[i], core::mem::zeroed::<T>()) };
    }
}

/// Wraps the provided closure. It manages initialization and dropping of the
/// sliced values and handles returning the value. This version initializes the slices with clones
/// of the prototype and therefore requires [`Clone`].
extern "C" fn wrapper_clone<T, F, R>(space: *mut u8, dto: *mut u8)
where
    T: Clone,
    F: FnOnce(&mut [T]) -> R,
{
    let dto: &mut SmallocaDto<T, F, R> = unsafe { core::mem::transmute(dto) };
    let smalloca: &mut [T] = unsafe { core::slice::from_raw_parts_mut(space as *mut T, dto.len) };

    clone_init(smalloca, dto.prototype.take().unwrap());

    // Call the actual inner function
    dto.ret = Some((dto.function.take().unwrap())(smalloca));

    // smalloca is out of scope now, however the data is still there, we can get a new handle to it.
    let smalloca: &mut [T] = unsafe { core::slice::from_raw_parts_mut(space as *mut T, dto.len) };
    zero_out(smalloca);
}

/// Wraps the provided closure. It manages initialization and dropping of the
/// sliced values and handles returning the value. This version initializes the slices with
/// the default value of the Object and therefore requires [`Default`]
extern "C" fn wrapper_default<T, F, R>(space: *mut u8, dto: *mut u8)
where
    T: Default,
    F: FnOnce(&mut [T]) -> R,
{
    let dto: &mut SmallocaDto<T, F, R> = unsafe { core::mem::transmute(dto) };
    let smalloca: &mut [T] = unsafe { core::slice::from_raw_parts_mut(space as *mut T, dto.len) };

    default_init(smalloca);

    // Call the actual inner function
    dto.ret = Some((dto.function.take().unwrap())(smalloca));

    // smalloca is out of scope now, however the data is still there, we can get a new handle to it.
    let smalloca: &mut [T] = unsafe { core::slice::from_raw_parts_mut(space as *mut T, dto.len) };
    zero_out(smalloca);
}

/// Builds a slice of a Prototype Objects with specified size, fills it with clones of the prototype
/// and passes it to the provided function. It runs the destructors after executing the function.
/// The prototype needs to implement [`Clone`]
pub fn smalloca<T, F, R>(prototype: T, len: usize, function: F) -> R
where
    T: Clone,
    F: FnOnce(&mut [T]) -> R,
{
    let size = core::mem::size_of::<T>() * len;

    let mut dto = SmallocaDto {
        size,
        len,
        ret: None,
        prototype: Some(prototype),
        function: Some(function),
    };

    unsafe {
        _smalloca(
            size,
            &mut dto as *mut SmallocaDto<T, F, R> as *mut u8,
            wrapper_clone::<T, F, R>,
        );
    }

    dto.ret.unwrap()
}

/// Builds a slice of a Prototype Objects with specified size, fills it with default values
/// and passes it to the provided function. It runs the destructors after executing the function.
/// The prototype needs to implement [`Default`]
pub fn smalloca_default<T, F, R>(len: usize, function: F) -> R
where
    T: Default,
    F: FnOnce(&mut [T]) -> R,
{
    let size = core::mem::size_of::<T>() * len;

    let mut dto = SmallocaDto {
        size,
        len,
        ret: None,
        prototype: None,
        function: Some(function),
    };

    unsafe {
        _smalloca(
            size,
            &mut dto as *mut SmallocaDto<T, F, R> as *mut u8,
            wrapper_default::<T, F, R>,
        );
    }

    dto.ret.unwrap()
}

#[cfg(test)]
mod tests {
    use crate::{smalloca, smalloca_default};
    use core::sync::atomic::{AtomicIsize, Ordering};

    static GENERATOR: AtomicIsize = AtomicIsize::new(0);

    #[derive(Debug)]
    struct MemSafetyTester {
        genesis: isize,
    }

    impl MemSafetyTester {
        fn new() -> Self {
            let id = GENERATOR.fetch_add(1, Ordering::Relaxed);
            MemSafetyTester { genesis: id }
        }
    }

    impl Clone for MemSafetyTester {
        fn clone(&self) -> Self {
            let id = GENERATOR.fetch_add(1, Ordering::SeqCst);
            MemSafetyTester { genesis: id }
        }
    }

    impl Default for MemSafetyTester {
        fn default() -> Self {
            let id = GENERATOR.fetch_add(1, Ordering::SeqCst);
            MemSafetyTester { genesis: id }
        }
    }

    impl Drop for MemSafetyTester {
        fn drop(&mut self) {
            GENERATOR.fetch_sub(1, Ordering::SeqCst);
        }
    }

    #[test]
    /// Tests, that all clones objects are properly dropped
    fn memsafety_test_clone() {
        assert_eq!(GENERATOR.load(Ordering::SeqCst), 0);
        smalloca(MemSafetyTester::new(), 8, |_a| {
            assert_eq!(GENERATOR.load(Ordering::SeqCst), 8);
        });
        assert_eq!(GENERATOR.load(Ordering::SeqCst), 0);
    }

    #[test]
    /// Thests, that the default version of smalloc is memory safe as well
    fn memsafety_test_default() {
        assert_eq!(GENERATOR.load(Ordering::SeqCst), 0);
        smalloca_default(8, |_a: &mut [MemSafetyTester]| {
            assert_eq!(GENERATOR.load(Ordering::SeqCst), 8);
        });
        assert_eq!(GENERATOR.load(Ordering::SeqCst), 0);
    }

    #[test]
    /// Tests that the return value gets properly returned.
    fn return_test() {
        let x = smalloca(8, 8, |_a| 5);

        assert_eq!(x, 5);
    }

    #[test]
    /// Tests the use of dynamically generated smallocas.
    /// Also tests iterators.
    fn dyn_smalloca_test() {
        for i in 0..32 {
            let x = smalloca(1, i, |a| {
                let mut sum = 0;
                for x in a {
                    sum += *x;
                }
                sum
            });

            assert_eq!(x, i);
        }
    }

    fn inner_function(values: &mut [u32]) -> u32 {
        let mut sum = 0;
        for i in 0..values.len() {
            sum += values[i];
        }
        sum
    }

    #[test]
    /// Tests, that inside the smalloca, we can also properly use the values and after
    /// returning, out stack is still intact.
    fn inner_fucntion_call_test() {
        let result = smalloca(1, 8, |a| inner_function(a));

        assert_eq!(result, 8);
    }

    #[test]
    #[should_panic]
    /// Tests that a panic is handled properly.
    fn panic_test() {
        smalloca(1, 8, |_a| panic!("This is a test"));
    }

}