feat: stable min heap data structure (#91)

This change introduces the `MinHeap` stable structure analogous to
`std::collections::BinaryHeap`.
Contrary to `BinaryHeap`, `MinHeap` prioritizes smaller items, which I
found much more useful in practice than prioritizing larger items.

.github/workflows/ci.yml

@@ -62,7 +62,7 @@ jobs:
       - name: Install DFX
         run: |
           wget --output-document install-dfx.sh "https://internetcomputer.org/install.sh"
-          DFX_VERSION=${DFX_VERSION:=0.11.2} bash install-dfx.sh < <(yes Y)
+          DFX_VERSION=${DFX_VERSION:=0.14.1} bash install-dfx.sh < <(yes Y)
           rm install-dfx.sh
           dfx cache install
           echo "$HOME/bin" >> $GITHUB_PATH

CHANGELOG.md

@@ -0,0 +1,11 @@
+# Changelog
+
+All notable changes to this project will be documented in this file.
+
+The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/),
+and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.html).
+
+## [0.5.5] - Unreleased
+
+### Added
+- The `MinHeap` stable data structure (#91)

README.md

@@ -26,6 +26,7 @@ For more information about the philosophy behind the library, see [Roman's tutor
 - [Vec]: A growable array
 - [Log]: An append-only list of variable-size entries
 - [Cell]: A serializable value
+- [MinHeap]: A priority queue.
 
 ## How it Works
 
@@ -94,7 +95,7 @@ Dependencies:
 [dependencies]
 ic-cdk = "0.6.8"
 ic-cdk-macros = "0.6.8"
-ic-stable-structures = "0.5.4"
+ic-stable-structures = "0.5.5"
 ```
 
 Code:

examples/Cargo.toml

@@ -4,4 +4,5 @@ members = [
     "src/vecs_and_strings",
     "src/custom_types_example",
     "src/quick_start",
+    "src/task_timer",
 ]

examples/dfx.json

@@ -1,5 +1,5 @@
 {
-  "dfx": "0.11.2",
+  "dfx": "0.14.1",
   "canisters": {
     "quick_start": {
       "candid": "src/quick_start/candid.did",
@@ -20,6 +20,11 @@
       "candid": "src/custom_types_example/candid.did",
       "package": "custom_types_example",
       "type": "rust"
+    },
+    "task_timer": {
+      "candid": "src/task_timer/candid.did",
+      "package": "task_timer",
+      "type": "rust"
     }
   },
   "defaults": {

examples/src/task_timer/Cargo.toml

@@ -0,0 +1,14 @@
+[package]
+name = "task_timer"
+version = "0.2.0"
+edition = "2018"
+
+[lib]
+crate-type = ["cdylib"]
+
+[dependencies]
+candid = "0.8.4"
+ic0 = "0.18"
+ic-cdk = "0.6.10"
+ic-cdk-macros = "0.6.10"
+ic-stable-structures = { path = "../../../" }

examples/src/task_timer/candid.did

@@ -0,0 +1,3 @@
+service : () -> {
+    schedule_task : (after_sec : nat64) -> ();
+}

examples/src/task_timer/src/lib.rs

@@ -0,0 +1,62 @@
+//! An example showcasing how to use a MinHeap for scheduled tasks.
+use ic_cdk_macros::{post_upgrade, update};
+use ic_stable_structures::memory_manager::{MemoryId, MemoryManager, VirtualMemory};
+use ic_stable_structures::{DefaultMemoryImpl, StableMinHeap};
+use std::cell::RefCell;
+
+type Memory = VirtualMemory<DefaultMemoryImpl>;
+
+thread_local! {
+    static MEMORY_MANAGER: RefCell<MemoryManager<DefaultMemoryImpl>> =
+        RefCell::new(MemoryManager::init(DefaultMemoryImpl::default()));
+
+    static TASKS: RefCell<StableMinHeap<u64, Memory>> =
+        MEMORY_MANAGER.with(|mm|
+            RefCell::new(
+                StableMinHeap::init(mm.borrow().get(MemoryId::new(1)))
+                .expect("failed to initialize the tasks"))
+        );
+}
+
+#[post_upgrade]
+fn post_upgrade() {
+    reschedule();
+}
+
+#[update]
+fn schedule_task(after_sec: u64) {
+    let task_time = ic_cdk::api::time() + after_sec * 1_000_000_000;
+    TASKS.with(|t| {
+        t.borrow_mut()
+            .push(&task_time)
+            .expect("failed to schedule a task")
+    });
+    reschedule();
+}
+
+#[export_name = "canister_global_timer"]
+fn timer() {
+    let now = ic_cdk::api::time();
+    while let Some(task_time) = TASKS.with(|t| t.borrow().peek()) {
+        if task_time > now {
+            reschedule();
+            return;
+        }
+        let _ = TASKS.with(|t| t.borrow_mut().pop());
+
+        execute_task(task_time, now);
+        reschedule();
+    }
+}
+
+fn execute_task(scheduled_at: u64, now: u64) {
+    ic_cdk::println!("executing task scheculed at {scheduled_at}, current time is {now}");
+}
+
+fn reschedule() {
+    if let Some(task_time) = TASKS.with(|t| t.borrow().peek()) {
+        unsafe {
+            ic0::global_timer_set(task_time as i64);
+        }
+    }
+}

src/lib.rs

@@ -9,12 +9,14 @@ mod ic0_memory; // Memory API for canisters.
 pub mod log;
 pub use log::{Log as StableLog, Log};
 pub mod memory_manager;
+pub mod min_heap;
 pub mod reader;
 pub mod storable;
 #[cfg(test)]
 mod tests;
 mod types;
 pub mod vec;
+pub use min_heap::{MinHeap, MinHeap as StableMinHeap};
 pub use vec::{Vec as StableVec, Vec};
 pub mod vec_mem;
 pub mod writer;

src/min_heap.rs

@@ -0,0 +1,209 @@
+use crate::base_vec::{BaseVec, InitError};
+use crate::storable::BoundedStorable;
+use crate::{GrowFailed, Memory};
+use std::fmt;
+
+#[cfg(test)]
+mod tests;
+
+const MAGIC: [u8; 3] = *b"SMH"; // Short for "stable min heap".
+
+/// An implementation of the [binary min heap](https://en.wikipedia.org/wiki/Binary_heap).
+// NB. Contrary to [std::collections::BinaryHeap], this heap is a min-heap (smallest items come first).
+// Motivation: max heaps are helpful for sorting, but most daily programming tasks require min
+// heaps.
+pub struct MinHeap<T: BoundedStorable + PartialOrd, M: Memory>(BaseVec<T, M>);
+
+// Note: Heap Invariant
+// ~~~~~~~~~~~~~~~~~~~~
+//
+// HeapInvariant(heap, i, j) :=
+//   ∀ k: i ≤ k ≤ j: LET p = (k - 1)/2 IN (p ≤ i) => heap[p] ≤ heap[k]
+
+impl<T, M> MinHeap<T, M>
+where
+    T: BoundedStorable + PartialOrd,
+    M: Memory,
+{
+    /// Creates a new empty heap in the specified memory,
+    /// overwriting any data structures the memory might have
+    /// contained.
+    ///
+    /// Complexity: O(1)
+    pub fn new(memory: M) -> Result<Self, GrowFailed> {
+        BaseVec::<T, M>::new(memory, MAGIC).map(Self)
+    }
+
+    /// Initializes a heap in the specified memory.
+    ///
+    /// Complexity: O(1)
+    ///
+    /// PRECONDITION: the memory is either empty or contains a valid
+    /// stable heap.
+    pub fn init(memory: M) -> Result<Self, InitError> {
+        BaseVec::<T, M>::init(memory, MAGIC).map(Self)
+    }
+
+    /// Returns the number of items in the heap.
+    ///
+    /// Complexity: O(1)
+    pub fn len(&self) -> u64 {
+        self.0.len()
+    }
+
+    /// Returns true if the heap is empty.
+    ///
+    /// Complexity: O(1)
+    pub fn is_empty(&self) -> bool {
+        self.0.is_empty()
+    }
+
+    /// Pushes an item onto the heap.
+    ///
+    /// Complexity: O(log(self.len()))
+    pub fn push(&mut self, item: &T) -> Result<(), GrowFailed> {
+        self.0.push(item)?;
+        self.bubble_up(self.0.len() - 1, item);
+        debug_assert_eq!(Ok(()), self.check_invariant());
+        Ok(())
+    }
+
+    /// Removes the smallest item from the heap and returns it.
+    /// Returns `None` if the heap is empty.
+    ///
+    /// Complexity: O(log(self.len()))
+    pub fn pop(&mut self) -> Option<T> {
+        let n = self.len();
+        match n {
+            0 => None,
+            1 => self.0.pop(),
+            _more => {
+                let smallest = self.0.get(0).unwrap();
+                let last = self.0.pop().unwrap();
+                self.0.set(0, &last);
+                self.bubble_down(0, n - 1, &last);
+                debug_assert_eq!(Ok(()), self.check_invariant());
+                Some(smallest)
+            }
+        }
+    }
+
+    /// Returns the smallest item in the heap.
+    /// Returns `None` if the heap is empty.
+    ///
+    /// Complexity: O(1)
+    pub fn peek(&self) -> Option<T> {
+        self.0.get(0)
+    }
+
+    /// Returns an iterator visiting all values in the underlying vector, in arbitrary order.
+    pub fn iter(&self) -> impl Iterator<Item = T> + '_ {
+        self.0.iter()
+    }
+
+    /// Returns the underlying memory instance.
+    pub fn into_memory(self) -> M {
+        self.0.into_memory()
+    }
+
+    #[allow(dead_code)]
+    /// Checks the HeapInvariant(self, 0, self.len() - 1)
+    fn check_invariant(&self) -> Result<(), String> {
+        let n = self.len();
+        for i in 1..n {
+            let p = (i - 1) / 2;
+            let item = self.0.get(i).unwrap();
+            let parent = self.0.get(p).unwrap();
+            if is_less(&item, &parent) {
+                return Err(format!(
+                    "Binary heap invariant violated in indices {i} and {p}"
+                ));
+            }
+        }
+        Ok(())
+    }
+
+    /// PRECONDITION: self.0.get(i) == item
+    fn bubble_up(&mut self, mut i: u64, item: &T) {
+        // We set the flag if self.0.get(i) does not contain the item anymore.
+        let mut swapped = false;
+        // LOOP INVARIANT: HeapInvariant(self, i, self.len() - 1)
+        while i > 0 {
+            let p = (i - 1) / 2;
+            let parent = self.0.get(p).unwrap();
+            if is_less(item, &parent) {
+                self.0.set(i, &parent);
+                swapped = true;
+            } else {
+                break;
+            }
+            i = p;
+        }
+        if swapped {
+            self.0.set(i, item);
+        }
+    }
+
+    /// PRECONDITION: self.0.get(i) == item
+    fn bubble_down(&mut self, mut i: u64, n: u64, item: &T) {
+        // We set the flag if self.0.get(i) does not contain the item anymore.
+        let mut swapped = false;
+        // LOOP INVARIANT: HeapInvariant(self, 0, i)
+        loop {
+            let l = i * 2 + 1;
+            let r = l + 1;
+
+            if n <= l {
+                break;
+            }
+
+            if n <= r {
+                // Only the left child is within the array bounds.
+
+                let left = self.0.get(l).unwrap();
+                if is_less(&left, item) {
+                    self.0.set(i, &left);
+                    swapped = true;
+                    i = l;
+                    continue;
+                }
+            } else {
+                // Both children are within the array bounds.
+
+                let left = self.0.get(l).unwrap();
+                let right = self.0.get(r).unwrap();
+
+                let (min_index, min_elem) = if is_less(&left, &right) {
+                    (l, &left)
+                } else {
+                    (r, &right)
+                };
+
+                if is_less(min_elem, item) {
+                    self.0.set(i, min_elem);
+                    swapped = true;
+                    i = min_index;
+                    continue;
+                }
+            }
+            break;
+        }
+        if swapped {
+            self.0.set(i, item);
+        }
+    }
+}
+
+fn is_less<T: PartialOrd>(x: &T, y: &T) -> bool {
+    x.partial_cmp(y) == Some(std::cmp::Ordering::Less)
+}
+
+impl<T, M> fmt::Debug for MinHeap<T, M>
+where
+    T: BoundedStorable + PartialOrd + fmt::Debug,
+    M: Memory,
+{
+    fn fmt(&self, fmt: &mut fmt::Formatter<'_>) -> fmt::Result {
+        self.0.fmt(fmt)
+    }
+}