Implementation of a generic binary heap data structure in Rust (0.11). Just for funsies.

heap.rs

pub struct Heap<T> {
    data: Vec<T>
}

impl<T: PartialOrd> Heap<T> {
    /// Create a new empty heap.
    pub fn new() -> Heap<T> {
        Heap { data: Vec::<T>::new() }
    }

    /// Create a new empty heap with an initial capacity.
    pub fn with_capacity(capacity: uint) -> Heap<T> {
        Heap { data: Vec::<T>::with_capacity(capacity) }
    }

    /// Create a new heap containing the given values.
    pub fn from_vec(values: Vec<T>) -> Heap<T> {
        let mut h = Heap { data: Vec::<T>::with_capacity(values.len()) };
        h.insert_vec(values);
        return h
    }

    pub fn len(&self) -> uint {
        self.data.len()
    }

    /// Remove the top item from the heap and return it,
    /// or None if the heap is empty.
    pub fn pop(&mut self) -> Option<T> {
        if self.data.is_empty() { None }
        else {
            let i = self.data.len() - 1;
            self.data.as_mut_slice().swap(0, i);
            let result = self.data.remove(i);
            self.sift_down(0);
            return result
        }
    }

    /// Insert the given value into the heap. The heap will automatically
    /// grow capacity if necessary.
    pub fn insert(&mut self, value: T) {
        // Double the capacity of the heap if it is currently full.
        if self.data.len() == self.data.capacity() {
            let cap = self.data.capacity();
            self.data.reserve_additional(cap);
        }
        self.data.push(value);
        let i = self.data.len() - 1;
        self.sift_up(i);
    }

    pub fn insert_vec(&mut self, values: Vec<T>) {
        for value in values.move_iter() { self.insert(value); }
    }


    /// Shrink the capacity of the heap as much as possible,
    /// so that there is no unused capacity.
    pub fn shrink_to_fit(&mut self) {
        self.data.shrink_to_fit();
    }


    fn sift_up(&mut self, i: uint) {
        match parent_index(i) {
            Some(parent) => {
                if self.data.get(i).lt(self.data.get(parent)) {
                    self.data.as_mut_slice().swap(i, parent);
                    self.sift_up(parent);
                }
            }
            None => {}
        }
    }

    fn sift_down(&mut self, i: uint) {
        let len = self.data.len();
        let l = left_child_index(i);
        let r = right_child_index(i);

        let child =
            if l >= len && r >= len { return } // no children, nothing to do
            else if l >= len { r }
            else if r >= len { l }
            else if self.data.get(l).lt(self.data.get(r)) { l }
            else { r };

        if self.data.get(child).lt(self.data.get(i)) {
            self.data.as_mut_slice().swap(child, i);
            self.sift_down(child);
        }
    }
}


impl<T: Clone> Heap<T> {
    /// Return a clone of the top item in the heap, or None if the heap is
    /// empty. Does not modify the heap.
    pub fn peek(&self) -> Option<T> {
        if self.data.is_empty() { None }
        else { Some(self.data.get(0).clone()) }
    }

    /// Return a new Vec containing clones of all values in the heap. The
    /// values in the new Vec retain the internal heap order.
    pub fn as_vec(&self) -> Vec<T> {
        self.data.clone()
    }
}


/// Returns the index of the given index's parent,
/// or None if it has no parent (i.e. i is 0).
fn parent_index(i: uint) -> Option<uint> {
    if i == 0 { None }
    else { Some((i - 1) / 2u) }
}

/// Returns the index of the given index's left child. The returned index
/// may be larger than the current capacity of the heap.
fn left_child_index(i: uint) -> uint {
    (i + 1) * 2 - 1
}

/// Returns the index of the given index's right child. The returned index
/// may be larger than the current capacity of the heap.
fn right_child_index(i: uint) -> uint {
    (i + 1) * 2
}