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)
130 return true;
131
132 if (!tree->root)
133 if (tree->height != 0)
134 panic("Tree has no root but nonzero height %u", tree->height);
135
136 int node_count = 0;
137 return radix_verify_node(tree, node: tree->root, NULL, level: 0, max_height: tree->height, prefix: 0,
138 node_count: &node_count);
139}
140
141static void radix_prune_up(struct radix_node *node, struct radix_tree *tree) {
142 while (node && node->present_mask == 0) {
143 struct radix_node *parent = node->parent;
144 if (!parent) {
145 if (tree->root == node) {
146 kfree(node);
147 tree->root = NULL;
148 }
149 break;
150 }
151
152 for (uint64_t i = 0; i < RADIX_SIZE; i++) {
153 if (parent->slots[i] == node) {
154 parent->slots[i] = NULL;
155 parent->present_mask &= ~(1ULL << i);
156 break;
157 }
158 }
159
160 kfree(node);
161 node = parent;
162 }
163}
164
165void *radix_delete(struct radix_tree *tree, uint64_t key) {
166 struct radix_node *node = tree->root;
167 uint64_t idx = 0;
168
169 for (int32_t level = tree->height; level > 1; level--) {
170 if (!node)
171 return NULL;
172 idx = radix_index(key, level: level - 1);
173 node = node->slots[idx];
174 }
175
176 if (!node)
177 return NULL;
178
179 idx = radix_index(key, level: 0);
180 void *item = node->slots[idx];
181 if (!item)
182 return NULL;
183
184 node->slots[idx] = NULL;
185 node->present_mask &= ~(1ULL << idx);
186
187 radix_prune_up(node, tree);
188 radix_verify_tree(tree);
189
190 return item;
191}
192
193struct radix_tree *radix_tree_init(struct radix_tree *r, radix_key_fn kfn,
194 int height) {
195 r->root = NULL;
196 r->key_fn = kfn;
197 r->height = height;
198 return r;
199}
200