Skip to main content

Module 0xdee9::critbit

use 0x2::table;
use 0x2::tx_context;
use 0xdee9::math;

Struct Leaf​

struct Leaf<V> has drop, store
Fields
key: u64
value: V
parent: u64

Struct InternalNode​

struct InternalNode has drop, store
Fields
mask: u64
left_child: u64
right_child: u64
parent: u64

Struct CritbitTree​

struct CritbitTree<V: store> has store
Fields
root: u64
internal_nodes: table::Table<u64, critbit::InternalNode>
leaves: table::Table<u64, critbit::Leaf<V>>
min_leaf: u64
max_leaf: u64
next_internal_node_index: u64
next_leaf_index: u64

Constants​

const EExceedCapacity: u64 = 2;


const EIndexOutOfRange: u64 = 7;


const EKeyAlreadyExist: u64 = 4;


const ELeafNotExist: u64 = 5;


const ENullParent: u64 = 8;


const ETreeNotEmpty: u64 = 3;


const MAX_CAPACITY: u64 = 9223372036854775807;


const MAX_U64: u64 = 18446744073709551615;


const PARTITION_INDEX: u64 = 9223372036854775808;

Function new​

public(friend) fun new<V: store>(ctx: &mut tx_context::TxContext): critbit::CritbitTree<V>
Implementation
public(package) fun new<V: store>(ctx: &mut TxContext): CritbitTree<V> {
    CritbitTree<V> {
        root: PARTITION_INDEX,
        internal_nodes: table::new(ctx),
        leaves: table::new(ctx),
        min_leaf: PARTITION_INDEX,
        max_leaf: PARTITION_INDEX,
        next_internal_node_index: 0,
        next_leaf_index: 0
    }
}

Function size​

public(friend) fun size<V: store>(tree: &critbit::CritbitTree<V>): u64
Implementation
public(package) fun size<V: store>(tree: &CritbitTree<V>): u64 {
    table::length(&tree.leaves)
}

Function is_empty​

public(friend) fun is_empty<V: store>(tree: &critbit::CritbitTree<V>): bool
Implementation
public(package) fun is_empty<V: store>(tree: &CritbitTree<V>): bool {
    table::is_empty(&tree.leaves)
}

Function min_leaf​

public fun min_leaf<V: store>(tree: &critbit::CritbitTree<V>): (u64, u64)
Implementation
public fun min_leaf<V: store>(tree: &CritbitTree<V>): (u64, u64) {
    assert!(!is_empty(tree), ELeafNotExist);
    let min_leaf = table::borrow(&tree.leaves, tree.min_leaf);
    return (min_leaf.key, tree.min_leaf)
}

Function max_leaf​

public fun max_leaf<V: store>(tree: &critbit::CritbitTree<V>): (u64, u64)
Implementation
public fun max_leaf<V: store>(tree: &CritbitTree<V>): (u64, u64) {
    assert!(!is_empty(tree), ELeafNotExist);
    let max_leaf = table::borrow(&tree.leaves, tree.max_leaf);
    return (max_leaf.key, tree.max_leaf)
}

Function previous_leaf​

public fun previous_leaf<V: store>(tree: &critbit::CritbitTree<V>, key: u64): (u64, u64)
Implementation
public fun previous_leaf<V: store>(tree: &CritbitTree<V>, key: u64): (u64, u64) {
    let (_, mut index) = find_leaf(tree, key);
    assert!(index != PARTITION_INDEX, ELeafNotExist);
    let mut ptr = MAX_U64 - index;
    let mut parent = table::borrow(&tree.leaves, index).parent;
    while (parent != PARTITION_INDEX && is_left_child(tree, parent, ptr)){
        ptr = parent;
        parent = table::borrow(&tree.internal_nodes, ptr).parent;
    };
    if(parent == PARTITION_INDEX) {
        return (0, PARTITION_INDEX)
    };
    index = MAX_U64 - right_most_leaf(tree, table::borrow(&tree.internal_nodes, parent).left_child);
    let key = table::borrow(&tree.leaves, index).key;
    return (key, index)
}

Function next_leaf​

