1#include <console/panic.h>
2#include <mem/alloc.h>
3#include <stdbool.h>
4#include <stdint.h>
5#include <structures/radix.h>
6
7static bool radix_verify_tree(struct radix_tree *tree);
8static void radix_prune_up(struct radix_node *node, struct radix_tree *tree);
9
10static inline uint64_t radix_index(uint64_t key, uint32_t level) {
11 uint32_t shift = level * RADIX_BITS;
12 return (key >> shift) & RADIX_MASK;
13}
14
15int32_t radix_insert(struct radix_tree *tree, void *item) {
16 uint64_t key = tree->key_fn(item);
17 int32_t level = tree->height;
18
19 if (!tree->root) {
20 tree->root = kmalloc(sizeof(struct radix_node), ALLOC_FLAGS_ZERO);
21 if (!tree->root)
22 return ERR_NO_MEM;
23 }
24
25 struct radix_node *node = tree->root;
26
27 for (; level > 1; level--) {
28 uint64_t idx = radix_index(key, level: level - 1);
29
30 if (!node->slots[idx]) {
31 struct radix_node *mid =
32 kmalloc(sizeof(struct radix_node), ALLOC_FLAGS_ZERO);
33 if (!mid) {
34 radix_prune_up(node, tree);
35 return ERR_NO_MEM;
36 }
37 mid->parent = node;
38 node->slots[idx] = mid;
39 node->present_mask |= (1ULL << idx);
40 }
41
42 node = node->slots[idx];
43 }
44
45 uint64_t idx = radix_index(key, level: 0);
46 if (node->slots[idx]) {
47 radix_verify_tree(tree);
48 return ERR_EXIST;
49 }
50
51 node->slots[idx] = item;
52 node->present_mask |= (1ULL << idx);
53 radix_verify_tree(tree);
54
55 return 0;
56}
57
58void *radix_lookup(struct radix_tree *tree, uint64_t key) {
59 struct radix_node *node = tree->root;
60 for (int32_t level = tree->height; level > 1; level--) {
61 if (!node)
62 return NULL;
63 uint64_t idx = radix_index(key, level: level - 1);
64 node = node->slots[idx];
65 }
66 if (!node)
67 return NULL;
68 return node->slots[radix_index(key, level: 0)];
69}
70
71static bool radix_verify_node(struct radix_tree *tree, struct radix_node *node,
72 struct radix_node *expected_parent, int level,
73 int max_height, uint64_t prefix,
74 int *node_count) {
75 if (!node)
76 return true;
77
78 if (node->parent != expected_parent)
79 panic("Node %p has incorrect parent %p (expected %p)", (void *) node,
80 (void *) node->parent, (void *) expected_parent);
81
82 if (node_count)
83 (*node_count)++;
84
85 uint64_t expected_mask = 0;
86 for (int i = 0; i < RADIX_SIZE; i++) {
87 if (node->slots[i])
88 expected_mask |= (1ULL << i);
89 }
90
91 if (expected_mask != node->present_mask)
92 panic("Node %p present_mask mismatch: expected 0x%llx, got 0x%llx",
93 (void *) node, (unsigned long long) expected_mask,
94 (unsigned long long) node->present_mask);
95
96 bool is_leaf_parent = (level + 1 == max_height);
97
98 for (int i = 0; i < RADIX_SIZE; i++) {
99 void *child = node->slots[i];
100 if (!child)
101 continue;
102
103 if (level >= max_height)
104 panic("Node %p at level %d has child beyond max height %d",
105 (void *) node, level, max_height);
106
107 uint32_t shift = (max_height - level - 1) * RADIX_BITS;
108 uint64_t child_prefix = prefix | ((uint64_t) i << shift);
109
110 if (is_leaf_parent) {
111 uint64_t item_key = tree->key_fn(child);
112 if (item_key != child_prefix)
113 panic("Leaf item at slot %d has key %llu (expected %llu)", i,
114 (unsigned long long) item_key,
115 (unsigned long long) child_prefix);
116
117 if (node_count)
118 (*node_count)++;
119 } else if (!radix_verify_node(tree, node: child, expected_parent: node, level: level + 1, max_height,
120 prefix: child_prefix, node_count)) {
121 return false;
122 }
123 }
124
125 return true;
126}
127
128static bool radix_verify_tree(struct radix_tree *tree) {
129 if (!tree || !tree->root)
130 return true;
131
132 int node_count = 0;
133 return radix_verify_node(tree, node: tree->root, NULL, level: 0, max_height: tree->height, prefix: 0,
134 node_count: &node_count);
135}
136
137static void radix_prune_up(struct radix_node *node, struct radix_tree *tree) {
138 while (node && node->present_mask == 0) {
139 struct radix_node *parent = node->parent;
140 if (!parent) {
141 if (tree->root == node) {
142 kfree(node);
143 tree->root = NULL;
144 }
145 break;
146 }
147
148 for (uint64_t i = 0; i < RADIX_SIZE; i++) {
149 if (parent->slots[i] == node) {
150 parent->slots[i] = NULL;
151 parent->present_mask &= ~(1ULL << i);
152 break;
153 }
154 }
155
156 kfree(node);
157 node = parent;
158 }
159}
160
161void *radix_delete(struct radix_tree *tree, uint64_t key) {
162 struct radix_node *node = tree->root;
163 uint64_t idx = 0;
164
165 for (int32_t level = tree->height; level > 1; level--) {
166 if (!node)
167 return NULL;
168 idx = radix_index(key, level: level - 1);
169 node = node->slots[idx];
170 }
171
172 if (!node)
173 return NULL;
174
175 idx = radix_index(key, level: 0);
176 void *item = node->slots[idx];
177 if (!item)
178 return NULL;
179
180 node->slots[idx] = NULL;
181 node->present_mask &= ~(1ULL << idx);
182
183 radix_prune_up(node, tree);
184 radix_verify_tree(tree);
185
186 return item;
187}
188
189struct radix_tree *radix_tree_init(struct radix_tree *r, radix_key_fn kfn,
190 int height) {
191 r->root = NULL;
192 r->key_fn = kfn;
193 r->height = height;
194 return r;
195}
196