time-event 0.0.1

pub struct Heap<T> {
    pub(crate) items: Vec<Node<T>>,
    index: Vec<Slot>,
    next_slot: usize,
}

#[derive(Debug)]
pub(crate) struct Node<T> {
    data: T,
    slot_index: usize,
}

enum Slot {
    // index of item in the items array
    Full(usize, usize),
    // index of the next empty slot
    Empty(usize, usize),
}

pub struct SlotHandle {
    slot: usize,
    version: usize,
}

impl<T: Ord> Heap<T> {
    pub fn new() -> Self {
        Heap {
            items: vec![],
            index: vec![],
            next_slot: 0,
        }
    }

    pub fn push(&mut self, data: T) -> SlotHandle {
        //let slot = Slot::Full(self.items.len());
        let (slot_index, version) = if self.next_slot == self.index.len() {
            self.index.push(Slot::Full(self.items.len(), 1));
            self.next_slot += 1;
            (self.index.len() - 1, 1)
        } else {
            if let Slot::Empty(_, version) = self.index[self.next_slot] {
                let slot = Slot::Full(self.items.len(), version + 1);
                let slot_idx = match std::mem::replace(&mut self.index[self.next_slot], slot) {
                    Slot::Full(_, _) => panic!("inconsistent"),
                    Slot::Empty(next, _) => std::mem::replace(&mut self.next_slot, next),
                };
                (slot_idx, version + 1)
            } else {
                unreachable!()
            }
        };
        self.items.push(Node {
            data: data,
            slot_index: slot_index,
        });
        self.percolate_up(self.items.len() - 1);
        SlotHandle {
            slot: slot_index,
            version,
        }
    }

    pub fn peek(&self) -> Option<&T> {
        self.items.get(0).map(|node| &node.data)
    }

    pub fn peek_mut(&mut self) -> Option<&mut T> {
        self.items.get_mut(0).map(|node| &mut node.data)
    }

    pub fn pop(&mut self) -> Option<T> {
        if self.items.is_empty() {
            None
        } else {
            self.remove(self.items[0].slot_index)
        }
    }

    pub fn update<U, F>(&mut self, handle: &SlotHandle, f: F) -> Option<U>
    where
        F: Fn(&mut T) -> U,
    {
        if let Slot::Full(idx, _) = self.index[handle.slot] {
            let u = f(&mut self.items[idx].data);
            if idx > 0 {
                if &self.items[(idx - 1) / 2].data > &self.items[idx].data {
                    self.percolate_down(idx);
                } else {
                    self.percolate_up(idx);
                }
            } else {
                self.percolate_down(idx);
            }
            Some(u)
        } else {
            None
        }
    }

    pub fn remove_with_handle(&mut self, handle: &SlotHandle) -> Option<T> {
        match self.index[handle.slot] {
            Slot::Full(_, version) if version != handle.version => return None,
            Slot::Empty(_, _) => return None,
            _ => return self.remove(handle.slot),
        };
    }

    pub fn remove(&mut self, idx: usize) -> Option<T> {
        let (item_idx, version) = match self.index[idx] {
            Slot::Full(item_idx, version) => (item_idx, version),
            Slot::Empty(_, _) => return None,
        };

        let removed = self.items.swap_remove(item_idx);
        self.index[removed.slot_index] = Slot::Empty(
            std::mem::replace(&mut self.next_slot, removed.slot_index),
            version + 1,
        );

        if self.items.is_empty() {
            return Some(removed.data);
        }

        let origin_last = &self.items[item_idx];
        match self.index[origin_last.slot_index] {
            Slot::Full(_, version) => {
                self.index[origin_last.slot_index] = Slot::Full(item_idx, version)
            }
            Slot::Empty(_, _) => unreachable!(),
        }

        if origin_last.data > removed.data {
            self.percolate_down(item_idx);
        } else {
            self.percolate_up(item_idx);
        }

        Some(removed.data)
    }

    pub(crate) fn percolate_down(&mut self, idx: usize) {
        let mut idx = idx;
        let mut left = idx * 2 + 1;
        let mut right = idx * 2 + 2;

        while left < self.items.len() {
            let min_idx = match self.items.get(right) {
                Some(_) => {
                    if self.items[left].data < self.items[right].data {
                        left
                    } else {
                        right
                    }
                }
                None => left,
            };
            if self.items[idx].data > self.items[min_idx].data {
                self.items.swap(idx, min_idx);
                let version = match self.index[self.items[idx].slot_index] {
                    Slot::Full(_, version) => version,
                    Slot::Empty(_, version) => version,
                };
                self.index[self.items[idx].slot_index] = Slot::Full(idx, version);

                let version = match self.index[self.items[min_idx].slot_index] {
                    Slot::Full(_, version) => version,
                    Slot::Empty(_, version) => version,
                };
                self.index[self.items[min_idx].slot_index] = Slot::Full(min_idx, version);
                idx = min_idx;
                left = min_idx * 2 + 1;
                right = min_idx * 2 + 2;
            } else {
                break;
            }
        }
    }