public fun next_leaf<V: store>(tree: &critbit::CritbitTree<V>, key: u64): (u64, u64)
Implementation
public fun next_leaf<V: store>(tree: &CritbitTree<V>, key: u64): (u64, u64) {
    let (_, mut index) = find_leaf(tree, key);
    assert!(index != PARTITION_INDEX, ELeafNotExist);
    let mut ptr = MAX_U64 - index;
    let mut parent = table::borrow(&tree.leaves, index).parent;
    while (parent != PARTITION_INDEX && !is_left_child(tree, parent, ptr)){
        ptr = parent;
        parent = table::borrow(&tree.internal_nodes, ptr).parent;
    };
    if(parent == PARTITION_INDEX) {
        return (0, PARTITION_INDEX)
    };
    index = MAX_U64 - left_most_leaf(tree, table::borrow(&tree.internal_nodes, parent).right_child);
    let key = table::borrow(&tree.leaves, index).key;
    return (key, index)
}

Function left_most_leaf​

fun left_most_leaf<V: store>(tree: &critbit::CritbitTree<V>, root: u64): u64
Implementation
fun left_most_leaf<V: store>(tree: &CritbitTree<V>, root: u64): u64 {
    let mut ptr = root;
    while (ptr < PARTITION_INDEX){
        ptr = table::borrow(& tree.internal_nodes, ptr).left_child;
    };
    ptr
}

Function right_most_leaf​

fun right_most_leaf<V: store>(tree: &critbit::CritbitTree<V>, root: u64): u64
Implementation
fun right_most_leaf<V: store>(tree: &CritbitTree<V>, root: u64): u64 {
    let mut ptr = root;
    while (ptr < PARTITION_INDEX){
        ptr = table::borrow(& tree.internal_nodes, ptr).right_child;
    };
    ptr
}

Function insert_leaf​

public(friend) fun insert_leaf<V: store>(tree: &mut critbit::CritbitTree<V>, key: u64, value: V): u64
Implementation
public(package) fun insert_leaf<V: store>(tree: &mut CritbitTree<V>, key: u64, value: V): u64 {
    let new_leaf = Leaf<V>{
        key,
        value,
        parent: PARTITION_INDEX,
    };
    let new_leaf_index = tree.next_leaf_index;
    tree.next_leaf_index = tree.next_leaf_index + 1;
    assert!(new_leaf_index < MAX_CAPACITY - 1, EExceedCapacity);
    table::add(&mut tree.leaves, new_leaf_index, new_leaf);

    let closest_leaf_index = get_closest_leaf_index_by_key(tree, key);

    // Handle the first insertion
    if (closest_leaf_index == PARTITION_INDEX) {
        assert!(new_leaf_index == 0, ETreeNotEmpty);
        tree.root = MAX_U64 - new_leaf_index;
        tree.min_leaf = new_leaf_index;
        tree.max_leaf = new_leaf_index;
        return 0
    };

    let closest_key = table::borrow(&tree.leaves, closest_leaf_index).key;
    assert!(closest_key != key, EKeyAlreadyExist);

    // Note that we reserve count_leading_zeros of form u128 for future use
    let critbit = 64 - (count_leading_zeros((closest_key ^ key) as u128) - 64);
    let new_mask = 1u64 << (critbit - 1);

    let new_internal_node= InternalNode {
        mask: new_mask,
        left_child: PARTITION_INDEX,
        right_child: PARTITION_INDEX,
        parent: PARTITION_INDEX,
    };
    let new_internal_node_index = tree.next_internal_node_index;
    tree.next_internal_node_index = tree.next_internal_node_index + 1;
    table::add(&mut tree.internal_nodes, new_internal_node_index, new_internal_node);

    let mut ptr = tree.root;
    let mut new_internal_node_parent_index = PARTITION_INDEX;
    // Search position of the new internal node
    while (ptr < PARTITION_INDEX) {
        let internal_node = table::borrow(&tree.internal_nodes, ptr);
        if (new_mask > internal_node.mask) {
            break
        };
        new_internal_node_parent_index = ptr;
        if (key & internal_node.mask == 0) {
            ptr = internal_node.left_child;
        } else {
            ptr = internal_node.right_child;
        };
    };

    // Update the child info of new internal node's parent
    if (new_internal_node_parent_index == PARTITION_INDEX){
        // if the new internal node is root
        tree.root = new_internal_node_index;
    } else{
        // In another case, we update the child field of the new internal node's parent
        // And the parent field of the new internal node
        let is_left_child = is_left_child(tree, new_internal_node_parent_index, ptr);
        update_child(tree, new_internal_node_parent_index, new_internal_node_index, is_left_child);
    };

    // Finally, update the child field of the new internal node
    let is_left_child = new_mask & key == 0;
    update_child(tree, new_internal_node_index, MAX_U64 - new_leaf_index, is_left_child);
    update_child(tree, new_internal_node_index, ptr, !is_left_child);

    if (table::borrow(&tree.leaves, tree.min_leaf).key > key) {
        tree.min_leaf = new_leaf_index;
    };
    if (table::borrow(&tree.leaves, tree.max_leaf).key < key) {
        tree.max_leaf = new_leaf_index;
    };
    new_leaf_index
}

