dynvec 0.1.4

use super::{Location, Region};

use std::alloc::Layout;

/// A contiguous region of the memory.
pub struct Block<H> {
    /// The locations of the objects allocated in this block.
    /// This vector is sorted by `ptr`.
    locations: Vec<Location<H>>,
    /// The actual owned data of this region.
    data: Box<[u8]>,
}

impl<H> Block<H> {
    /// Creates a new `Block` that is `length` bytes long.
    pub fn new(length: usize) -> Self {
        let data = unsafe {
            let layout = Layout::from_size_align_unchecked(length, 1);
            std::alloc::alloc(layout)
        };

        if data.is_null() {
            panic!("Failed to allocate the block");
        }

        // Safety: the data was just allocated
        let slice = unsafe { std::slice::from_raw_parts_mut(data, length) };

        Self {
            locations: Vec::new(),

            // TODO: replace with `Box::new_uninit_slice` when stable
            data: unsafe { Box::from_raw(slice) },
        }
    }

    /// Creates a new empty `Block` that is large enough to store
    /// each layout.
    pub fn for_layouts(iter: impl IntoIterator<Item = Layout>) -> Self {
        let mut max_align = 1;
        let mut size = 0;

        for layout in iter {
            if layout.align() > max_align {
                max_align = layout.align();
            }

            size += layout.size();
        }

        // the allocated block might not be properly aligned
        // to our max alignement
        // we need to make sure that we have enough bytes in the worst case
        let length = max_align + size - 1;
        Self::new(length)
    }

    /// Returns a pointer to the first byte that is not part of the block.
    fn _end_ptr(&self) -> *mut u8 {
        unsafe { self.data.as_ptr().add(self.data.len()) as *mut u8 }
    }
}

impl<H> Region<H> for Block<H> {
    type Pointers = std::vec::IntoIter<*mut u8>;

    fn allocate(&mut self, layout: Layout, header: H) -> Result<*mut u8, H> {
        let mut ptr = unsafe {
            let cursor = if let Some(loc) = self.locations.last() {
                loc.ptr.add(loc.layout.size())
            } else {
                self.data.as_ptr()
            };

            let offset = cursor as usize % layout.align();
            cursor.add(offset) as *mut u8
        };

        if self.locations.is_empty() {
            // no object is allocated
            // we just have to check whether the block is large enough
            if unsafe { ptr.add(layout.size()) } > self._end_ptr() {
                return Err(header);
            } else {
                // a location was found, we can add our location
                self.locations.push(Location {
                    ptr,
                    layout,
                    header,
                });
            }
        } else {
            let mut loc = 0usize;
            let mut retry = true;

            loop {
                unsafe {
                    let location = self.locations.get_unchecked(loc);
                    while ptr < location.ptr.add(location.layout.size()) {
                        ptr = ptr.add(layout.align());
                    }
                }

                // `ptr` is still properly aligned

                loc += 1;

                if loc >= self.locations.len() {
                    // no more allocated objects here
                    // we just need to check whether the block is large enough
                    if unsafe { ptr.add(layout.size()) } > self._end_ptr() {
                        // we might want to retry from index 0.
                        if retry {
                            let offset = self.data.as_ptr() as usize % layout.align();
                            ptr = unsafe { self.data.as_mut_ptr().add(offset) };
                            retry = false;
                            loc = 0usize;
                        } else {
                            return Err(header);
                        }
                    } else {
                        break;
                    }
                }

                unsafe {
                    let location = self.locations.get_unchecked(loc);
                    if ptr.add(layout.size()) <= location.ptr {
                        // we found a location between two existing blocks
                        break;
                    }
                }
            }

            // a pointer was found, we can add the location
            self.locations.insert(
                loc,
                Location {
                    ptr,
                    layout,
                    header,
                },
            );
        }

        Ok(ptr)
    }

    fn deallocate(&mut self, ptr: *mut u8) -> Option<H> {
        // we just need to find the location that have this specific pointer and
        // remove it from our vector.
        match self.locations.iter().position(|loc| loc.ptr == ptr) {
            Some(idx) => {
                let loc = self.locations.remove(idx);
                Some(loc.header)
            }
            None => None,
        }
    }

    fn has_allocated(&self, ptr: *mut u8) -> bool {
        self.locations.iter().any(|loc| loc.ptr == ptr)
    }

    fn count(&self) -> usize {
        self.locations.len()
    }

