==> src/omap.c <==1/* The AVL tree behind @ajax/mesh-omap: recursive insert and delete, navigation by key. */
2#include "mesh/omap.h"
3
4#include <stdio.h>
5#include <stdlib.h>
6#include <string.h>
7
8typedef struct node {
9 struct node *l, *r;
10 int h;
11 uint64_t data[]; /* the key's words, then the value */
12} node;
13
14struct momap {
15 node *root;
16 size_t len;
17 unsigned kw;
18 size_t vsize;
19 momap_free_fn free_val;
20 momap_copy_fn copy_val;
21};
22
23static void *xalloc(size_t n) {
24 void *p = calloc(1, n ? n : 1);
25 if (!p) {
26 fputs("mesh-omap: out of memory\n", stderr);
27 abort();
28 }
29 return p;
30}
31
32static void *val_of(const momap *m, const node *n) { return (void *)(n->data + m->kw); }
33
34static inline int cmp(const momap *m, const uint64_t *a, const uint64_t *b) {
35 if (a[0] != b[0]) return a[0] < b[0] ? -1 : 1;
36 for (unsigned i = 1; i < m->kw; i++)
37 if (a[i] != b[i]) return a[i] < b[i] ? -1 : 1;
38 return 0;
39}
40
41static int height(const node *n) { return n ? n->h : 0; }
42static void fix(node *n) {
43 int a = height(n->l), b = height(n->r);
44 n->h = (a > b ? a : b) + 1;
45}
46static node *rot_right(node *n) {
47 node *l = n->l;
48 n->l = l->r;
49 l->r = n;
50 fix(n);
51 fix(l);
52 return l;
53}
54static node *rot_left(node *n) {
55 node *r = n->r;
56 n->r = r->l;
57 r->l = n;
58 fix(n);
59 fix(r);
60 return r;
61}
62static node *balance(node *n) {
63 fix(n);
64 int b = height(n->l) - height(n->r);
65 if (b > 1) {
66 if (height(n->l->l) < height(n->l->r)) n->l = rot_left(n->l);
67 return rot_right(n);
68 }
69 if (b < -1) {
70 if (height(n->r->r) < height(n->r->l)) n->r = rot_right(n->r);
71 return rot_left(n);
72 }
73 return n;
74}
75
76momap *momap_new(unsigned key_words, size_t val_size, momap_free_fn free_val, momap_copy_fn copy_val) {
77 if (key_words < 1 || key_words > MOMAP_MAX_KEY) {
78 fprintf(stderr, "mesh-omap: %u key words (1..%d)\n", key_words, MOMAP_MAX_KEY);
79 abort();
80 }
81 momap *m = xalloc(sizeof *m);
82 m->kw = key_words;
83 m->vsize = val_size;
84 m->free_val = free_val;
85 m->copy_val = copy_val;
86 return m;
87}
88
89static void drop_tree(momap *m, node *n) {
90 while (n) {
91 drop_tree(m, n->l);
92 node *r = n->r;
93 if (m->free_val) m->free_val(val_of(m, n));
94 free(n);
95 n = r;
96 }
97}
98
99void momap_clear(momap *m) {
100 if (!m) return;
101 drop_tree(m, m->root);
102 m->root = NULL;
103 m->len = 0;
104}
105
106void momap_free(momap *m) {
107 if (!m) return;
108 momap_clear(m);
109 free(m);
110}
111
112static node *new_node(const momap *m, const uint64_t *key) {
113 node *n = xalloc(sizeof(node) + 8 * m->kw + (m->vsize + 7) / 8 * 8);
114 memcpy(n->data, key, 8 * m->kw);
115 n->h = 1;
116 return n;
117}
118
119static node *clone_tree(const momap *m, const node *n) {
120 if (!n) return NULL;
121 node *c = new_node(m, n->data);
122 c->h = n->h;
123 if (m->copy_val) m->copy_val(val_of(m, c), val_of(m, n));
124 else memcpy(val_of(m, c), val_of(m, n), m->vsize);
125 c->l = clone_tree(m, n->l);
126 c->r = clone_tree(m, n->r);
127 return c;
128}
129
130momap *momap_clone(const momap *m) {
131 momap *c = momap_new(m->kw, m->vsize, m->free_val, m->copy_val);
132 c->root = clone_tree(m, m->root);
133 c->len = m->len;
134 return c;
135}
136
137size_t momap_len(const momap *m) { return m ? m->len : 0; }
138unsigned momap_key_words(const momap *m) { return m->kw; }
139size_t momap_val_size(const momap *m) { return m->vsize; }
140
141void *momap_get(const momap *m, const uint64_t *key) {
142 const node *n = m->root;
143 if (m->kw == 1) {
144 uint64_t k = key[0];
145 while (n) {
146 uint64_t x = n->data[0];
147 if (k == x) return val_of(m, n);
148 n = k < x ? n->l : n->r;
149 }
150 return NULL;
151 }
152 while (n) {
153 int c = cmp(m, key, n->data);
154 if (c == 0) return val_of(m, n);
155 n = c < 0 ? n->l : n->r;
156 }
157 return NULL;
158}
159
160static node *insert(momap *m, node *n, const uint64_t *key, node **out) {
161 if (!n) {
162 *out = new_node(m, key);
163 m->len++;
164 return *out;
165 }
166 int c = cmp(m, key, n->data);
167 if (c == 0) {
168 *out = n;
169 return n;
170 }
171 if (c < 0) n->l = insert(m, n->l, key, out);
172 else n->r = insert(m, n->r, key, out);
173 return balance(n);
174}
175
176void *momap_put(momap *m, const uint64_t *key, bool *fresh) {
177 size_t before = m->len;
178 node *got = NULL;
179 m->root = insert(m, m->root, key, &got);
180 if (fresh) *fresh = m->len != before;
181 return val_of(m, got);
182}
183
184/* Detaches the smallest node of n's subtree into *min. */
185static node *pop_min(node *n, node **min) {
186 if (!n->l) {
187 *min = n;
188 return n->r;
189 }
190 n->l = pop_min(n->l, min);
191 return balance(n);
192}
193
194static node *remove_node(momap *m, node *n, const uint64_t *key, node **gone) {
195 if (!n) return NULL;
196 int c = cmp(m, key, n->data);
197 if (c < 0) n->l = remove_node(m, n->l, key, gone);
198 else if (c > 0) n->r = remove_node(m, n->r, key, gone);
199 else {
200 *gone = n;
201 if (!n->l) return n->r;
202 if (!n->r) return n->l;
203 node *min = NULL;
204 node *r = pop_min(n->r, &min);
205 min->l = n->l;
206 min->r = r;
207 return balance(min);
208 }
209 return balance(n);
210}
211
212bool momap_take(momap *m, const uint64_t *key, void *out) {
213 node *gone = NULL;
214 m->root = remove_node(m, m->root, key, &gone);
215 if (!gone) return false;
216 m->len--;
217 if (out) memcpy(out, val_of(m, gone), m->vsize);
218 else if (m->free_val) m->free_val(val_of(m, gone));
219 free(gone);
220 return true;
221}
222
223bool momap_del(momap *m, const uint64_t *key) {
224 node *gone = NULL;
225 m->root = remove_node(m, m->root, key, &gone);
226 if (!gone) return false;
227 m->len--;
228 if (m->free_val) m->free_val(val_of(m, gone));
229 free(gone);
230 return true;
231}
232
233static void *give(const momap *m, const node *n, uint64_t *key_out) {
234 if (!n) return NULL;
235 if (key_out) memmove(key_out, n->data, 8 * m->kw);
236 return val_of(m, n);
237}
238
239void *momap_first(const momap *m, uint64_t *key_out) {
240 const node *n = m->root;
241 while (n && n->l) n = n->l;
242 return give(m, n, key_out);
243}
244
245void *momap_last(const momap *m, uint64_t *key_out) {
246 const node *n = m->root;
247 while (n && n->r) n = n->r;
248 return give(m, n, key_out);
249}
250
251/* The nearest node to key: strictly after (dir > 0, or at it too when eq), or before (dir < 0). */
252static const node *near(const momap *m, const uint64_t *key, int dir, bool eq) {
253 const node *n = m->root, *best = NULL;
254 while (n) {
255 int c = cmp(m, n->data, key);
256 if (c == 0 && eq) return n;
257 if (dir > 0) {
258 if (c > 0) {
259 best = n;
260 n = n->l;
261 } else n = n->r;
262 } else {
263 if (c < 0) {
264 best = n;
265 n = n->r;
266 } else n = n->l;
267 }
268 }
269 return best;
270}
271
272void *momap_next(const momap *m, const uint64_t *after, uint64_t *key_out) { return give(m, near(m, after, 1, false), key_out); }
273void *momap_prev(const momap *m, const uint64_t *before, uint64_t *key_out) { return give(m, near(m, before, -1, false), key_out); }
274void *momap_floor(const momap *m, const uint64_t *at, uint64_t *key_out) { return give(m, near(m, at, -1, true), key_out); }
275void *momap_ceil(const momap *m, const uint64_t *at, uint64_t *key_out) { return give(m, near(m, at, 1, true), key_out); }
276
277momap *momap_split_off(momap *m, const uint64_t *at) {
278 momap *o = momap_new(m->kw, m->vsize, m->free_val, m->copy_val);
279 uint64_t k[MOMAP_MAX_KEY];
280 for (;;) {
281 void *v = momap_ceil(m, at, k);
282 if (!v) break;
283 void *dst = momap_put(o, k, NULL);
284 momap_take(m, k, dst);
285 }
286 return o;
287}
288
289static void push_left(momap_iter *it, const node *n) {
290 for (; n; n = n->l) it->stack[it->top++] = n;
291}
292
293void momap_iter_init(momap_iter *it, const momap *m) {
294 it->top = 0;
295 it->m = m;
296 push_left(it, m->root);
297}
298
299void *momap_iter_next(momap_iter *it, const uint64_t **key) {
300 if (!it->top) return NULL;
301 const node *n = it->stack[--it->top];
302 push_left(it, n->r);
303 if (key) *key = n->data;
304 return val_of(it->m, n);
305}
306
307void *momap_iter_next1(momap_iter *it, uint64_t *key) {
308 const uint64_t *k;
309 void *v = momap_iter_next(it, &k);
310 if (v && key) *key = k[0];
311 return v;
312}