Function find_leaf​

public fun find_leaf<V: store>(tree: &critbit::CritbitTree<V>, key: u64): (bool, u64)
Implementation
public fun find_leaf<V: store>(tree: & CritbitTree<V>, key: u64): (bool, u64) {
    if (is_empty(tree)) {
        return (false, PARTITION_INDEX)
    };
    let closest_leaf_index = get_closest_leaf_index_by_key(tree, key);
    let closeset_leaf = table::borrow(&tree.leaves, closest_leaf_index);
    if (closeset_leaf.key != key){
        return (false, PARTITION_INDEX)
    } else{
        return (true, closest_leaf_index)
    }
}

Function find_closest_key​

public(friend) fun find_closest_key<V: store>(tree: &critbit::CritbitTree<V>, key: u64): u64
Implementation
public(package) fun find_closest_key<V: store>(tree: & CritbitTree<V>, key: u64): u64 {
    if (is_empty(tree)) {
        return 0
    };
    let closest_leaf_index = get_closest_leaf_index_by_key(tree, key);
    let closeset_leaf = table::borrow(&tree.leaves, closest_leaf_index);
    closeset_leaf.key
}

Function remove_leaf_by_index​

public(friend) fun remove_leaf_by_index<V: store>(tree: &mut critbit::CritbitTree<V>, index: u64): V
Implementation
public(package) fun remove_leaf_by_index<V: store>(tree: &mut CritbitTree<V>, index: u64): V {
    let key = table::borrow(& tree.leaves, index).key;
    if (tree.min_leaf == index) {
        let (_, index) = next_leaf(tree, key);
        tree.min_leaf = index;
    };
    if (tree.max_leaf == index) {
        let (_, index) = previous_leaf(tree, key);
        tree.max_leaf = index;
    };

    let mut is_left_child_;
    let Leaf<V> {key: _, value, parent: removed_leaf_parent_index} = table::remove(&mut tree.leaves, index);

    if (size(tree) == 0) {
        tree.root = PARTITION_INDEX;
        tree.min_leaf = PARTITION_INDEX;
        tree.max_leaf = PARTITION_INDEX;
        tree.next_internal_node_index = 0;
        tree.next_leaf_index = 0;
    } else {
        assert!(removed_leaf_parent_index != PARTITION_INDEX, EIndexOutOfRange);
        let removed_leaf_parent = table::borrow(&tree.internal_nodes, removed_leaf_parent_index);
        let removed_leaf_grand_parent_index = removed_leaf_parent.parent;

        // Note that sibling of the removed leaf can be a leaf or an internal node
        is_left_child_ = is_left_child(tree, removed_leaf_parent_index, MAX_U64 - index);
        let sibling_index = if (is_left_child_) { removed_leaf_parent.right_child }
        else { removed_leaf_parent.left_child };

        if (removed_leaf_grand_parent_index == PARTITION_INDEX) {
            // Parent of the removed leaf is the tree root
            // Update the parent of the sibling node and set sibling as the tree root
            if (sibling_index < PARTITION_INDEX) {
                // sibling is an internal node
                table::borrow_mut(&mut tree.internal_nodes, sibling_index).parent = PARTITION_INDEX;
            } else{
                // sibling is a leaf
                table::borrow_mut(&mut tree.leaves, MAX_U64 - sibling_index).parent = PARTITION_INDEX;
            };
            tree.root = sibling_index;
        } else {
            // grand parent of the removed leaf is a internal node
            // set sibling as the child of the grand parent of the removed leaf
            is_left_child_ = is_left_child(tree, removed_leaf_grand_parent_index, removed_leaf_parent_index);
            update_child(tree, removed_leaf_grand_parent_index, sibling_index, is_left_child_);
        };
        table::remove(&mut tree.internal_nodes, removed_leaf_parent_index);
    };
    value
}