    fn allocations(&self) -> Self::Pointers {
        self.locations
            .iter()
            .map(|loc| loc.ptr)
            .collect::<Vec<*mut u8>>()
            .into_iter()
    }
}

#[cfg(test)]
#[test]
fn reuse_deallocated_memory() {
    let mut block = Block::new(1024);

    unsafe {
        let ptr = block.allocate(Layout::new::<u64>(), ()).unwrap() as *mut u64;

        ptr.write(6u64);

        block.deallocate(ptr as *mut u8);

        let another_ptr = block.allocate(Layout::new::<u64>(), ()).unwrap() as *mut u64;

        // the memory should be the same as before because the allocator
        // use the same location.
        assert_eq!(another_ptr.read(), 6u64);
    }
}

#[cfg(test)]
#[test]
fn valid_alignement() {
    for _ in 0..200 {
        let mut block = Block::new(32);
        let ptr = block.allocate(Layout::new::<u128>(), ()).unwrap();

        assert_eq!(ptr as usize % std::mem::align_of::<u128>(), 0);
    }
}

#[cfg(test)]
#[test]
fn distinct_allocations() {
    let mut block = Block::new(50);

    let allocations: Vec<*mut u8> = (0..50)
        .map(|_| block.allocate(Layout::new::<u8>(), ()))
        .map(|r| r.unwrap())
        .collect();

    for (i, &a) in allocations.iter().enumerate() {
        for &b in allocations.iter().skip(i + 1) {
            assert_ne!(a, b);
        }
    }
}

#[cfg(test)]
#[test]
fn invalid_deallocation() {
    let mut block = Block::new(3);

    let a = block.allocate(Layout::new::<u16>(), 5u8).unwrap();
    assert_eq!(block.deallocate(a), Some(5u8));
    assert_eq!(block.deallocate(a), None);
    assert_eq!(block.deallocate(12usize as _), None);
}

#[cfg(test)]
#[test]
fn has_allocated() {
    let mut block = Block::new(10);

    let a = block.allocate(Layout::new::<u16>(), ()).unwrap();
    let b = block.allocate(Layout::new::<u16>(), ()).unwrap();
    let c = block.allocate(Layout::new::<u8>(), ()).unwrap();

    assert!(block.has_allocated(a));
    assert!(block.has_allocated(b));
    assert!(block.has_allocated(c));

    block.deallocate(a).unwrap();
    block.deallocate(b).unwrap();

    assert!(!block.has_allocated(a));
    assert!(!block.has_allocated(b));
    assert!(block.has_allocated(c));

    block.deallocate(c).unwrap();

    assert!(!block.has_allocated(c));
}

#[cfg(test)]
#[test]
fn for_layouts() {
    let layouts = vec![
        Layout::new::<u8>(),
        Layout::new::<u8>(),
        Layout::new::<u8>(),
        Layout::new::<u16>(),
        Layout::new::<u16>(),
        Layout::new::<u32>(),
        Layout::new::<u32>(),
        Layout::new::<u32>(),
        Layout::new::<u8>(),
        Layout::new::<u128>(),
    ];

    let mut block = Block::for_layouts(layouts);

    block.allocate(Layout::new::<u8>(), ()).unwrap();
    block.allocate(Layout::new::<u128>(), ()).unwrap();
    block.allocate(Layout::new::<u8>(), ()).unwrap();
    block.allocate(Layout::new::<u16>(), ()).unwrap();
    block.allocate(Layout::new::<u32>(), ()).unwrap();
    block.allocate(Layout::new::<u32>(), ()).unwrap();
    block.allocate(Layout::new::<u16>(), ()).unwrap();
    block.allocate(Layout::new::<u32>(), ()).unwrap();
    block.allocate(Layout::new::<u8>(), ()).unwrap();
    block.allocate(Layout::new::<u8>(), ()).unwrap();
}

#[cfg(test)]
#[test]
fn allocations_count() {
    let mut block = Block::new(20);

    block.allocate(Layout::new::<u8>(), ()).unwrap();
    block.allocate(Layout::new::<u16>(), ()).unwrap();
    block.allocate(Layout::new::<u32>(), ()).unwrap();
    block.allocate(Layout::new::<u8>(), ()).unwrap();
    block.allocate(Layout::new::<u8>(), ()).unwrap();

    assert_eq!(block.count(), 5);
}