mirror of
https://github.com/Ed94/Odin.git
synced 2026-08-05 15:18:49 +00:00
Add core:container/rbtree
Add a red-black tree with configurable $Key and $Value. Also includes tests that verify it maintains RB invariants, doesn't leak. Originally based on the CC0 implementation from literateprograms.org. But reworked to the same API used by @Yawning's excellent `core:container/avl` for ease of use.
This commit is contained in:
@@ -5,13 +5,10 @@ The implementation is non-intrusive, and non-recursive.
|
|||||||
*/
|
*/
|
||||||
package container_avl
|
package container_avl
|
||||||
|
|
||||||
import "base:intrinsics"
|
@(require) import "base:intrinsics"
|
||||||
import "base:runtime"
|
@(require) import "base:runtime"
|
||||||
import "core:slice"
|
import "core:slice"
|
||||||
|
|
||||||
_ :: intrinsics
|
|
||||||
_ :: runtime
|
|
||||||
|
|
||||||
// Originally based on the CC0 implementation by Eric Biggers
|
// Originally based on the CC0 implementation by Eric Biggers
|
||||||
// See: https://github.com/ebiggers/avl_tree/
|
// See: https://github.com/ebiggers/avl_tree/
|
||||||
|
|
||||||
@@ -675,4 +672,4 @@ iterator_first :: proc "contextless" (it: ^Iterator($Value)) {
|
|||||||
if it._cur != nil {
|
if it._cur != nil {
|
||||||
it._next = node_next_or_prev_in_order(it._cur, it._direction)
|
it._next = node_next_or_prev_in_order(it._cur, it._direction)
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
@@ -0,0 +1,567 @@
|
|||||||
|
// This package implements a red-black tree
|
||||||
|
package container_intrusive_rbtree
|
||||||
|
|
||||||
|
@(require) import "base:intrinsics"
|
||||||
|
@(require) import "base:runtime"
|
||||||
|
import "core:slice"
|
||||||
|
|
||||||
|
// Originally based on the CC0 implementation from literateprograms.org
|
||||||
|
// But with API design mimicking `core:container/avl` for ease of use.
|
||||||
|
|
||||||
|
// Direction specifies the traversal direction for a tree iterator.
|
||||||
|
Direction :: enum i8 {
|
||||||
|
// Backward is the in-order backwards direction.
|
||||||
|
Backward = -1,
|
||||||
|
// Forward is the in-order forwards direction.
|
||||||
|
Forward = 1,
|
||||||
|
}
|
||||||
|
|
||||||
|
Ordering :: slice.Ordering
|
||||||
|
|
||||||
|
// Tree is a red-black tree
|
||||||
|
Tree :: struct($Key: typeid, $Value: typeid) {
|
||||||
|
// user_data is a parameter that will be passed to the on_remove
|
||||||
|
// callback.
|
||||||
|
user_data: rawptr,
|
||||||
|
// on_remove is an optional callback that can be called immediately
|
||||||
|
// after a node is removed from the tree.
|
||||||
|
on_remove: proc(key: Key, value: Value, user_data: rawptr),
|
||||||
|
|
||||||
|
_root: ^Node(Key, Value),
|
||||||
|
_node_allocator: runtime.Allocator,
|
||||||
|
_cmp_fn: proc(Key, Key) -> Ordering,
|
||||||
|
_size: int,
|
||||||
|
}
|
||||||
|
|
||||||
|
// Node is a red-black tree node.
|
||||||
|
//
|
||||||
|
// WARNING: It is unsafe to mutate value if the node is part of a tree
|
||||||
|
// if doing so will alter the Node's sort position relative to other
|
||||||
|
// elements in the tree.
|
||||||
|
Node :: struct($Key: typeid, $Value: typeid) {
|
||||||
|
key: Key,
|
||||||
|
value: Value,
|
||||||
|
|
||||||
|
_parent: ^Node(Key, Value),
|
||||||
|
_left: ^Node(Key, Value),
|
||||||
|
_right: ^Node(Key, Value),
|
||||||
|
_color: Color,
|
||||||
|
}
|
||||||
|
|
||||||
|
// Might store this in the node pointer in the future, but that'll require a decent amount of rework to pass ^^N instead of ^N
|
||||||
|
Color :: enum uintptr {Black = 0, Red = 1}
|
||||||
|
|
||||||
|
// Iterator is a tree iterator.
|
||||||
|
//
|
||||||
|
// WARNING: It is unsafe to modify the tree while iterating, except via
|
||||||
|
// the iterator_remove method.
|
||||||
|
Iterator :: struct($Key: typeid, $Value: typeid) {
|
||||||
|
_tree: ^Tree(Key, Value),
|
||||||
|
_cur: ^Node(Key, Value),
|
||||||
|
_next: ^Node(Key, Value),
|
||||||
|
_direction: Direction,
|
||||||
|
_called_next: bool,
|
||||||
|
}
|
||||||
|
|
||||||
|
// init initializes a tree.
|
||||||
|
init :: proc {
|
||||||
|
init_ordered,
|
||||||
|
init_cmp,
|
||||||
|
}
|
||||||
|
|
||||||
|
// init_cmp initializes a tree.
|
||||||
|
init_cmp :: proc(t: ^$T/Tree($Key, $Value), cmp_fn: proc(a, b: Key) -> Ordering, node_allocator := context.allocator) {
|
||||||
|
t._root = nil
|
||||||
|
t._node_allocator = node_allocator
|
||||||
|
t._cmp_fn = cmp_fn
|
||||||
|
t._size = 0
|
||||||
|
}
|
||||||
|
|
||||||
|
// init_ordered initializes a tree containing ordered keys, with
|
||||||
|
// a comparison function that results in an ascending order sort.
|
||||||
|
init_ordered :: proc(t: ^$T/Tree($Key, $Value), node_allocator := context.allocator) where intrinsics.type_is_ordered_numeric(Key) {
|
||||||
|
init_cmp(t, slice.cmp_proc(Key), node_allocator)
|
||||||
|
}
|
||||||
|
|
||||||
|
// destroy de-initializes a tree.
|
||||||
|
destroy :: proc(t: ^$T/Tree($Key, $Value), call_on_remove: bool = true) {
|
||||||
|
iter := iterator(t, .Forward)
|
||||||
|
for _ in iterator_next(&iter) {
|
||||||
|
iterator_remove(&iter, call_on_remove)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
len :: proc "contextless" (t: ^$T/Tree($Key, $Value)) -> (node_count: int) {
|
||||||
|
return t._size
|
||||||
|
}
|
||||||
|
|
||||||
|
// first returns the first node in the tree (in-order) or nil iff
|
||||||
|
// the tree is empty.
|
||||||
|
first :: proc "contextless" (t: ^$T/Tree($Key, $Value)) -> ^Node(Key, Value) {
|
||||||
|
return tree_first_or_last_in_order(t, Direction.Backward)
|
||||||
|
}
|
||||||
|
|
||||||
|
// last returns the last element in the tree (in-order) or nil iff
|
||||||
|
// the tree is empty.
|
||||||
|
last :: proc "contextless" (t: ^$T/Tree($Key, $Value)) -> ^Node(Key, Value) {
|
||||||
|
return tree_first_or_last_in_order(t, Direction.Forward)
|
||||||
|
}
|
||||||
|
|
||||||
|
// find finds the key in the tree, and returns the corresponding node, or nil iff the value is not present.
|
||||||
|
find :: proc(t: ^$T/Tree($Key, $Value), key: Key) -> (node: ^Node(Key, Value)) {
|
||||||
|
node = t._root
|
||||||
|
for node != nil {
|
||||||
|
switch t._cmp_fn(key, node.key) {
|
||||||
|
case .Equal: return node
|
||||||
|
case .Less: node = node._left
|
||||||
|
case .Greater: node = node._right
|
||||||
|
}
|
||||||
|
}
|
||||||
|
return node
|
||||||
|
}
|
||||||
|
|
||||||
|
// find_value finds the key in the tree, and returns the corresponding value, or nil iff the value is not present.
|
||||||
|
find_value :: proc(t: ^$T/Tree($Key, $Value), key: Key) -> (value: Value, ok: bool) #optional_ok {
|
||||||
|
if n := find(t, key); n != nil {
|
||||||
|
return n.value, true
|
||||||
|
}
|
||||||
|
return
|
||||||
|
}
|
||||||
|
|
||||||
|
// find_or_insert attempts to insert the value into the tree, and returns
|
||||||
|
// the node, a boolean indicating if the value was inserted, and the
|
||||||
|
// node allocator error if relevant. If the value is already present, the existing node is updated.
|
||||||
|
find_or_insert :: proc(t: ^$T/Tree($Key, $Value), key: Key, value: Value) -> (n: ^Node(Key, Value), inserted: bool) {
|
||||||
|
n_ptr := &t._root
|
||||||
|
for n_ptr^ != nil {
|
||||||
|
n = n_ptr^
|
||||||
|
switch t._cmp_fn(key, n.key) {
|
||||||
|
case .Less:
|
||||||
|
n_ptr = &n._left
|
||||||
|
case .Greater:
|
||||||
|
n_ptr = &n._right
|
||||||
|
case .Equal:
|
||||||
|
return
|
||||||
|
}
|
||||||
|
}
|
||||||
|
_parent := n
|
||||||
|
n = new_clone(Node(Key, Value){key=key, value=value, _parent=_parent, _color=.Red}, t._node_allocator) // or_return
|
||||||
|
n_ptr^ = n
|
||||||
|
insert_case1(t, n)
|
||||||
|
t._size += 1
|
||||||
|
return n, true
|
||||||
|
}
|
||||||
|
|
||||||
|
// remove removes a node or value from the tree, and returns true iff the
|
||||||
|
// removal was successful. While the node's value will be left intact,
|
||||||
|
// the node itself will be freed via the tree's node allocator.
|
||||||
|
remove :: proc {
|
||||||
|
remove_key,
|
||||||
|
remove_node,
|
||||||
|
}
|
||||||
|
|
||||||
|
// remove_value removes a value from the tree, and returns true iff the
|
||||||
|
// removal was successful. While the node's key + value will be left intact,
|
||||||
|
// the node itself will be freed via the tree's node allocator.
|
||||||
|
remove_key :: proc(t: ^$T/Tree($Key, $Value), key: Key, call_on_remove := true) -> bool {
|
||||||
|
n := find(t, key)
|
||||||
|
if n == nil {
|
||||||
|
return false // Key not found, nothing to do
|
||||||
|
}
|
||||||
|
return remove_node(t, n, call_on_remove)
|
||||||
|
}
|
||||||
|
|
||||||
|
// remove_node removes a node from the tree, and returns true iff the
|
||||||
|
// removal was successful. While the node's key + value will be left intact,
|
||||||
|
// the node itself will be freed via the tree's node allocator.
|
||||||
|
remove_node :: proc(t: ^$T/Tree($Key, $Value), node: ^$N/Node(Key, Value), call_on_remove := true) -> (found: bool) {
|
||||||
|
if node._parent == node || (node._parent == nil && t._root != node) {
|
||||||
|
return false // Don't touch self-parented or dangling nodes.
|
||||||
|
}
|
||||||
|
node := node
|
||||||
|
if node._left != nil && node._right != nil {
|
||||||
|
// Copy key + value from predecessor and delete it instead
|
||||||
|
predecessor := maximum_node(node._left)
|
||||||
|
node.key = predecessor.key
|
||||||
|
node.value = predecessor.value
|
||||||
|
node = predecessor
|
||||||
|
}
|
||||||
|
|
||||||
|
child := node._right == nil ? node._left : node._right
|
||||||
|
if node_color(node) == .Black {
|
||||||
|
node._color = node_color(child)
|
||||||
|
remove_case1(t, node)
|
||||||
|
}
|
||||||
|
replace_node(t, node, child)
|
||||||
|
if node._parent == nil && child != nil {
|
||||||
|
child._color = .Black // root should be black
|
||||||
|
}
|
||||||
|
|
||||||
|
if call_on_remove && t.on_remove != nil {
|
||||||
|
t.on_remove(node.key, node.value, t.user_data)
|
||||||
|
}
|
||||||
|
free(node, t._node_allocator)
|
||||||
|
t._size -= 1
|
||||||
|
return true
|
||||||
|
}
|
||||||
|
|
||||||
|
// iterator returns a tree iterator in the specified direction.
|
||||||
|
iterator :: proc "contextless" (t: ^$T/Tree($Key, $Value), direction: Direction) -> Iterator(Key, Value) {
|
||||||
|
it: Iterator(Key, Value)
|
||||||
|
it._tree = transmute(^Tree(Key, Value))t
|
||||||
|
it._direction = direction
|
||||||
|
|
||||||
|
iterator_first(&it)
|
||||||
|
|
||||||
|
return it
|
||||||
|
}
|
||||||
|
|
||||||
|
// iterator_from_pos returns a tree iterator in the specified direction,
|
||||||
|
// spanning the range [pos, last] (inclusive).
|
||||||
|
iterator_from_pos :: proc "contextless" (t: ^$T/Tree($Key, $Value), pos: ^Node(Key, Value), direction: Direction) -> Iterator(Key, Value) {
|
||||||
|
it: Iterator(Key, Value)
|
||||||
|
it._tree = transmute(^Tree(Key, Value))t
|
||||||
|
it._direction = direction
|
||||||
|
it._next = nil
|
||||||
|
it._called_next = false
|
||||||
|
|
||||||
|
if it._cur = pos; pos != nil {
|
||||||
|
it._next = node_next_or_prev_in_order(it._cur, it._direction)
|
||||||
|
}
|
||||||
|
|
||||||
|
return it
|
||||||
|
}
|
||||||
|
|
||||||
|
// iterator_get returns the node currently pointed to by the iterator,
|
||||||
|
// or nil iff the node has been removed, the tree is empty, or the end
|
||||||
|
// of the tree has been reached.
|
||||||
|
iterator_get :: proc "contextless" (it: ^$I/Iterator($Key, $Value)) -> ^Node(Key, Value) {
|
||||||
|
return it._cur
|
||||||
|
}
|
||||||
|
|
||||||
|
// iterator_remove removes the node currently pointed to by the iterator,
|
||||||
|
// and returns true iff the removal was successful. Semantics are the
|
||||||
|
// same as the Tree remove.
|
||||||
|
iterator_remove :: proc(it: ^$I/Iterator($Key, $Value), call_on_remove: bool = true) -> bool {
|
||||||
|
if it._cur == nil {
|
||||||
|
return false
|
||||||
|
}
|
||||||
|
|
||||||
|
ok := remove_node(it._tree, it._cur , call_on_remove)
|
||||||
|
if ok {
|
||||||
|
it._cur = nil
|
||||||
|
}
|
||||||
|
|
||||||
|
return ok
|
||||||
|
}
|
||||||
|
|
||||||
|
// iterator_next advances the iterator and returns the (node, true) or
|
||||||
|
// or (nil, false) iff the end of the tree has been reached.
|
||||||
|
//
|
||||||
|
// Note: The first call to iterator_next will return the first node instead
|
||||||
|
// of advancing the iterator.
|
||||||
|
iterator_next :: proc "contextless" (it: ^$I/Iterator($Key, $Value)) -> (^Node(Key, Value), bool) {
|
||||||
|
// This check is needed so that the first element gets returned from
|
||||||
|
// a brand-new iterator, and so that the somewhat contrived case where
|
||||||
|
// iterator_remove is called before the first call to iterator_next
|
||||||
|
// returns the correct value.
|
||||||
|
if !it._called_next {
|
||||||
|
it._called_next = true
|
||||||
|
|
||||||
|
// There can be the contrived case where iterator_remove is
|
||||||
|
// called before ever calling iterator_next, which needs to be
|
||||||
|
// handled as an actual call to next.
|
||||||
|
//
|
||||||
|
// If this happens it._cur will be nil, so only return the
|
||||||
|
// first value, if it._cur is valid.
|
||||||
|
if it._cur != nil {
|
||||||
|
return it._cur, true
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
if it._next == nil {
|
||||||
|
return nil, false
|
||||||
|
}
|
||||||
|
|
||||||
|
it._cur = it._next
|
||||||
|
it._next = node_next_or_prev_in_order(it._cur, it._direction)
|
||||||
|
|
||||||
|
return it._cur, true
|
||||||
|
}
|
||||||
|
|
||||||
|
@(private)
|
||||||
|
tree_first_or_last_in_order :: proc "contextless" (t: ^$T/Tree($Key, $Value), direction: Direction) -> ^Node(Key, Value) {
|
||||||
|
first, sign := t._root, i8(direction)
|
||||||
|
if first != nil {
|
||||||
|
for {
|
||||||
|
tmp := node_get_child(first, sign)
|
||||||
|
if tmp == nil {
|
||||||
|
break
|
||||||
|
}
|
||||||
|
first = tmp
|
||||||
|
}
|
||||||
|
}
|
||||||
|
return first
|
||||||
|
}
|
||||||
|
|
||||||
|
@(private)
|
||||||
|
node_get_child :: #force_inline proc "contextless" (n: ^Node($Key, $Value), sign: i8) -> ^Node(Key, Value) {
|
||||||
|
if sign < 0 {
|
||||||
|
return n._left
|
||||||
|
}
|
||||||
|
return n._right
|
||||||
|
}
|
||||||
|
|
||||||
|
@(private)
|
||||||
|
node_next_or_prev_in_order :: proc "contextless" (n: ^Node($Key, $Value), direction: Direction) -> ^Node(Key, Value) {
|
||||||
|
next, tmp: ^Node(Key, Value)
|
||||||
|
sign := i8(direction)
|
||||||
|
|
||||||
|
if next = node_get_child(n, +sign); next != nil {
|
||||||
|
for {
|
||||||
|
tmp = node_get_child(next, -sign)
|
||||||
|
if tmp == nil {
|
||||||
|
break
|
||||||
|
}
|
||||||
|
next = tmp
|
||||||
|
}
|
||||||
|
} else {
|
||||||
|
tmp, next = n, n._parent
|
||||||
|
for next != nil && tmp == node_get_child(next, +sign) {
|
||||||
|
tmp, next = next, next._parent
|
||||||
|
}
|
||||||
|
}
|
||||||
|
return next
|
||||||
|
}
|
||||||
|
|
||||||
|
@(private)
|
||||||
|
iterator_first :: proc "contextless" (it: ^Iterator($Key, $Value)) {
|
||||||
|
// This is private because behavior when the user manually calls
|
||||||
|
// iterator_first followed by iterator_next is unintuitive, since
|
||||||
|
// the first call to iterator_next MUST return the first node
|
||||||
|
// instead of advancing so that `for node in iterator_next(&next)`
|
||||||
|
// works as expected.
|
||||||
|
|
||||||
|
switch it._direction {
|
||||||
|
case .Forward:
|
||||||
|
it._cur = tree_first_or_last_in_order(it._tree, .Backward)
|
||||||
|
case .Backward:
|
||||||
|
it._cur = tree_first_or_last_in_order(it._tree, .Forward)
|
||||||
|
}
|
||||||
|
|
||||||
|
it._next = nil
|
||||||
|
it._called_next = false
|
||||||
|
|
||||||
|
if it._cur != nil {
|
||||||
|
it._next = node_next_or_prev_in_order(it._cur, it._direction)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
@(private)
|
||||||
|
grand_parent :: proc(n: ^$N/Node($Key, $Value)) -> (g: ^N) {
|
||||||
|
return n._parent._parent
|
||||||
|
}
|
||||||
|
|
||||||
|
@(private)
|
||||||
|
sibling :: proc(n: ^$N/Node($Key, $Value)) -> (s: ^N) {
|
||||||
|
if n == n._parent._left {
|
||||||
|
return n._parent._right
|
||||||
|
} else {
|
||||||
|
return n._parent._left
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
@(private)
|
||||||
|
uncle :: proc(n: ^$N/Node($Key, $Value)) -> (u: ^N) {
|
||||||
|
return sibling(n._parent)
|
||||||
|
}
|
||||||
|
|
||||||
|
@(private)
|
||||||
|
rotate__left :: proc(t: ^$T/Tree($Key, $Value), n: ^$N/Node(Key, Value)) {
|
||||||
|
r := n._right
|
||||||
|
replace_node(t, n, r)
|
||||||
|
n._right = r._left
|
||||||
|
if r._left != nil {
|
||||||
|
r._left._parent = n
|
||||||
|
}
|
||||||
|
r._left = n
|
||||||
|
n._parent = r
|
||||||
|
}
|
||||||
|
|
||||||
|
@(private)
|
||||||
|
rotate__right :: proc(t: ^$T/Tree($Key, $Value), n: ^$N/Node(Key, Value)) {
|
||||||
|
l := n._left
|
||||||
|
replace_node(t, n, l)
|
||||||
|
n._left = l._right
|
||||||
|
if l._right != nil {
|
||||||
|
l._right._parent = n
|
||||||
|
}
|
||||||
|
l._right = n
|
||||||
|
n._parent = l
|
||||||
|
}
|
||||||
|
|
||||||
|
@(private)
|
||||||
|
replace_node :: proc(t: ^$T/Tree($Key, $Value), old_n: ^$N/Node(Key, Value), new_n: ^N) {
|
||||||
|
if old_n._parent == nil {
|
||||||
|
t._root = new_n
|
||||||
|
} else {
|
||||||
|
if (old_n == old_n._parent._left) {
|
||||||
|
old_n._parent._left = new_n
|
||||||
|
} else {
|
||||||
|
old_n._parent._right = new_n
|
||||||
|
}
|
||||||
|
}
|
||||||
|
if new_n != nil {
|
||||||
|
new_n._parent = old_n._parent
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
@(private)
|
||||||
|
insert_case1 :: proc(t: ^$T/Tree($Key, $Value), n: ^$N/Node(Key, Value)) {
|
||||||
|
if n._parent == nil {
|
||||||
|
n._color = .Black
|
||||||
|
} else {
|
||||||
|
insert_case2(t, n)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
@(private)
|
||||||
|
insert_case2 :: proc(t: ^$T/Tree($Key, $Value), n: ^$N/Node(Key, Value)) {
|
||||||
|
if node_color(n._parent) == .Black {
|
||||||
|
return // Tree is still valid
|
||||||
|
} else {
|
||||||
|
insert_case3(t, n)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
@(private)
|
||||||
|
insert_case3 :: proc(t: ^$T/Tree($Key, $Value), n: ^$N/Node(Key, Value)) {
|
||||||
|
if node_color(uncle(n)) == .Red {
|
||||||
|
n._parent._color = .Black
|
||||||
|
uncle(n)._color = .Black
|
||||||
|
grand_parent(n)._color = .Red
|
||||||
|
insert_case1(t, grand_parent(n))
|
||||||
|
} else {
|
||||||
|
insert_case4(t, n)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
@(private)
|
||||||
|
insert_case4 :: proc(t: ^$T/Tree($Key, $Value), n: ^$N/Node(Key, Value)) {
|
||||||
|
n := n
|
||||||
|
if n == n._parent._right && n._parent == grand_parent(n)._left {
|
||||||
|
rotate__left(t, n._parent)
|
||||||
|
n = n._left
|
||||||
|
} else if n == n._parent._left && n._parent == grand_parent(n)._right {
|
||||||
|
rotate__right(t, n._parent)
|
||||||
|
n = n._right
|
||||||
|
}
|
||||||
|
insert_case5(t, n)
|
||||||
|
}
|
||||||
|
|
||||||
|
@(private)
|
||||||
|
insert_case5 :: proc(t: ^$T/Tree($Key, $Value), n: ^$N/Node(Key, Value)) {
|
||||||
|
n._parent._color = .Black
|
||||||
|
grand_parent(n)._color = .Red
|
||||||
|
if n == n._parent._left && n._parent == grand_parent(n)._left {
|
||||||
|
rotate__right(t, grand_parent(n))
|
||||||
|
} else {
|
||||||
|
rotate__left(t, grand_parent(n))
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// The maximum_node() helper function just walks _right until it reaches the last non-leaf:
|
||||||
|
@(private)
|
||||||
|
maximum_node :: proc(n: ^$N/Node($Key, $Value)) -> (max_node: ^N) {
|
||||||
|
n := n
|
||||||
|
for n._right != nil {
|
||||||
|
n = n._right
|
||||||
|
}
|
||||||
|
return n
|
||||||
|
}
|
||||||
|
|
||||||
|
@(private)
|
||||||
|
remove_case1 :: proc(t: ^$T/Tree($Key, $Value), n: ^$N/Node(Key, Value)) {
|
||||||
|
if n._parent == nil {
|
||||||
|
return
|
||||||
|
} else {
|
||||||
|
remove_case2(t, n)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
@(private)
|
||||||
|
remove_case2 :: proc(t: ^$T/Tree($Key, $Value), n: ^$N/Node(Key, Value)) {
|
||||||
|
if node_color(sibling(n)) == .Red {
|
||||||
|
n._parent._color = .Red
|
||||||
|
sibling(n)._color = .Black
|
||||||
|
if n == n._parent._left {
|
||||||
|
rotate__left(t, n._parent)
|
||||||
|
} else {
|
||||||
|
rotate__right(t, n._parent)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
remove_case3(t, n)
|
||||||
|
}
|
||||||
|
|
||||||
|
@(private)
|
||||||
|
remove_case3 :: proc(t: ^$T/Tree($Key, $Value), n: ^$N/Node(Key, Value)) {
|
||||||
|
if node_color(n._parent) == .Black &&
|
||||||
|
node_color(sibling(n)) == .Black &&
|
||||||
|
node_color(sibling(n)._left) == .Black &&
|
||||||
|
node_color(sibling(n)._right) == .Black {
|
||||||
|
sibling(n)._color = .Red
|
||||||
|
remove_case1(t, n._parent)
|
||||||
|
} else {
|
||||||
|
remove_case4(t, n)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
@(private)
|
||||||
|
remove_case4 :: proc(t: ^$T/Tree($Key, $Value), n: ^$N/Node(Key, Value)) {
|
||||||
|
if node_color(n._parent) == .Red &&
|
||||||
|
node_color(sibling(n)) == .Black &&
|
||||||
|
node_color(sibling(n)._left) == .Black &&
|
||||||
|
node_color(sibling(n)._right) == .Black {
|
||||||
|
sibling(n)._color = .Red
|
||||||
|
n._parent._color = .Black
|
||||||
|
} else {
|
||||||
|
remove_case5(t, n)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
@(private)
|
||||||
|
remove_case5 :: proc(t: ^$T/Tree($Key, $Value), n: ^$N/Node(Key, Value)) {
|
||||||
|
if n == n._parent._left &&
|
||||||
|
node_color(sibling(n)) == .Black &&
|
||||||
|
node_color(sibling(n)._left) == .Red &&
|
||||||
|
node_color(sibling(n)._right) == .Black {
|
||||||
|
sibling(n)._color = .Red
|
||||||
|
sibling(n)._left._color = .Black
|
||||||
|
rotate__right(t, sibling(n))
|
||||||
|
} else if n == n._parent._right &&
|
||||||
|
node_color(sibling(n)) == .Black &&
|
||||||
|
node_color(sibling(n)._right) == .Red &&
|
||||||
|
node_color(sibling(n)._left) == .Black {
|
||||||
|
sibling(n)._color = .Red
|
||||||
|
sibling(n)._right._color = .Black
|
||||||
|
rotate__left(t, sibling(n))
|
||||||
|
}
|
||||||
|
remove_case6(t, n)
|
||||||
|
}
|
||||||
|
|
||||||
|
@(private)
|
||||||
|
remove_case6 :: proc(t: ^$T/Tree($Key, $Value), n: ^$N/Node(Key, Value)) {
|
||||||
|
sibling(n)._color = node_color(n._parent)
|
||||||
|
n._parent._color = .Black
|
||||||
|
if n == n._parent._left {
|
||||||
|
sibling(n)._right._color = .Black
|
||||||
|
rotate__left(t, n._parent)
|
||||||
|
} else {
|
||||||
|
sibling(n)._left._color = .Black
|
||||||
|
rotate__right(t, n._parent)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
node_color :: proc(n: ^$N/Node($Key, $Value)) -> (c: Color) {
|
||||||
|
return n == nil ? .Black : n._color
|
||||||
|
}
|
||||||
+40
-45
@@ -3,30 +3,15 @@ set COMMON=-no-bounds-check -vet -strict-style
|
|||||||
set COLLECTION=-collection:tests=..
|
set COLLECTION=-collection:tests=..
|
||||||
set PATH_TO_ODIN==..\..\odin
|
set PATH_TO_ODIN==..\..\odin
|
||||||
python3 download_assets.py
|
python3 download_assets.py
|
||||||
echo ---
|
|
||||||
echo Running core:image tests
|
|
||||||
echo ---
|
|
||||||
%PATH_TO_ODIN% run image %COMMON% -out:test_core_image.exe || exit /b
|
|
||||||
|
|
||||||
echo ---
|
echo ---
|
||||||
echo Running core:compress tests
|
echo Running core:compress tests
|
||||||
echo ---
|
echo ---
|
||||||
%PATH_TO_ODIN% run compress %COMMON% -out:test_core_compress.exe || exit /b
|
%PATH_TO_ODIN% run compress %COMMON% -out:test_core_compress.exe || exit /b
|
||||||
|
|
||||||
echo ---
|
echo ---
|
||||||
echo Running core:strings tests
|
echo Running core:container tests
|
||||||
echo ---
|
echo ---
|
||||||
%PATH_TO_ODIN% run strings %COMMON% -out:test_core_strings.exe || exit /b
|
%PATH_TO_ODIN% run container %COMMON% %COLLECTION% -out:test_core_container.exe || exit /b
|
||||||
|
|
||||||
echo ---
|
|
||||||
echo Running core:hash tests
|
|
||||||
echo ---
|
|
||||||
%PATH_TO_ODIN% run hash %COMMON% -o:size -out:test_core_hash.exe || exit /b
|
|
||||||
|
|
||||||
echo ---
|
|
||||||
echo Running core:odin tests
|
|
||||||
echo ---
|
|
||||||
%PATH_TO_ODIN% run odin %COMMON% -o:size -out:test_core_odin.exe || exit /b
|
|
||||||
|
|
||||||
echo ---
|
echo ---
|
||||||
echo Running core:crypto tests
|
echo Running core:crypto tests
|
||||||
@@ -45,9 +30,19 @@ rem %PATH_TO_ODIN% run encoding/hxa %COMMON% %COLLECTION% -out:test_hxa.exe |
|
|||||||
%PATH_TO_ODIN% run encoding/base64 %COMMON% -out:test_base64.exe || exit /b
|
%PATH_TO_ODIN% run encoding/base64 %COMMON% -out:test_base64.exe || exit /b
|
||||||
|
|
||||||
echo ---
|
echo ---
|
||||||
echo Running core:math/noise tests
|
echo Running core:fmt tests
|
||||||
echo ---
|
echo ---
|
||||||
%PATH_TO_ODIN% run math/noise %COMMON% -out:test_noise.exe || exit /b
|
%PATH_TO_ODIN% run fmt %COMMON% %COLLECTION% -out:test_core_fmt.exe || exit /b
|
||||||
|
|
||||||
|
echo ---
|
||||||
|
echo Running core:hash tests
|
||||||
|
echo ---
|
||||||
|
%PATH_TO_ODIN% run hash %COMMON% -o:size -out:test_core_hash.exe || exit /b
|
||||||
|
|
||||||
|
echo ---
|
||||||
|
echo Running core:image tests
|
||||||
|
echo ---
|
||||||
|
%PATH_TO_ODIN% run image %COMMON% -out:test_core_image.exe || exit /b
|
||||||
|
|
||||||
echo ---
|
echo ---
|
||||||
echo Running core:math tests
|
echo Running core:math tests
|
||||||
@@ -59,6 +54,21 @@ echo Running core:math/linalg/glsl tests
|
|||||||
echo ---
|
echo ---
|
||||||
%PATH_TO_ODIN% run math/linalg/glsl %COMMON% %COLLECTION% -out:test_linalg_glsl.exe || exit /b
|
%PATH_TO_ODIN% run math/linalg/glsl %COMMON% %COLLECTION% -out:test_linalg_glsl.exe || exit /b
|
||||||
|
|
||||||
|
echo ---
|
||||||
|
echo Running core:math/noise tests
|
||||||
|
echo ---
|
||||||
|
%PATH_TO_ODIN% run math/noise %COMMON% -out:test_noise.exe || exit /b
|
||||||
|
|
||||||
|
echo ---
|
||||||
|
echo Running core:net
|
||||||
|
echo ---
|
||||||
|
%PATH_TO_ODIN% run net %COMMON% -out:test_core_net.exe || exit /b
|
||||||
|
|
||||||
|
echo ---
|
||||||
|
echo Running core:odin tests
|
||||||
|
echo ---
|
||||||
|
%PATH_TO_ODIN% run odin %COMMON% -o:size -out:test_core_odin.exe || exit /b
|
||||||
|
|
||||||
echo ---
|
echo ---
|
||||||
echo Running core:path/filepath tests
|
echo Running core:path/filepath tests
|
||||||
echo ---
|
echo ---
|
||||||
@@ -69,47 +79,32 @@ echo Running core:reflect tests
|
|||||||
echo ---
|
echo ---
|
||||||
%PATH_TO_ODIN% run reflect %COMMON% %COLLECTION% -out:test_core_reflect.exe || exit /b
|
%PATH_TO_ODIN% run reflect %COMMON% %COLLECTION% -out:test_core_reflect.exe || exit /b
|
||||||
|
|
||||||
|
echo ---
|
||||||
|
echo Running core:runtime tests
|
||||||
|
echo ---
|
||||||
|
%PATH_TO_ODIN% run runtime %COMMON% %COLLECTION% -out:test_core_runtime.exe || exit /b
|
||||||
|
|
||||||
echo ---
|
echo ---
|
||||||
echo Running core:slice tests
|
echo Running core:slice tests
|
||||||
echo ---
|
echo ---
|
||||||
%PATH_TO_ODIN% run slice %COMMON% -out:test_core_slice.exe || exit /b
|
%PATH_TO_ODIN% run slice %COMMON% -out:test_core_slice.exe || exit /b
|
||||||
|
|
||||||
|
echo ---
|
||||||
|
echo Running core:strings tests
|
||||||
|
echo ---
|
||||||
|
%PATH_TO_ODIN% run strings %COMMON% -out:test_core_strings.exe || exit /b
|
||||||
|
|
||||||
echo ---
|
echo ---
|
||||||
echo Running core:text/i18n tests
|
echo Running core:text/i18n tests
|
||||||
echo ---
|
echo ---
|
||||||
%PATH_TO_ODIN% run text\i18n %COMMON% -out:test_core_i18n.exe || exit /b
|
%PATH_TO_ODIN% run text\i18n %COMMON% -out:test_core_i18n.exe || exit /b
|
||||||
|
|
||||||
echo ---
|
|
||||||
echo Running core:net
|
|
||||||
echo ---
|
|
||||||
%PATH_TO_ODIN% run net %COMMON% -out:test_core_net.exe || exit /b
|
|
||||||
|
|
||||||
echo ---
|
|
||||||
echo Running core:slice tests
|
|
||||||
echo ---
|
|
||||||
%PATH_TO_ODIN% run slice %COMMON% -out:test_core_slice.exe || exit /b
|
|
||||||
|
|
||||||
echo ---
|
|
||||||
echo Running core:container tests
|
|
||||||
echo ---
|
|
||||||
%PATH_TO_ODIN% run container %COMMON% %COLLECTION% -out:test_core_container.exe || exit /b
|
|
||||||
|
|
||||||
echo ---
|
echo ---
|
||||||
echo Running core:thread tests
|
echo Running core:thread tests
|
||||||
echo ---
|
echo ---
|
||||||
%PATH_TO_ODIN% run thread %COMMON% %COLLECTION% -out:test_core_thread.exe || exit /b
|
%PATH_TO_ODIN% run thread %COMMON% %COLLECTION% -out:test_core_thread.exe || exit /b
|
||||||
|
|
||||||
echo ---
|
|
||||||
echo Running core:runtime tests
|
|
||||||
echo ---
|
|
||||||
%PATH_TO_ODIN% run runtime %COMMON% %COLLECTION% -out:test_core_runtime.exe || exit /b
|
|
||||||
|
|
||||||
echo ---
|
echo ---
|
||||||
echo Running core:time tests
|
echo Running core:time tests
|
||||||
echo ---
|
echo ---
|
||||||
%PATH_TO_ODIN% run time %COMMON% %COLLECTION% -out:test_core_time.exe || exit /b
|
%PATH_TO_ODIN% run time %COMMON% %COLLECTION% -out:test_core_time.exe || exit /b
|
||||||
|
|
||||||
echo ---
|
|
||||||
echo Running core:fmt tests
|
|
||||||
echo ---
|
|
||||||
%PATH_TO_ODIN% run fmt %COMMON% %COLLECTION% -out:test_core_fmt.exe || exit /b
|
|
||||||
@@ -4,12 +4,12 @@ import "core:container/avl"
|
|||||||
import "core:math/rand"
|
import "core:math/rand"
|
||||||
import "core:slice"
|
import "core:slice"
|
||||||
import "core:testing"
|
import "core:testing"
|
||||||
|
import "core:fmt"
|
||||||
import tc "tests:common"
|
import tc "tests:common"
|
||||||
|
|
||||||
@(test)
|
@(test)
|
||||||
test_avl :: proc(t: ^testing.T) {
|
test_avl :: proc(t: ^testing.T) {
|
||||||
tc.log(t, "Testing avl")
|
tc.log(t, fmt.tprintf("Testing avl, using random seed %v, add -define:RANDOM_SEED=%v to reuse it.", random_seed, random_seed))
|
||||||
|
|
||||||
// Initialization.
|
// Initialization.
|
||||||
tree: avl.Tree(int)
|
tree: avl.Tree(int)
|
||||||
@@ -21,11 +21,14 @@ test_avl :: proc(t: ^testing.T) {
|
|||||||
iter := avl.iterator(&tree, avl.Direction.Forward)
|
iter := avl.iterator(&tree, avl.Direction.Forward)
|
||||||
tc.expect(t, avl.iterator_get(&iter) == nil, "empty/iterator: first node should be nil")
|
tc.expect(t, avl.iterator_get(&iter) == nil, "empty/iterator: first node should be nil")
|
||||||
|
|
||||||
|
r: rand.Rand
|
||||||
|
rand.init(&r, random_seed)
|
||||||
|
|
||||||
// Test insertion.
|
// Test insertion.
|
||||||
NR_INSERTS :: 32 + 1 // Ensure at least 1 collision.
|
NR_INSERTS :: 32 + 1 // Ensure at least 1 collision.
|
||||||
inserted_map := make(map[int]^avl.Node(int))
|
inserted_map := make(map[int]^avl.Node(int))
|
||||||
for i := 0; i < NR_INSERTS; i += 1 {
|
for i := 0; i < NR_INSERTS; i += 1 {
|
||||||
v := int(rand.uint32() & 0x1f)
|
v := int(rand.uint32(&r) & 0x1f)
|
||||||
existing_node, in_map := inserted_map[v]
|
existing_node, in_map := inserted_map[v]
|
||||||
|
|
||||||
n, ok, _ := avl.find_or_insert(&tree, v)
|
n, ok, _ := avl.find_or_insert(&tree, v)
|
||||||
@@ -38,7 +41,7 @@ test_avl :: proc(t: ^testing.T) {
|
|||||||
}
|
}
|
||||||
nrEntries := len(inserted_map)
|
nrEntries := len(inserted_map)
|
||||||
tc.expect(t, avl.len(&tree) == nrEntries, "insert: len after")
|
tc.expect(t, avl.len(&tree) == nrEntries, "insert: len after")
|
||||||
tree_validate(t, &tree)
|
validate_avl(t, &tree)
|
||||||
|
|
||||||
// Ensure that all entries can be found.
|
// Ensure that all entries can be found.
|
||||||
for k, v in inserted_map {
|
for k, v in inserted_map {
|
||||||
@@ -74,7 +77,7 @@ test_avl :: proc(t: ^testing.T) {
|
|||||||
tc.expect(t, visited == nrEntries, "iterator/backward: visited")
|
tc.expect(t, visited == nrEntries, "iterator/backward: visited")
|
||||||
|
|
||||||
// Test removal.
|
// Test removal.
|
||||||
rand.shuffle(inserted_values[:])
|
rand.shuffle(inserted_values[:], &r)
|
||||||
for v, i in inserted_values {
|
for v, i in inserted_values {
|
||||||
node := avl.find(&tree, v)
|
node := avl.find(&tree, v)
|
||||||
tc.expect(t, node != nil, "remove: find (pre)")
|
tc.expect(t, node != nil, "remove: find (pre)")
|
||||||
@@ -82,7 +85,7 @@ test_avl :: proc(t: ^testing.T) {
|
|||||||
ok := avl.remove(&tree, v)
|
ok := avl.remove(&tree, v)
|
||||||
tc.expect(t, ok, "remove: succeeds")
|
tc.expect(t, ok, "remove: succeeds")
|
||||||
tc.expect(t, nrEntries - (i + 1) == avl.len(&tree), "remove: len (post)")
|
tc.expect(t, nrEntries - (i + 1) == avl.len(&tree), "remove: len (post)")
|
||||||
tree_validate(t, &tree)
|
validate_avl(t, &tree)
|
||||||
|
|
||||||
tc.expect(t, nil == avl.find(&tree, v), "remove: find (post")
|
tc.expect(t, nil == avl.find(&tree, v), "remove: find (post")
|
||||||
}
|
}
|
||||||
@@ -114,7 +117,7 @@ test_avl :: proc(t: ^testing.T) {
|
|||||||
tc.expect(t, ok == (avl.len(&tree) > 0), "iterator/remove: next should return false")
|
tc.expect(t, ok == (avl.len(&tree) > 0), "iterator/remove: next should return false")
|
||||||
tc.expect(t, node == avl.first(&tree), "iterator/remove: next should return first")
|
tc.expect(t, node == avl.first(&tree), "iterator/remove: next should return first")
|
||||||
|
|
||||||
tree_validate(t, &tree)
|
validate_avl(t, &tree)
|
||||||
}
|
}
|
||||||
tc.expect(t, avl.len(&tree) == nrEntries - 1, "iterator/remove: len should drop by 1")
|
tc.expect(t, avl.len(&tree) == nrEntries - 1, "iterator/remove: len should drop by 1")
|
||||||
|
|
||||||
@@ -123,7 +126,7 @@ test_avl :: proc(t: ^testing.T) {
|
|||||||
}
|
}
|
||||||
|
|
||||||
@(private)
|
@(private)
|
||||||
tree_validate :: proc(t: ^testing.T, tree: ^avl.Tree($Value)) {
|
validate_avl :: proc(t: ^testing.T, tree: ^avl.Tree($Value)) {
|
||||||
tree_check_invariants(t, tree, tree._root, nil)
|
tree_check_invariants(t, tree, tree._root, nil)
|
||||||
}
|
}
|
||||||
|
|
||||||
|
|||||||
@@ -20,7 +20,7 @@ main :: proc() {
|
|||||||
t := testing.T{}
|
t := testing.T{}
|
||||||
|
|
||||||
test_avl(&t)
|
test_avl(&t)
|
||||||
|
test_rbtree(&t)
|
||||||
test_small_array(&t)
|
test_small_array(&t)
|
||||||
|
|
||||||
tc.report(&t)
|
tc.report(&t)
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -0,0 +1,244 @@
|
|||||||
|
package test_core_container
|
||||||
|
|
||||||
|
import rb "core:container/rbtree"
|
||||||
|
import "core:math/rand"
|
||||||
|
import "core:testing"
|
||||||
|
import "core:fmt"
|
||||||
|
import "base:intrinsics"
|
||||||
|
import "core:mem"
|
||||||
|
import "core:slice"
|
||||||
|
import tc "tests:common"
|
||||||
|
|
||||||
|
RANDOM_SEED :: #config(RANDOM_SEED, 0)
|
||||||
|
random_seed := u64(intrinsics.read_cycle_counter()) when RANDOM_SEED == 0 else u64(RANDOM_SEED)
|
||||||
|
|
||||||
|
test_rbtree_integer :: proc(t: ^testing.T, $Key: typeid, $Value: typeid) {
|
||||||
|
track: mem.Tracking_Allocator
|
||||||
|
mem.tracking_allocator_init(&track, context.allocator)
|
||||||
|
defer mem.tracking_allocator_destroy(&track)
|
||||||
|
context.allocator = mem.tracking_allocator(&track)
|
||||||
|
|
||||||
|
r: rand.Rand
|
||||||
|
rand.init(&r, random_seed)
|
||||||
|
|
||||||
|
tc.log(t, fmt.tprintf("Testing Red-Black Tree($Key=%v,$Value=%v), using random seed %v, add -define:RANDOM_SEED=%v to reuse it.", type_info_of(Key), type_info_of(Value), random_seed, random_seed))
|
||||||
|
tree: rb.Tree(Key, Value)
|
||||||
|
rb.init(&tree)
|
||||||
|
|
||||||
|
tc.expect(t, rb.len(&tree) == 0, "empty: len should be 0")
|
||||||
|
tc.expect(t, rb.first(&tree) == nil, "empty: first should be nil")
|
||||||
|
tc.expect(t, rb.last(&tree) == nil, "empty: last should be nil")
|
||||||
|
iter := rb.iterator(&tree, .Forward)
|
||||||
|
tc.expect(t, rb.iterator_get(&iter) == nil, "empty/iterator: first node should be nil")
|
||||||
|
|
||||||
|
// Test insertion.
|
||||||
|
NR_INSERTS :: 32 + 1 // Ensure at least 1 collision.
|
||||||
|
inserted_map := make(map[Key]^rb.Node(Key, Value))
|
||||||
|
|
||||||
|
min_key := max(Key)
|
||||||
|
max_key := min(Key)
|
||||||
|
|
||||||
|
for i := 0; i < NR_INSERTS; i += 1 {
|
||||||
|
k := Key(rand.uint32(&r)) & 0x1f
|
||||||
|
min_key = min(min_key, k); max_key = max(max_key, k)
|
||||||
|
v := Value(rand.uint32(&r))
|
||||||
|
|
||||||
|
existing_node, in_map := inserted_map[k]
|
||||||
|
n, inserted := rb.find_or_insert(&tree, k, v)
|
||||||
|
tc.expect(t, in_map != inserted, "insert: inserted should match inverse of map lookup")
|
||||||
|
if inserted {
|
||||||
|
inserted_map[k] = n
|
||||||
|
} else {
|
||||||
|
tc.expect(t, existing_node == n, "insert: expecting existing node")
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
entry_count := len(inserted_map)
|
||||||
|
tc.expect(t, rb.len(&tree) == entry_count, "insert: len after")
|
||||||
|
validate_rbtree(t, &tree)
|
||||||
|
|
||||||
|
first := rb.first(&tree)
|
||||||
|
last := rb.last(&tree)
|
||||||
|
tc.expect(t, first != nil && first.key == min_key, fmt.tprintf("insert: first should be present with key %v", min_key))
|
||||||
|
tc.expect(t, last != nil && last.key == max_key, fmt.tprintf("insert: last should be present with key %v", max_key))
|
||||||
|
|
||||||
|
// Ensure that all entries can be found.
|
||||||
|
for k, v in inserted_map {
|
||||||
|
tc.expect(t, v == rb.find(&tree, k), "Find(): Node")
|
||||||
|
tc.expect(t, k == v.key, "Find(): Node key")
|
||||||
|
}
|
||||||
|
|
||||||
|
// Test the forward/backward iterators.
|
||||||
|
inserted_keys: [dynamic]Key
|
||||||
|
for k in inserted_map {
|
||||||
|
append(&inserted_keys, k)
|
||||||
|
}
|
||||||
|
slice.sort(inserted_keys[:])
|
||||||
|
|
||||||
|
iter = rb.iterator(&tree, rb.Direction.Forward)
|
||||||
|
visited: int
|
||||||
|
for node in rb.iterator_next(&iter) {
|
||||||
|
k, idx := node.key, visited
|
||||||
|
tc.expect(t, inserted_keys[idx] == k, "iterator/forward: key")
|
||||||
|
tc.expect(t, node == rb.iterator_get(&iter), "iterator/forward: get")
|
||||||
|
visited += 1
|
||||||
|
}
|
||||||
|
tc.expect(t, visited == entry_count, "iterator/forward: visited")
|
||||||
|
|
||||||
|
slice.reverse(inserted_keys[:])
|
||||||
|
iter = rb.iterator(&tree, rb.Direction.Backward)
|
||||||
|
visited = 0
|
||||||
|
for node in rb.iterator_next(&iter) {
|
||||||
|
k, idx := node.key, visited
|
||||||
|
tc.expect(t, inserted_keys[idx] == k, "iterator/backward: key")
|
||||||
|
visited += 1
|
||||||
|
}
|
||||||
|
tc.expect(t, visited == entry_count, "iterator/backward: visited")
|
||||||
|
|
||||||
|
// Test removal (and on_remove callback)
|
||||||
|
rand.shuffle(inserted_keys[:], &r)
|
||||||
|
callback_count := entry_count
|
||||||
|
tree.user_data = &callback_count
|
||||||
|
tree.on_remove = proc(key: Key, value: Value, user_data: rawptr) {
|
||||||
|
(^int)(user_data)^ -= 1
|
||||||
|
}
|
||||||
|
for k, i in inserted_keys {
|
||||||
|
node := rb.find(&tree, k)
|
||||||
|
tc.expect(t, node != nil, "remove: find (pre)")
|
||||||
|
|
||||||
|
ok := rb.remove(&tree, k)
|
||||||
|
tc.expect(t, ok, "remove: succeeds")
|
||||||
|
tc.expect(t, entry_count - (i + 1) == rb.len(&tree), "remove: len (post)")
|
||||||
|
validate_rbtree(t, &tree)
|
||||||
|
|
||||||
|
tc.expect(t, nil == rb.find(&tree, k), "remove: find (post")
|
||||||
|
}
|
||||||
|
tc.expect(t, rb.len(&tree) == 0, "remove: len should be 0")
|
||||||
|
tc.expect(t, callback_count == 0, fmt.tprintf("remove: on_remove should've been called %v times, it was %v", entry_count, callback_count))
|
||||||
|
tc.expect(t, rb.first(&tree) == nil, "remove: first should be nil")
|
||||||
|
tc.expect(t, rb.last(&tree) == nil, "remove: last should be nil")
|
||||||
|
|
||||||
|
// Refill the tree.
|
||||||
|
for k in inserted_keys {
|
||||||
|
rb.find_or_insert(&tree, k, 42)
|
||||||
|
}
|
||||||
|
|
||||||
|
// Test that removing the node doesn't break the iterator.
|
||||||
|
callback_count = entry_count
|
||||||
|
iter = rb.iterator(&tree, rb.Direction.Forward)
|
||||||
|
if node := rb.iterator_get(&iter); node != nil {
|
||||||
|
k := node.key
|
||||||
|
|
||||||
|
ok := rb.iterator_remove(&iter)
|
||||||
|
tc.expect(t, ok, "iterator/remove: success")
|
||||||
|
|
||||||
|
ok = rb.iterator_remove(&iter)
|
||||||
|
tc.expect(t, !ok, "iterator/remove: redundant removes should fail")
|
||||||
|
|
||||||
|
tc.expect(t, rb.find(&tree, k) == nil, "iterator/remove: node should be gone")
|
||||||
|
tc.expect(t, rb.iterator_get(&iter) == nil, "iterator/remove: get should return nil")
|
||||||
|
|
||||||
|
// Ensure that iterator_next still works.
|
||||||
|
node, ok = rb.iterator_next(&iter)
|
||||||
|
tc.expect(t, ok == (rb.len(&tree) > 0), "iterator/remove: next should return false")
|
||||||
|
tc.expect(t, node == rb.first(&tree), "iterator/remove: next should return first")
|
||||||
|
|
||||||
|
validate_rbtree(t, &tree)
|
||||||
|
}
|
||||||
|
tc.expect(t, rb.len(&tree) == entry_count - 1, "iterator/remove: len should drop by 1")
|
||||||
|
|
||||||
|
rb.destroy(&tree)
|
||||||
|
tc.expect(t, rb.len(&tree) == 0, "destroy: len should be 0")
|
||||||
|
tc.expect(t, callback_count == 0, fmt.tprintf("remove: on_remove should've been called %v times, it was %v", entry_count, callback_count))
|
||||||
|
|
||||||
|
// print_tree_node(tree._root)
|
||||||
|
delete(inserted_map)
|
||||||
|
delete(inserted_keys)
|
||||||
|
tc.expect(t, len(track.allocation_map) == 0, fmt.tprintf("Expected 0 leaks, have %v", len(track.allocation_map)))
|
||||||
|
tc.expect(t, len(track.bad_free_array) == 0, fmt.tprintf("Expected 0 bad frees, have %v", len(track.bad_free_array)))
|
||||||
|
return
|
||||||
|
}
|
||||||
|
|
||||||
|
@(test)
|
||||||
|
test_rbtree :: proc(t: ^testing.T) {
|
||||||
|
test_rbtree_integer(t, u16, u16)
|
||||||
|
}
|
||||||
|
|
||||||
|
print_tree_node :: proc(n: ^$N/rb.Node($Key, $Value), indent := 0) {
|
||||||
|
if n == nil {
|
||||||
|
fmt.println("<empty tree>")
|
||||||
|
return
|
||||||
|
}
|
||||||
|
if n.right != nil {
|
||||||
|
print_tree_node(n.right, indent + 1)
|
||||||
|
}
|
||||||
|
for _ in 0..<indent {
|
||||||
|
fmt.printf("\t")
|
||||||
|
}
|
||||||
|
if n.color == .Black {
|
||||||
|
fmt.printfln("%v", n.key)
|
||||||
|
} else {
|
||||||
|
fmt.printfln("<%v>", n.key)
|
||||||
|
}
|
||||||
|
if n.left != nil {
|
||||||
|
print_tree_node(n.left, indent + 1)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
validate_rbtree :: proc(t: ^testing.T, tree: ^$T/rb.Tree($Key, $Value)) {
|
||||||
|
verify_rbtree_propery_1(t, tree._root)
|
||||||
|
verify_rbtree_propery_2(t, tree._root)
|
||||||
|
/* Property 3 is implicit */
|
||||||
|
verify_rbtree_propery_4(t, tree._root)
|
||||||
|
verify_rbtree_propery_5(t, tree._root)
|
||||||
|
}
|
||||||
|
|
||||||
|
verify_rbtree_propery_1 :: proc(t: ^testing.T, n: ^$N/rb.Node($Key, $Value)) {
|
||||||
|
tc.expect(t, rb.node_color(n) == .Black || rb.node_color(n) == .Red, "Property #1: Each node is either red or black.")
|
||||||
|
if n == nil {
|
||||||
|
return
|
||||||
|
}
|
||||||
|
verify_rbtree_propery_1(t, n._left)
|
||||||
|
verify_rbtree_propery_1(t, n._right)
|
||||||
|
}
|
||||||
|
|
||||||
|
verify_rbtree_propery_2 :: proc(t: ^testing.T, root: ^$N/rb.Node($Key, $Value)) {
|
||||||
|
tc.expect(t, rb.node_color(root) == .Black, "Property #2: Root node should be black.")
|
||||||
|
}
|
||||||
|
|
||||||
|
verify_rbtree_propery_4 :: proc(t: ^testing.T, n: ^$N/rb.Node($Key, $Value)) {
|
||||||
|
if rb.node_color(n) == .Red {
|
||||||
|
// A red node's left, right and parent should be black
|
||||||
|
all_black := rb.node_color(n._left) == .Black && rb.node_color(n._right) == .Black && rb.node_color(n._parent) == .Black
|
||||||
|
tc.expect(t, all_black, "Property #3: Red node's children + parent must be black.")
|
||||||
|
}
|
||||||
|
if n == nil {
|
||||||
|
return
|
||||||
|
}
|
||||||
|
verify_rbtree_propery_4(t, n._left)
|
||||||
|
verify_rbtree_propery_4(t, n._right)
|
||||||
|
}
|
||||||
|
|
||||||
|
verify_rbtree_propery_5 :: proc(t: ^testing.T, root: ^$N/rb.Node($Key, $Value)) {
|
||||||
|
black_count_path := -1
|
||||||
|
verify_rbtree_propery_5_helper(t, root, 0, &black_count_path)
|
||||||
|
}
|
||||||
|
verify_rbtree_propery_5_helper :: proc(t: ^testing.T, n: ^$N/rb.Node($Key, $Value), black_count: int, path_black_count: ^int) {
|
||||||
|
black_count := black_count
|
||||||
|
|
||||||
|
if rb.node_color(n) == .Black {
|
||||||
|
black_count += 1
|
||||||
|
}
|
||||||
|
if n == nil {
|
||||||
|
if path_black_count^ == -1 {
|
||||||
|
path_black_count^ = black_count
|
||||||
|
} else {
|
||||||
|
tc.expect(t, black_count == path_black_count^, "Property #5: Paths from a node to its leaves contain same black count.")
|
||||||
|
}
|
||||||
|
return
|
||||||
|
}
|
||||||
|
verify_rbtree_propery_5_helper(t, n._left, black_count, path_black_count)
|
||||||
|
verify_rbtree_propery_5_helper(t, n._right, black_count, path_black_count)
|
||||||
|
}
|
||||||
|
// Properties 4 and 5 together guarantee that no path in the tree is more than about twice as long as any other path,
|
||||||
|
// which guarantees that it has O(log n) height.
|
||||||
Reference in New Issue
Block a user