    fn percolate_up(&mut self, idx: usize) {
        let mut cur_idx = idx;
        while cur_idx > 0 {
            let parent = (cur_idx - 1) / 2;
            if self.items[parent].data <= self.items[cur_idx].data {
                break;
            }
            self.items.swap(parent, cur_idx);
            let version = match self.index[self.items[cur_idx].slot_index] {
                Slot::Full(_, version) => version,
                Slot::Empty(_, version) => version,
            };
            self.index[self.items[cur_idx].slot_index] = Slot::Full(cur_idx, version);

            let version = match self.index[self.items[parent].slot_index] {
                Slot::Full(_, version) => version,
                Slot::Empty(_, version) => version,
            };
            self.index[self.items[parent].slot_index] = Slot::Full(parent, version);
            cur_idx = parent;
        }
    }
}

#[cfg(test)]
mod tests {
    use super::Heap;

    #[test]
    fn simple() {
        let mut h = Heap::new();
        h.push(1);
        h.push(2);
        h.push(8);
        h.push(4);
        assert_eq!(h.pop(), Some(1));
        assert_eq!(h.pop(), Some(2));
        assert_eq!(h.pop(), Some(4));
        assert_eq!(h.pop(), Some(8));
        assert_eq!(h.pop(), None);
        assert_eq!(h.pop(), None);
    }

    #[test]
    fn simple2() {
        let mut h = Heap::new();
        h.push(5);
        h.push(4);
        h.push(3);
        h.push(2);
        h.push(1);
        assert_eq!(h.pop(), Some(1));
        h.push(8);
        assert_eq!(h.pop(), Some(2));
        h.push(1);
        assert_eq!(h.pop(), Some(1));
        assert_eq!(h.pop(), Some(3));
        assert_eq!(h.pop(), Some(4));
        h.push(5);
        assert_eq!(h.pop(), Some(5));
        assert_eq!(h.pop(), Some(5));
        assert_eq!(h.pop(), Some(8));
    }

    #[test]
    fn remove() {
        let mut h = Heap::new();
        h.push(5);
        h.push(4);
        h.push(3);
        let two = h.push(2);
        h.push(1);
        assert_eq!(h.pop(), Some(1));
        assert_eq!(h.remove_with_handle(&two), Some(2));
        h.push(1);
        assert_eq!(h.pop(), Some(1));
        assert_eq!(h.pop(), Some(3));
    }

    fn vec2heap<T: Ord>(v: Vec<T>) -> Heap<T> {
        let mut h = Heap::new();
        for t in v {
            h.push(t);
        }
        return h;
    }

    #[test]
    fn test_peek_and_pop() {
        let data = vec![2, 4, 6, 2, 1, 8, 10, 3, 5, 7, 0, 9, 1];
        let mut sorted = data.clone();
        sorted.sort();
        let mut heap = vec2heap(data);
        while heap.peek().is_some() {
            assert_eq!(heap.peek().unwrap(), sorted.first().unwrap());
            assert_eq!(heap.pop().unwrap(), sorted.remove(0));
        }
    }

    #[test]
    fn test_push() {
        let mut heap = Heap::new();
        heap.push(-2);
        heap.push(-4);
        heap.push(-9);
        assert!(*heap.peek().unwrap() == -9);
        heap.push(-11);
        assert!(*heap.peek().unwrap() == -11);
        heap.push(-5);
        assert!(*heap.peek().unwrap() == -11);
        heap.push(-27);
        assert!(*heap.peek().unwrap() == -27);
        heap.push(-3);
        assert!(*heap.peek().unwrap() == -27);
        heap.push(-103);
        assert!(*heap.peek().unwrap() == -103);
    }

    fn check_to_vec(mut data: Vec<i32>) {
        let mut heap = Heap::new();
        for data in data.iter() {
            heap.push(*data);
        }
        data.sort();
        let mut v = Vec::new();
        while let Some(i) = heap.pop() {
            v.push(i);
        }
        assert_eq!(v, data);
    }

    #[test]
    fn test_to_vec() {
        check_to_vec(vec![]);
        check_to_vec(vec![5]);
        check_to_vec(vec![3, 2]);
        check_to_vec(vec![2, 3]);
        check_to_vec(vec![5, 1, 2]);
        check_to_vec(vec![1, 100, 2, 3]);
        check_to_vec(vec![1, 3, 5, 7, 9, 2, 4, 6, 8, 0]);
        check_to_vec(vec![2, 4, 6, 2, 1, 8, 10, 3, 5, 7, 0, 9, 1]);
        check_to_vec(vec![9, 11, 9, 9, 9, 9, 11, 2, 3, 4, 11, 9, 0, 0, 0, 0]);
        check_to_vec(vec![0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10]);
        check_to_vec(vec![10, 9, 8, 7, 6, 5, 4, 3, 2, 1, 0]);
        check_to_vec(vec![0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 0, 0, 0, 1, 2]);
        check_to_vec(vec![5, 4, 3, 2, 1, 5, 4, 3, 2, 1, 5, 4, 3, 2, 1]);
    }

    #[test]
    fn test_empty_pop() {
        let mut heap = Heap::<i32>::new();
        assert!(heap.pop().is_none());
    }

    #[test]
    fn test_empty_peek() {
        let empty = Heap::<i32>::new();
        assert!(empty.peek().is_none());
    }
}