diff --git a/aptos-move/framework/supra-stdlib/sources/heap.move b/aptos-move/framework/supra-stdlib/sources/heap.move new file mode 100644 index 0000000000000..555b65e383ed1 --- /dev/null +++ b/aptos-move/framework/supra-stdlib/sources/heap.move @@ -0,0 +1,183 @@ +module supra_std::heap { + use std::vector; + use std::option::{Self, Option}; + + const EHEAP_INDEX_OUT_OF_BOUNDS: u64 = 1; + const EREMOVAL_FROM_EMPTY_HEAP: u64 = 2; + const ENON_EMPTY_HEAP: u64 = 3; + + struct Heap has store { + elements: vector, + } + + public fun new(): Heap { + Heap { elements: vector::empty(), } + } + + // Read-only borrow of the element at `index` + public fun borrow(heap: &Heap, index: u64): &T { + assert!(size(heap) > index, EHEAP_INDEX_OUT_OF_BOUNDS); + vector::borrow(&heap.elements, index) + } + + // Swaps elements at indices `i` and `j` + // CAUTION: This method must not be used or called outside of this module. This method has been made public + // only to allow `add` and `remove` inline methods to be able to call this with generic `cmp` comparision function + // DO NOT USE THIS METHOD DIRECTLY, We will break this in future once `Function Values` are supported + public fun unsafe_swap(heap: &mut Heap, i: u64, j: u64) { + vector::swap(&mut heap.elements, i, j); + } + + // Precondition: n == vector::length(&heap.elements) AND `heap` already maintains heap property with respect to `cmp` comparision function except + // at `index` where `index` may be sub-optimal with respect to its descendents with respect to `cmp` + // Postcondition: `heap` maintains heap property with respect to `cmp` function + + // public inline fun sift_down(heap: &mut Heap, index: u64, n: u64, cmp:|&T,&T|bool) { + // let opt = index; + // let left = 2 * index + 1; + // let right = 2 * index + 2; + // + // while(opt < n) { + // let left = 2 * opt + 1; + // let right = 2 * opt + 2; + // if (left < n && cmp(supra_std::heap::borrow(heap,left), supra_std::heap::borrow(heap,opt))) { + // opt = left; + // }; + // if (right < n && cmp(supra_std::heap::borrow(heap,right), supra_std::heap::borrow(heap,opt))) { + // opt = right; + // }; + // if (opt != index) { + // supra_std::heap::unsafe_swap(heap, index, opt); + // } + // else { + // break; + // }; + // } + // + // } + + // Returns the `size` of the `heap` + public fun size(heap: &Heap): u64 { + vector::length(&heap.elements) + } + + // Returns the read-only reference to the elements of the heap + public fun borrow_elems(heap: &Heap): &vector { + &heap.elements + } + + //CAUTION: This method must never be called from outside of this module + //This method is made public only to allow it to be called by `add`, which has to be an `inline` + // method to receive a lambda for comparison + // DO NOT USE THIS METHOD DIRECTLY, We will break this in future once `Function Values` are supported + public fun unsafe_push_back(heap: &mut Heap, value: T) { + vector::push_back(&mut heap.elements, value); + } + + // Precondition: `heap` must already have the `heap` property with respect to `cmp` comparison function + // Postcondition: A new element `value` is added to the `heap` while retaining the `heap` property with respect to `cmp` function + public inline fun add(heap: &mut Heap, value: T, cmp: |&T, &T| bool) { + supra_std::heap::unsafe_push_back(heap, value); + let index = supra_std::heap::size(heap) - 1; + while (index > 0) { + let parent: u64 = if (index & 1 == 1) (index - 1) >> 1 else index >> 1; + + if (cmp( + supra_std::heap::borrow(heap, index), + supra_std::heap::borrow(heap, parent), + )) { + supra_std::heap::unsafe_swap(heap, index, parent); + index = parent; + } else { break } + } + } + + //Precondition: `heap` must be having `heap` property with respect to `cmp` comparison function + //Postcondition: If `heap` is non-empty it will return `option::some` of the optimum element (`max` or `min` as defined by `cmp` function), and it will + // leave the `heap` that continues to maintain `heap` property with respect to `cmp` function + public inline fun remove(heap: &mut Heap, cmp: |&T, &T| bool): Option { + + let result = option::none(); + let n = supra_std::heap::size(heap); + if (n > 0) { + supra_std::heap::unsafe_swap(heap, 0, n - 1); + let index = 0; + result = option::some(supra_std::heap::pop_back(heap)); + n = n - 1; + let opt = index; + + //The loop is guaranteed to terminate because if `opt` does not increase, it must necessarily go in the `else` branch and `break` + while (opt <= (n / 2)) { + let left = 2 * opt + 1; + let right = 2 * opt + 2; + if (left < n + && cmp( + supra_std::heap::borrow(heap, left), + supra_std::heap::borrow(heap, opt), + )) { + opt = left; + }; + if (right < n + && cmp( + supra_std::heap::borrow(heap, right), + supra_std::heap::borrow(heap, opt), + )) { + opt = right; + }; + if (opt != index) { + supra_std::heap::unsafe_swap(heap, index, opt); + index = opt; + } else { + break; + }; + } + }; + + result + } + + // Removes the last element from the heap, this should still maintain the heap property + // Precondition: `heap` must not be empty + public fun pop_back(heap: &mut Heap): T { + vector::pop_back(&mut heap.elements) + } + + // Returns true if the heap is empty, false otherwise + public fun is_empty(h: &Heap): bool { + vector::is_empty(&h.elements) + } + + // Destroys an empty heap + public fun destroy_empty(h: Heap) { + assert!(is_empty(&h), ENON_EMPTY_HEAP); + vector::destroy_empty(h.elements); + let Heap { elements } = h; + } + + // Destroy a heap and removes/drop all the elements if it is non-empty + public fun destroy(h: Heap) { + clear(&mut h); + destroy_empty(h); + } + + // Returns a copy of the heap contents + public fun get_heap_contents(heap: &Heap): vector { + heap.elements + } + + // Removes all the elements from the heap + public fun clear(h: &mut Heap) { + while (vector::length(&h.elements) > 0) { + vector::pop_back(&mut h.elements); + } + } + + // If the heap is non-empty, it returns the copy of the first element + public fun get_top_element(heap: &Heap): Option { + if (is_empty(heap)) { + option::none() + } else { + option::some(*vector::borrow(&heap.elements, 0)) + } + } +} diff --git a/aptos-move/framework/supra-stdlib/tests/heap_tests.move b/aptos-move/framework/supra-stdlib/tests/heap_tests.move new file mode 100644 index 0000000000000..d72de27041302 --- /dev/null +++ b/aptos-move/framework/supra-stdlib/tests/heap_tests.move @@ -0,0 +1,119 @@ +#[test_only] +module supra_std::heap_tests { + use supra_std::heap; + use std::vector; + use std::option; + use aptos_std::debug; + + #[test_only] + inline fun is_heap(v: &vector, cmp: |&T, &T| bool): bool { + let n = vector::length(v); + let result = true; + for (i in 0..n) { + let left = 2 * i + 1; + let right = 2 * i + 2; + if (left < n && !cmp(vector::borrow(v, i), vector::borrow(v, left))) { + result = false; + }; + if (right < n && !cmp(vector::borrow(v, i), vector::borrow(v, right))) { + result = false; + } + }; + result + } + + #[test] + public fun test_heapify() { + let v = vector[78, 32, 33, 34, 12, 21, 113, 321, 897]; + let heap = heap::new(); + vector::for_each( + v, + |item| { + heap::add(&mut heap, item, |a, b| { *a < *b }); + } + ); + + assert!(is_heap(&heap::get_heap_contents(&heap), |a, b| { *a < *b }), 1); + heap::clear(&mut heap); + heap::destroy_empty(heap); + } + + #[test] + public fun test_heap_add_remove() { + let v = vector[78, 32, 33, 34, 12, 21, 113, 321, 897]; + let heap = heap::new(); + vector::for_each( + v, + |item| { + heap::add(&mut heap, item, |a, b| { *a < *b }); + } + ); + + assert!(is_heap(&heap::get_heap_contents(&heap), |a, b| { *a < *b }), 1); + + let removed = heap::remove(&mut heap, |a, b| { *a < *b }); + assert!(option::is_some(&removed), 3); + assert!(option::extract(&mut removed) == 12, 2); + + removed = heap::remove(&mut heap, |a, b| { *a < *b }); + assert!(option::is_some(&removed), 3); + assert!(option::extract(&mut removed) == 21, 2); + + heap::add(&mut heap, 2, |a, b| { *a < *b }); + removed = heap::remove(&mut heap, |a, b| { *a < *b }); + assert!(option::is_some(&removed), 3); + assert!(option::extract(&mut removed) == 2, 2); + + assert!(is_heap(&heap::get_heap_contents(&heap), |a, b| { *a < *b }), 1); + heap::clear(&mut heap); + heap::destroy_empty(heap); + } + + #[test] + public fun test_heap_pop() { + let v = vector[78, 32, 33, 34, 12, 21, 113, 321, 897]; + let heap = heap::new(); + vector::for_each( + v, + |item| { + heap::add(&mut heap, item, |a, b| { *a < *b }); + }, + ); + + assert!(is_heap(&heap::get_heap_contents(&heap), |a, b| { *a < *b }), 1); + + let popped = heap::pop_back(&mut heap); + assert!(popped == 897, 2); + assert!(is_heap(&heap::get_heap_contents(&heap), |a, b| { *a < *b }), 1); + heap::clear(&mut heap); + heap::destroy_empty(heap); + } + + #[test] + public fun test_not_heap_afer_unsafe_swap() { + let v = vector[78, 32, 33, 34, 12, 21, 113, 321, 897]; + let heap = heap::new(); + vector::for_each( + v, + |item| { + heap::add(&mut heap, item, |a, b| { *a < *b }); + } + ); + + assert!(is_heap(&heap::get_heap_contents(&heap), |a, b| { *a < *b }), 1); + + supra_std::heap::unsafe_swap(&mut heap, 0, 1); + assert!(!is_heap(&heap::get_heap_contents(&heap), |a, b| { *a < *b }), 1); + //heap::clear(&mut heap); + //heap::destroy_empty(heap); + heap::destroy(heap); + } + + #[test] + public fun test_removal_from_empty_heap_returns_none() { + let heap = heap::new(); + let removed = heap::remove(&mut heap, |a, b| { *a < *b }); + assert!(option::is_none(&removed), 1); + heap::destroy_empty(heap); + } +}