memory.rs

use std::collections::BTreeMap;
use std::mem;
use std::ptr;

const LIVE: usize = 1 << 16;
const BLOCK: usize = 1 << 9;

/// An unmanaged pointer to a managed object.
pub type Pointer = *mut usize;

unsafe fn is_live(ptr: Pointer) -> bool { *ptr & LIVE == 0 }

unsafe fn size(ptr: Pointer) -> usize { *ptr & 0x3f }

/// Prevent this object and any other reachable from it, from being reclaimed
/// during the next garbage collection cycle.
pub unsafe fn mark(ptr: Pointer) {
    let mut vec = vec!(ptr);
    
    // While the stack of yet-to-mark objects is not empty...
    loop {
        let ptr = match vec.pop() {
            None => break,
            Some(ptr) => ptr
        };
        
        // Mark the object as non-garbage.
        if !is_live(ptr) {
            *ptr |= LIVE;
            
            // Schedule neighboring objects for marking.
            for idx in 1 .. size(ptr) {
                let mask = LIVE << idx;
                if *ptr & mask == mask {
                    let ptr = ptr.offset(idx as isize);
                    vec.push(*ptr as *mut usize);
                }
            }
        }
    }
}

type Adjust = BTreeMap<usize, Pointer>;

unsafe fn adjust(ptr: Pointer, adj: &Adjust) {
    
    // Clear the live flag.
    assert!(is_live(ptr), "Corrupted heap");
    *ptr &= !LIVE;
    
    // Adjust pointer fields.
    for idx in 1 .. size(ptr) {
        let mask = LIVE << idx;
        if *ptr & mask == mask {
            let ptr = ptr.offset(idx as isize);
            let old = *ptr;
            for &new in adj.get(&old) {
                *ptr = new as usize;
            }
        }
    }
}

struct BlockIter {
    curr: usize,
    data: Vec<usize>
}

impl Iterator for BlockIter {
    type Item = Pointer;
    
    fn next(&mut self) -> Option<Pointer> {
        while self.curr != self.data.len() {
            unsafe {
                let ptr = self.data.get_unchecked_mut(self.curr) as Pointer;
                self.curr += size(ptr);
                if is_live(ptr) { return Some(ptr) }
            }
        }
        None
    }
}

struct Block { data: Vec<usize> }

impl Block {
    fn new() -> Block {
        Block { data: Vec::with_capacity(BLOCK) }
    }
    
    fn allocate(&mut self, size: usize) -> Option<Pointer> {
        let old = self.data.len();
        let new = old + size;
        if new > self.data.capacity() { return None }
        unsafe {
            self.data.set_len(new);
            Some(self.data.get_unchecked_mut(old) as Pointer)
        }
    }
}

impl IntoIterator for Block {
    type Item = Pointer;
    type IntoIter = BlockIter;
    
    fn into_iter(self) -> BlockIter {
        BlockIter {
            curr: 0,
            data: self.data
        }
    }
}

/// A managed collection of objects that can be compacted by discarding objects
/// deemed unreachable.
pub struct Heap {
    prev: Vec<Block>,
    curr: Block
}

impl Heap {
    
    /// Create an empty heap.
    pub fn new() -> Heap {
        Heap {
            prev: Vec::new(),
            curr: Block::new()
        }
    }
    
    /// Allocate a new object.
    pub fn allocate(&mut self, size: usize) -> Pointer {
        loop {
            for ptr in self.curr.allocate(size) {
                return ptr
            }
            
            let mut new = Block::new();
            mem::swap(&mut self.curr, &mut new);
            self.prev.push(new);
        }
    }
    
    unsafe fn merge_object(&mut self, src: Pointer, adj: &mut Adjust) {
        let size = size(src);
        let dest = self.allocate(size);
        adj.insert(src as usize, dest);
        ptr::copy_nonoverlapping(src, dest, size);
    }
    
    unsafe fn merge_block(&mut self, block: Block, adj: &mut Adjust) {
        for src in block {
            self.merge_object(src, adj);
        }
    }
    
    /// 
    pub unsafe fn merge(&mut self, heap: Heap) {
        let mut adj = BTreeMap::new();
        
        // 
        for block in heap.prev {
            self.merge_block(block, &mut adj)
        }
        self.merge_block(heap.curr, &mut adj);
        
        // Adjust every reallocated object.
        for &ptr in adj.values() {
            adjust(ptr, &adj)
        }
    }
    
    /// 
    pub unsafe fn compact(&mut self) {
        let mut new = Heap::new();
        mem::swap(self, &mut new);
        self.merge(new);
    }
}