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:
Jeroen van Rijn
2024-05-23 23:00:00 +02:00
parent 7dc1f114b9
commit 410876b36a
6 changed files with 866 additions and 60 deletions
+40 -45
View File
@@ -3,30 +3,15 @@ set COMMON=-no-bounds-check -vet -strict-style
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 Running core:container 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 ---
%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
@@ -59,6 +54,21 @@ echo Running core:math/linalg/glsl tests
echo ---
%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 Running core:path/filepath tests
echo ---
@@ -69,47 +79,32 @@ echo Running core:reflect tests
echo ---
%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 Running core:slice tests
echo ---
%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 Running core:text/i18n tests
echo ---
%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 Running core:thread tests
echo ---
%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 Running core:time tests
echo ---
%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
%PATH_TO_ODIN% run time %COMMON% %COLLECTION% -out:test_core_time.exe || exit /b
+11 -8
View File
@@ -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)
}
@@ -20,7 +20,7 @@ main :: proc() {
t := testing.T{}
test_avl(&t)
test_rbtree(&t)
test_small_array(&t)
tc.report(&t)
}
+244
View File
@@ -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.