diff options
| author | Jeroen van Rijn <Kelimion@users.noreply.github.com> | 2024-05-24 14:04:56 +0200 |
|---|---|---|
| committer | GitHub <noreply@github.com> | 2024-05-24 14:04:56 +0200 |
| commit | 479d301e92c133d42c6ff643a148fab9283c3772 (patch) | |
| tree | 61d27b814a36a575f7113beae713265a44230dd7 | |
| parent | 0658778a307fdfbfa522040d793df6dbb1da0211 (diff) | |
| parent | 11e57fd3fd290c1237f33e83f0142b3d98a7e2d9 (diff) | |
Merge pull request #3614 from Kelimion/rbtree
Add `core:container/rbtree`
| -rw-r--r-- | core/container/avl/avl.odin | 9 | ||||
| -rw-r--r-- | core/container/rbtree/rbtree.odin | 568 | ||||
| -rw-r--r-- | tests/core/build.bat | 77 | ||||
| -rw-r--r-- | tests/core/container/test_core_avl.odin | 19 | ||||
| -rw-r--r-- | tests/core/container/test_core_container.odin | 2 | ||||
| -rw-r--r-- | tests/core/container/test_core_rbtree.odin | 244 |
6 files changed, 863 insertions, 56 deletions
diff --git a/core/container/avl/avl.odin b/core/container/avl/avl.odin index eecc1b756..582cd87fd 100644 --- a/core/container/avl/avl.odin +++ b/core/container/avl/avl.odin @@ -5,13 +5,10 @@ The implementation is non-intrusive, and non-recursive. */ package container_avl -import "base:intrinsics" -import "base:runtime" +@(require) import "base:intrinsics" +@(require) import "base:runtime" import "core:slice" -_ :: intrinsics -_ :: runtime - // Originally based on the CC0 implementation by Eric Biggers // See: https://github.com/ebiggers/avl_tree/ @@ -675,4 +672,4 @@ iterator_first :: proc "contextless" (it: ^Iterator($Value)) { if it._cur != nil { it._next = node_next_or_prev_in_order(it._cur, it._direction) } -} +}
\ No newline at end of file diff --git a/core/container/rbtree/rbtree.odin b/core/container/rbtree/rbtree.odin new file mode 100644 index 000000000..8ab131b3b --- /dev/null +++ b/core/container/rbtree/rbtree.odin @@ -0,0 +1,568 @@ +// This package implements a red-black tree +package container_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, err: runtime.Allocator_Error) { + 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, nil +} + +// 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 = cast(^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 +}
\ No newline at end of file diff --git a/tests/core/build.bat b/tests/core/build.bat index 5dd2d1932..7871e52e2 100644 --- a/tests/core/build.bat +++ b/tests/core/build.bat @@ -4,29 +4,14 @@ set COLLECTION=-collection:tests=.. set PATH_TO_ODIN==..\..\odin 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 Running core:compress tests echo --- %PATH_TO_ODIN% run compress %COMMON% -out:test_core_compress.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 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 Running core:container tests echo --- -%PATH_TO_ODIN% run odin %COMMON% -o:size -out:test_core_odin.exe || exit /b +%PATH_TO_ODIN% run container %COMMON% %COLLECTION% -out:test_core_container.exe || exit /b echo --- 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 echo --- -echo Running core:math/noise tests +echo Running core:fmt tests 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 Running core:math tests @@ -60,56 +55,56 @@ echo --- %PATH_TO_ODIN% run math/linalg/glsl %COMMON% %COLLECTION% -out:test_linalg_glsl.exe || exit /b echo --- -echo Running core:path/filepath tests +echo Running core:math/noise tests echo --- -%PATH_TO_ODIN% run path/filepath %COMMON% %COLLECTION% -out:test_core_filepath.exe || exit /b +%PATH_TO_ODIN% run math/noise %COMMON% -out:test_noise.exe || exit /b echo --- -echo Running core:reflect tests +echo Running core:net echo --- -%PATH_TO_ODIN% run reflect %COMMON% %COLLECTION% -out:test_core_reflect.exe || exit /b +%PATH_TO_ODIN% run net %COMMON% -out:test_core_net.exe || exit /b echo --- -echo Running core:slice tests +echo Running core:odin tests echo --- -%PATH_TO_ODIN% run slice %COMMON% -out:test_core_slice.exe || exit /b +%PATH_TO_ODIN% run odin %COMMON% -o:size -out:test_core_odin.exe || exit /b echo --- -echo Running core:text/i18n tests +echo Running core:path/filepath tests echo --- -%PATH_TO_ODIN% run text\i18n %COMMON% -out:test_core_i18n.exe || exit /b +%PATH_TO_ODIN% run path/filepath %COMMON% %COLLECTION% -out:test_core_filepath.exe || exit /b echo --- -echo Running core:net +echo Running core:reflect tests echo --- -%PATH_TO_ODIN% run net %COMMON% -out:test_core_net.exe || exit /b +%PATH_TO_ODIN% run reflect %COMMON% %COLLECTION% -out:test_core_reflect.exe || exit /b echo --- -echo Running core:slice tests +echo Running core:runtime tests echo --- -%PATH_TO_ODIN% run slice %COMMON% -out:test_core_slice.exe || exit /b +%PATH_TO_ODIN% run runtime %COMMON% %COLLECTION% -out:test_core_runtime.exe || exit /b echo --- -echo Running core:container tests +echo Running core:slice tests echo --- -%PATH_TO_ODIN% run container %COMMON% %COLLECTION% -out:test_core_container.exe || exit /b +%PATH_TO_ODIN% run slice %COMMON% -out:test_core_slice.exe || exit /b echo --- -echo Running core:thread tests +echo Running core:strings tests echo --- -%PATH_TO_ODIN% run thread %COMMON% %COLLECTION% -out:test_core_thread.exe || exit /b +%PATH_TO_ODIN% run strings %COMMON% -out:test_core_strings.exe || exit /b echo --- -echo Running core:runtime tests +echo Running core:text/i18n tests echo --- -%PATH_TO_ODIN% run runtime %COMMON% %COLLECTION% -out:test_core_runtime.exe || exit /b +%PATH_TO_ODIN% run text\i18n %COMMON% -out:test_core_i18n.exe || exit /b echo --- -echo Running core:time tests +echo Running core:thread tests echo --- -%PATH_TO_ODIN% run time %COMMON% %COLLECTION% -out:test_core_time.exe || exit /b +%PATH_TO_ODIN% run thread %COMMON% %COLLECTION% -out:test_core_thread.exe || exit /b echo --- -echo Running core:fmt tests +echo Running core:time tests echo --- -%PATH_TO_ODIN% run fmt %COMMON% %COLLECTION% -out:test_core_fmt.exe || exit /b
\ No newline at end of file +%PATH_TO_ODIN% run time %COMMON% %COLLECTION% -out:test_core_time.exe || exit /b
\ No newline at end of file diff --git a/tests/core/container/test_core_avl.odin b/tests/core/container/test_core_avl.odin index f6343c5ea..2244ab7f6 100644 --- a/tests/core/container/test_core_avl.odin +++ b/tests/core/container/test_core_avl.odin @@ -4,12 +4,12 @@ import "core:container/avl" import "core:math/rand" import "core:slice" import "core:testing" - +import "core:fmt" import tc "tests:common" @(test) 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. tree: avl.Tree(int) @@ -21,11 +21,14 @@ test_avl :: proc(t: ^testing.T) { iter := avl.iterator(&tree, avl.Direction.Forward) 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. NR_INSERTS :: 32 + 1 // Ensure at least 1 collision. inserted_map := make(map[int]^avl.Node(int)) 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] n, ok, _ := avl.find_or_insert(&tree, v) @@ -38,7 +41,7 @@ test_avl :: proc(t: ^testing.T) { } nrEntries := len(inserted_map) 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. for k, v in inserted_map { @@ -74,7 +77,7 @@ test_avl :: proc(t: ^testing.T) { tc.expect(t, visited == nrEntries, "iterator/backward: visited") // Test removal. - rand.shuffle(inserted_values[:]) + rand.shuffle(inserted_values[:], &r) for v, i in inserted_values { node := avl.find(&tree, v) tc.expect(t, node != nil, "remove: find (pre)") @@ -82,7 +85,7 @@ test_avl :: proc(t: ^testing.T) { ok := avl.remove(&tree, v) tc.expect(t, ok, "remove: succeeds") 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") } @@ -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, 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") @@ -123,7 +126,7 @@ test_avl :: proc(t: ^testing.T) { } @(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) } diff --git a/tests/core/container/test_core_container.odin b/tests/core/container/test_core_container.odin index f816a6bcb..7dd4a3628 100644 --- a/tests/core/container/test_core_container.odin +++ b/tests/core/container/test_core_container.odin @@ -20,7 +20,7 @@ main :: proc() { t := testing.T{} test_avl(&t) + test_rbtree(&t) test_small_array(&t) - tc.report(&t) } diff --git a/tests/core/container/test_core_rbtree.odin b/tests/core/container/test_core_rbtree.odin new file mode 100644 index 000000000..89742b1d0 --- /dev/null +++ b/tests/core/container/test_core_rbtree.odin @@ -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.
\ No newline at end of file |