git.ozlabs.org Git - ccan/blob - 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_before_(const struct intmap *map, intmap_index_t *indexp)
 235 {
 236         const struct intmap *n, *prev = NULL;
 237         intmap_index_t index = (*indexp) - 1;
 238
 239         /* Special case of overflow */
 240         if (index == (intmap_index_t)-1ULL)
 241                 goto none_left;
 242
 243         /* Special case of empty map */
 244         if (intmap_empty_(map))
 245                 goto none_left;
 246
 247         /* Follow down, until prefix differs. */
 248         n = map;
 249         while (!n->v) {
 250                 int crit = critbit(n);
 251                 u8 direction;
 252                 intmap_index_t prefix, idx;
 253
 254                 idx = (index >> crit);
 255                 direction = idx & 1;
 256
 257                 /* Leave critbit in place: we can't shift by 64 anyway */
 258                 idx |= 1;
 259                 prefix = n->u.n->prefix_and_critbit >> crit;
 260
 261                 /* If this entire tree is less than index, take last */
 262                 if (idx > prefix)
 263                         return intmap_last_(n, indexp);
 264                 /* If this entire tree is greater than index, we're past it. */
 265                 else if (idx < prefix)
 266                         goto try_lesser_tree;
 267
 268                 /* Remember lesser tree for backtracking */
 269                 if (direction)
 270                         prev = n;
 271                 n = &n->u.n->child[direction];
 272         }
 273
 274         /* Found a predecessor? */
 275         if (n->u.i <= index) {
 276                 errno = 0;
 277                 *indexp = n->u.i;
 278                 return n->v;
 279         }
 280
 281 try_lesser_tree:
 282         /* If we ever took a lesser branch, go back to lesser branch */
 283         if (prev)
 284                 return intmap_last_(&prev->u.n->child[0], indexp);
 285
 286 none_left:
 287         errno = ENOENT;
 288         return NULL;
 289 }
 290
 291 void *intmap_last_(const struct intmap *map, intmap_index_t *indexp)
 292 {
 293         const struct intmap *n;
 294
 295         if (intmap_empty_(map)) {
 296                 errno = ENOENT;
 297                 return NULL;
 298         }
 299
 300         n = map;
 301         /* Anything with NULL value is a node. */
 302         while (!n->v)
 303                 n = &n->u.n->child[1];
 304         errno = 0;
 305         *indexp = n->u.i;
 306         return n->v;
 307 }
 308
 309 static void clear(struct intmap n)
 310 {
 311         if (!n.v) {
 312                 clear(n.u.n->child[0]);
 313                 clear(n.u.n->child[1]);
 314                 free(n.u.n);
 315         }
 316 }
 317
 318 void intmap_clear_(struct intmap *map)
 319 {
 320         if (!intmap_empty_(map))
 321                 clear(*map);
 322         intmap_init_(map);
 323 }
 324
 325 bool intmap_iterate_(const struct intmap *n,
 326                      bool (*handle)(intmap_index_t, void *, void *),
 327                      void *data,
 328                      intmap_index_t offset)
 329 {
 330         /* Can only happen at root */
 331         if (intmap_empty_(n))
 332                 return true;
 333
 334         if (n->v)
 335                 return handle(n->u.i - offset, n->v, data);
 336
 337         return intmap_iterate_(&n->u.n->child[0], handle, data, offset)
 338                 && intmap_iterate_(&n->u.n->child[1], handle, data, offset);
 339 }