Function borrow_mut_leaf_by_index​

public(friend) fun borrow_mut_leaf_by_index<V: store>(tree: &mut critbit::CritbitTree<V>, index: u64): &mut V
Implementation
public(package) fun borrow_mut_leaf_by_index<V: store>(tree: &mut CritbitTree<V>, index: u64): &mut V {
    let entry = table::borrow_mut(&mut tree.leaves, index);
    &mut entry.value
}

Function borrow_leaf_by_index​

public fun borrow_leaf_by_index<V: store>(tree: &critbit::CritbitTree<V>, index: u64): &V
Implementation
public fun borrow_leaf_by_index<V: store>(tree: & CritbitTree<V>, index: u64): &V {
    let entry = table::borrow(&tree.leaves, index);
    &entry.value
}

Function borrow_leaf_by_key​

public fun borrow_leaf_by_key<V: store>(tree: &critbit::CritbitTree<V>, key: u64): &V
Implementation
public fun borrow_leaf_by_key<V: store>(tree: & CritbitTree<V>, key: u64): &V {
    let (is_exist, index) = find_leaf(tree, key);
    assert!(is_exist, ELeafNotExist);
    borrow_leaf_by_index(tree, index)
}

Function drop​

public(friend) fun drop<V: drop, store>(tree: critbit::CritbitTree<V>)
Implementation
public(package) fun drop<V: store + drop>(tree: CritbitTree<V>) {
    let CritbitTree<V> {
        root: _,
        internal_nodes,
        leaves,
        min_leaf: _,
        max_leaf: _,
        next_internal_node_index: _,
        next_leaf_index: _,

    } = tree;
    table::drop(internal_nodes);
    table::drop(leaves);
}

Function destroy_empty​

public(friend) fun destroy_empty<V: store>(tree: critbit::CritbitTree<V>)
Implementation
public(package) fun destroy_empty<V: store>(tree: CritbitTree<V>) {
    assert!(table::length(&tree.leaves) == 0, 0);

    let CritbitTree<V> {
        root: _,
        leaves,
        internal_nodes,
        min_leaf: _,
        max_leaf: _,
        next_internal_node_index: _,
        next_leaf_index: _
    } = tree;

    table::destroy_empty(leaves);
    table::destroy_empty(internal_nodes);
}

Function get_closest_leaf_index_by_key​

fun get_closest_leaf_index_by_key<V: store>(tree: &critbit::CritbitTree<V>, key: u64): u64
Implementation
fun get_closest_leaf_index_by_key<V: store>(tree: &CritbitTree<V>, key: u64): u64 {
    let mut ptr = tree.root;
    // if tree is empty, return the patrition index
    if(ptr == PARTITION_INDEX) return PARTITION_INDEX;
    while (ptr < PARTITION_INDEX){
        let node = table::borrow(&tree.internal_nodes, ptr);
        if (key & node.mask == 0){
            ptr = node.left_child;
        } else {
            ptr = node.right_child;
        }
    };
    return (MAX_U64 - ptr)
}

Function update_child​

fun update_child<V: store>(tree: &mut critbit::CritbitTree<V>, parent_index: u64, new_child: u64, is_left_child: bool)
Implementation
fun update_child<V: store>(tree: &mut CritbitTree<V>, parent_index: u64, new_child: u64, is_left_child: bool) {
    assert!(parent_index != PARTITION_INDEX, ENullParent);
    if (is_left_child) {
        table::borrow_mut(&mut tree.internal_nodes, parent_index).left_child = new_child;
    } else{
        table::borrow_mut(&mut tree.internal_nodes, parent_index).right_child = new_child;
    };
    if (new_child > PARTITION_INDEX) {
        table::borrow_mut(&mut tree.leaves, MAX_U64 - new_child).parent = parent_index;
    } else {
        table::borrow_mut(&mut tree.internal_nodes, new_child).parent = parent_index;
    }
}

Function is_left_child​

fun is_left_child<V: store>(tree: &critbit::CritbitTree<V>, parent_index: u64, index: u64): bool
Implementation
fun is_left_child<V: store>(tree: &CritbitTree<V>, parent_index: u64, index: u64): bool {
    table::borrow(&tree.internal_nodes, parent_index).left_child == index
}