shortpointer/src/main.rs

use core::alloc::Layout;
use core::marker::PhantomData;
use core::mem::size_of;
use core::ops::{Deref, DerefMut};

use std::fmt::Formatter;
use std::mem::size_of_val;
use std::sync::Mutex;

use lazy_static::lazy_static;

lazy_static! {
    static ref POOL: Pool = Pool::default();
}

struct Pool {
    inner: Mutex<PoolInner>,
    base_ptr: usize,
}

impl Default for Pool {
    fn default() -> Self {
        let inner = Mutex::new(PoolInner { data: [0; 0x40000], begin: 0, end: 0, last_position: 0 });
        let base_ptr = inner.lock().unwrap().data.as_ptr() as usize;
        Pool {
            inner,
            base_ptr,
        }
    }
}

impl Pool {
    fn alloc(&self, layout: Layout) -> u16 {
        self.inner.lock().unwrap().alloc(layout)
    }
    fn dealloc(&self, sptr: u16) {
        self.inner.lock().unwrap().dealloc(sptr)
    }
}

impl std::fmt::Debug for Pool {
    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
        self.inner.lock().unwrap().fmt(f)
    }
}

#[repr(C, align(4))]
struct PoolInner {
    data: [u8; 0x40000],
    begin: usize,
    end: usize,
    last_position: usize,
}

// halfassed bump allocator just to prove a point.
impl PoolInner {
    fn alloc(&mut self, layout: Layout) -> u16 {
        let size = layout.size();
        let align = layout.align().max(4);

        let mut position = self.end + size_of::<u16>(); // space for skip marker
        let alignmask = align - 1;
        if position & alignmask != 0 {
            position += ((position ^ alignmask) & alignmask) + 1;
        }
        let new_end = position + size;
        assert!(new_end < self.data.len(), "ShortPointer memory pool exhausted");
        // using most significant bit as whether it's been freed
        assert!(position - self.last_position < self.data.len() / 2, "Alignment caused pool size overrun");

        let skip_marker = 0x8000_u16.to_ne_bytes();
        self.data[position - 2..position].copy_from_slice(&skip_marker);
        self.end = new_end;
        if self.begin == 0 {
            self.begin = position;
        }
        if self.last_position != 0 {
            unsafe {
                let last_skip_marker = self.data.as_mut_ptr().add(self.last_position - 2) as *mut u16;
                assert_eq!(*last_skip_marker & 0x7fff, 0);
                *last_skip_marker |= ((position - self.last_position) / 4) as u16;
            }
        }
        self.last_position = position;
        (position / 4) as u16
    }

    fn dealloc(&mut self, sptr: u16) {
        let mut position = (sptr * 4) as usize;
        unsafe {
            let skip_marker_ptr = self.data.as_mut_ptr().add(position - 2) as *mut u16;
            *skip_marker_ptr &= 0x7fff;
            // see how much we can free
            let mut skip_marker = *skip_marker_ptr;
            let update_begin = self.begin == position;
            while skip_marker & 0x8000 == 0 {
                if skip_marker == 0 {
                    // last object, we can reclaim memory back to sptr
                    self.end = (sptr * 4) as usize;
                    break;
                } else {
                    position += skip_marker as usize * 4;
                }
                skip_marker = *(self.data.as_ptr().add(position - 2) as *const u16);
            }
            if update_begin {
                self.begin = position;
            }
        }
        if self.begin == self.end {
            // all objects freed, we can reset the entire state
            self.begin = 0;
            self.end = 0;
            self.last_position = 0;
        } else {
            // TODO: restore self.last_position linearly from self.begin
        }
    }
}

impl std::fmt::Debug for PoolInner {
    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
        write!(
            f, "Pool {{ data: [_; {}], begin: {}, end: {}, last_position: {} }}",
            self.data.len(),
            self.begin,
            self.end,
            self.last_position,
        )
    }
}

pub struct ShortPointer<T>(u16, PhantomData<T>);

impl<T> Drop for ShortPointer<T> {
    fn drop(&mut self) {
        POOL.dealloc(self.0)
    }
}

impl<T> ShortPointer<T> {
    fn deref_inner(&self) -> &mut T {
        unsafe {
            ((POOL.base_ptr + (self.0 as usize * 4)) as *mut T).as_mut().unwrap()
        }
    }

    pub fn new(obj: T) -> Self {
        let layout = Layout::for_value(&obj);
        let mut sptr = Self(POOL.alloc(layout), PhantomData::default());
        let uninitialized = core::mem::replace(sptr.deref_mut(), obj);
        core::mem::forget(uninitialized);
        sptr
    }
}

impl<T> Deref for ShortPointer<T> {
    type Target = T;

    fn deref(&self) -> &Self::Target {
        self.deref_inner()
    }
}

impl<T> DerefMut for ShortPointer<T> {
    fn deref_mut(&mut self) -> &mut Self::Target {
        self.deref_inner()
    }
}

#[allow(unused_must_use)] // dbg! statements
fn main() {
    dbg!(&*POOL);
    {
        let foo = ShortPointer::new([1, 2, 3, 4, 5]);
        dbg!(&*POOL);
        {
            let bar = ShortPointer::new(0x900dCafeDa7e5_u64);
            dbg!(&*POOL);
            println!("{} | {:?} | {:x}", size_of_val(&foo), *foo, *bar);
        }
        dbg!(&*POOL);
    }
    dbg!(&*POOL);
}