intmap: reimplement so that intmap_after works.
[ccan] / ccan / intmap / intmap.c
1 /* CC0 license (public domain) - see LICENSE file for details */
2 /* This code is based on ccan/strmap.c. */
3 #include <ccan/bitops/bitops.h>
4 #include <ccan/intmap/intmap.h>
5 #include <ccan/short_types/short_types.h>
6 #include <ccan/str/str.h>
7 #include <assert.h>
8 #include <stdlib.h>
9 #include <errno.h>
10
11 struct node {
12         /* These point to strings or nodes. */
13         struct intmap child[2];
14         /* Encoding both prefix and critbit: 1 is appended to prefix. */
15         intmap_index_t prefix_and_critbit;
16 };
17
18 static int critbit(const struct intmap *n)
19 {
20         return bitops_ls64(n->u.n->prefix_and_critbit);
21 }
22
23 static intmap_index_t prefix_mask(int critbit)
24 {
25         /* Mask does not include critbit itself, but can't shift by critbit+1 */
26         return -2ULL << critbit;
27 }
28
29 static intmap_index_t prefix_and_critbit(intmap_index_t v, int n)
30 {
31         intmap_index_t critbit = ((intmap_index_t)1 << n);
32         return (v & ~(critbit - 1)) | critbit;
33 }
34
35 void *intmap_get_(const struct intmap *map, intmap_index_t index)
36 {
37         /* Not empty map? */
38         if (!intmap_empty_(map)) {
39                 const struct intmap *n = map;
40                 /* Anything with NULL value is a node. */
41                 while (!n->v) {
42                         /* FIXME: compare cmp prefix, if not equal, ENOENT */
43                         u8 direction = (index >> critbit(n)) & 1;
44                         n = &n->u.n->child[direction];
45                 }
46                 if (index == n->u.i)
47                         return n->v;
48         }
49         errno = ENOENT;
50         return NULL;
51 }
52
53 static bool split_node(struct intmap *n, intmap_index_t nodeindex,
54                        intmap_index_t index, const void *value)
55 {
56         struct node *newn;
57         int new_dir;
58
59         /* Find highest bit where they differ. */
60         unsigned int critbit = bitops_hs64(nodeindex ^ index);
61         assert(critbit < CHAR_BIT*sizeof(index));
62
63         /* Which direction do we go at this bit? */
64         new_dir = (index >> critbit) & 1;
65
66         /* Allocate new node. */
67         newn = malloc(sizeof(*newn));
68         if (!newn) {
69                 errno = ENOMEM;
70                 return false;
71         }
72         newn->prefix_and_critbit = prefix_and_critbit(index, critbit);
73         newn->child[new_dir].v = (void *)value;
74         newn->child[new_dir].u.i = index;
75         newn->child[!new_dir] = *n;
76
77         n->u.n = newn;
78         n->v = NULL;
79         return true;
80 }
81
82 bool intmap_add_(struct intmap *map, intmap_index_t index, const void *value)
83 {
84         struct intmap *n;
85
86         assert(value);
87
88         /* Empty map? */
89         if (intmap_empty_(map)) {
90                 map->u.i = index;
91                 map->v = (void *)value;
92                 return true;
93         }
94
95         n = map;
96         /* Anything with NULL value is a node. */
97         while (!n->v) {
98                 int crit = critbit(n);
99                 intmap_index_t mask = prefix_mask(crit);
100                 u8 direction = (index >> crit) & 1;
101
102                 if ((index & mask) != (n->u.n->prefix_and_critbit & mask))
103                         return split_node(n, n->u.n->prefix_and_critbit & mask,
104                                           index, value);
105                 n = &n->u.n->child[direction];
106         }
107
108         if (index == n->u.i) {
109                 errno = EEXIST;
110                 return false;
111         }
112
113         return split_node(n, n->u.i, index, value);
114 }
115
116 void *intmap_del_(struct intmap *map, intmap_index_t index)
117 {
118         struct intmap *parent = NULL, *n;
119         u8 direction;
120         void *value;
121
122         /* Empty map? */
123         if (intmap_empty_(map)) {
124                 errno = ENOENT;
125                 return NULL;
126         }
127
128         /* Find closest, but keep track of parent. */
129         n = map;
130         /* Anything with NULL value is a node. */
131         while (!n->v) {
132                 /* FIXME: compare cmp prefix, if not equal, ENOENT */
133                 parent = n;
134                 direction = (index >> critbit(n)) & 1;
135                 n = &n->u.n->child[direction];
136         }
137
138         /* Did we find it? */
139         if (index != n->u.i) {
140                 errno = ENOENT;
141                 return NULL;
142         }
143
144         value = n->v;
145
146         if (!parent) {
147                 /* We deleted last node. */
148                 intmap_init_(map);
149         } else {
150                 struct node *old = parent->u.n;
151                 /* Raise other node to parent. */
152                 *parent = old->child[!direction];
153                 free(old);
154         }
155         errno = 0;
156         return value;
157 }
158
159 void *intmap_first_(const struct intmap *map, intmap_index_t *indexp)
160 {
161         const struct intmap *n;
162
163         if (intmap_empty_(map)) {
164                 errno = ENOENT;
165                 return NULL;
166         }
167         
168         n = map;
169         /* Anything with NULL value is a node. */
170         while (!n->v)
171                 n = &n->u.n->child[0];
172         errno = 0;
173         *indexp = n->u.i;
174         return n->v;
175 }
176                 
177 void *intmap_after_(const struct intmap *map, intmap_index_t *indexp)
178 {
179         const struct intmap *n, *prev = NULL;
180         intmap_index_t index = (*indexp) + 1;
181
182         /* Special case of overflow */
183         if (index == 0)
184                 goto none_left;
185
186         /* Special case of empty map */
187         if (intmap_empty_(map))
188                 goto none_left;
189
190         /* Follow down, until prefix differs. */
191         n = map;
192         while (!n->v) {
193                 int crit = critbit(n);
194                 u8 direction;
195                 intmap_index_t prefix, idx;
196
197                 idx = (index >> crit);
198                 direction = idx & 1;
199
200                 /* Leave critbit in place: we can't shift by 64 anyway */
201                 idx |= 1;
202                 prefix = n->u.n->prefix_and_critbit >> crit;
203
204                 /* If this entire tree is greater than index, take first */
205                 if (idx < prefix)
206                         return intmap_first_(n, indexp);
207                 /* If this entire tree is less than index, we're past it. */
208                 else if (idx > prefix)
209                         goto try_greater_tree;
210
211                 /* Remember greater tree for backtracking */
212                 if (!direction)
213                         prev = n;
214                 n = &n->u.n->child[direction];
215         }
216
217         /* Found a successor? */
218         if (n->u.i >= index) {
219                 errno = 0;
220                 *indexp = n->u.i;
221                 return n->v;
222         }
223
224 try_greater_tree:
225         /* If we ever took a lesser branch, go back to greater branch */
226         if (prev)
227                 return intmap_first_(&prev->u.n->child[1], indexp);
228
229 none_left:
230         errno = ENOENT;
231         return NULL;
232 }
233
234 void *intmap_last_(const struct intmap *map, intmap_index_t *indexp)
235 {
236         const struct intmap *n;
237
238         if (intmap_empty_(map)) {
239                 errno = ENOENT;
240                 return NULL;
241         }
242
243         n = map;
244         /* Anything with NULL value is a node. */
245         while (!n->v)
246                 n = &n->u.n->child[1];
247         errno = 0;
248         *indexp = n->u.i;
249         return n->v;
250 }
251
252 static void clear(struct intmap n)
253 {
254         if (!n.v) {
255                 clear(n.u.n->child[0]);
256                 clear(n.u.n->child[1]);
257                 free(n.u.n);
258         }
259 }
260
261 void intmap_clear_(struct intmap *map)
262 {
263         if (!intmap_empty_(map))
264                 clear(*map);
265         intmap_init_(map);
266 }