1#include <console/panic.h>
2#include <math/bit.h>
3#include <mem/alloc.h>
4#include <stdbool.h>
5#include <stdint.h>
6#include <structures/radix.h>
7
8static bool radix_verify_tree(struct radix_tree *tree);
9static void radix_prune_up(struct radix_node *node, struct radix_tree *tree);
10
11static inline uint64_t radix_index(uint64_t key, uint32_t level) {
12 uint32_t shift = level * RADIX_BITS;
13 return (key >> shift) & RADIX_MASK;
14}
15
16int32_t radix_insert(struct radix_tree *tree, void *item) {
17 uint64_t key = tree->key_fn(item);
18 int32_t level = tree->height;
19
20 if (!tree->root) {
21 tree->root = vnew(tree->root, ALLOC_ZERO);
22 if (!tree->root)
23 return ERR_NO_MEM;
24 }
25
26 struct radix_node *node = tree->root;
27
28 for (; level > 1; level--) {
29 uint64_t idx = radix_index(key, level: level - 1);
30
31 if (!node->slots[idx]) {
32 struct radix_node *mid = vnew(mid, ALLOC_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 |= BIT(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 |= BIT(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(CORRUPTION, "Node %p has incorrect parent %p (expected %p)",
80 (void *) node, (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 |= BIT(i);
89 }
90
91 if (expected_mask != node->present_mask)
92 panic(CORRUPTION,
93 "Node %p present_mask mismatch: expected 0x%llx, got 0x%llx",
94 (void *) node, (unsigned long long) expected_mask,
95 (unsigned long long) node->present_mask);
96
97 bool is_leaf_parent = (level + 1 == max_height);
98
99 for (int i = 0; i < RADIX_SIZE; i++) {
100 void *child = node->slots[i];
101 if (!child)
102 continue;
103
104 if (level >= max_height)
105 panic(CORRUPTION,
106 "Node %p at level %d has child beyond max height %d",
107 (void *) node, level, max_height);
108
109 uint32_t shift = (max_height - level - 1) * RADIX_BITS;
110 uint64_t child_prefix = prefix | ((uint64_t) i << shift);
111
112 if (is_leaf_parent) {
113 uint64_t item_key = tree->key_fn(child);
114 if (item_key != child_prefix)
115 panic(CORRUPTION,
116 "Leaf item at slot %d has key %llu (expected %llu)", i,
117 (unsigned long long) item_key,
118 (unsigned long long) child_prefix);
119
120 if (node_count)
121 (*node_count)++;
122 } else if (!radix_verify_node(tree, node: child, expected_parent: node, level: level + 1, max_height,
123 prefix: child_prefix, node_count)) {
124 return false;
125 }
126 }
127
128 return true;
129}
130
131static bool radix_verify_tree(struct radix_tree *tree) {
132 if (!tree || !tree->root)
133 return true;
134
135 int node_count = 0;
136 return radix_verify_node(tree, node: tree->root, NULL, level: 0, max_height: tree->height, prefix: 0,
137 node_count: &node_count);
138}
139
140static void radix_prune_up(struct radix_node *node, struct radix_tree *tree) {
141 enforce(!tree->root || !tree->root->parent, CORRUPTION);
142 while (node && node->present_mask == 0) {
143 struct radix_node *parent = node->parent;
144 if (!parent) {
145 if (tree->root == node) {
146 vfree(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 &= ~BIT(i);
156 break;
157 }
158 }
159
160 vfree(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 &= ~BIT(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