lib/generic-radix-tree.c

   1
   2 #include <linux/atomic.h>
   3 #include <linux/export.h>
   4 #include <linux/generic-radix-tree.h>
   5 #include <linux/gfp.h>
   6 #include <linux/kmemleak.h>
   7
   8 #define GENRADIX_ARY            (GENRADIX_NODE_SIZE / sizeof(struct genradix_node *))
   9 #define GENRADIX_ARY_SHIFT      ilog2(GENRADIX_ARY)
  10
  11 struct genradix_node {
  12         union {
  13                 /* Interior node: */
  14                 struct genradix_node    *children[GENRADIX_ARY];
  15
  16                 /* Leaf: */
  17                 u8                      data[GENRADIX_NODE_SIZE];
  18         };
  19 };
  20
  21 static inline int genradix_depth_shift(unsigned depth)
  22 {
  23         return GENRADIX_NODE_SHIFT + GENRADIX_ARY_SHIFT * depth;
  24 }
  25
  26 /*
  27  * Returns size (of data, in bytes) that a tree of a given depth holds:
  28  */
  29 static inline size_t genradix_depth_size(unsigned depth)
  30 {
  31         return 1UL << genradix_depth_shift(depth);
  32 }
  33
  34 /* depth that's needed for a genradix that can address up to ULONG_MAX: */
  35 #define GENRADIX_MAX_DEPTH      \
  36         DIV_ROUND_UP(BITS_PER_LONG - GENRADIX_NODE_SHIFT, GENRADIX_ARY_SHIFT)
  37
  38 #define GENRADIX_DEPTH_MASK                             \
  39         ((unsigned long) (roundup_pow_of_two(GENRADIX_MAX_DEPTH + 1) - 1))
  40
  41 static inline unsigned genradix_root_to_depth(struct genradix_root *r)
  42 {
  43         return (unsigned long) r & GENRADIX_DEPTH_MASK;
  44 }
  45
  46 static inline struct genradix_node *genradix_root_to_node(struct genradix_root *r)
  47 {
  48         return (void *) ((unsigned long) r & ~GENRADIX_DEPTH_MASK);
  49 }
  50
  51 /*
  52  * Returns pointer to the specified byte @offset within @radix, or NULL if not
  53  * allocated
  54  */
  55 void *__genradix_ptr(struct __genradix *radix, size_t offset)
  56 {
  57         struct genradix_root *r = READ_ONCE(radix->root);
  58         struct genradix_node *n = genradix_root_to_node(r);
  59         unsigned level          = genradix_root_to_depth(r);
  60
  61         if (ilog2(offset) >= genradix_depth_shift(level))
  62                 return NULL;
  63
  64         while (1) {
  65                 if (!n)
  66                         return NULL;
  67                 if (!level)
  68                         break;
  69
  70                 level--;
  71
  72                 n = n->children[offset >> genradix_depth_shift(level)];
  73                 offset &= genradix_depth_size(level) - 1;
  74         }
  75
  76         return &n->data[offset];
  77 }
  78 EXPORT_SYMBOL(__genradix_ptr);
  79
  80 static inline struct genradix_node *genradix_alloc_node(gfp_t gfp_mask)
  81 {
  82         return kzalloc(GENRADIX_NODE_SIZE, gfp_mask);
  83 }
  84
  85 static inline void genradix_free_node(struct genradix_node *node)
  86 {
  87         kfree(node);
  88 }
  89
  90 /*
  91  * Returns pointer to the specified byte @offset within @radix, allocating it if
  92  * necessary - newly allocated slots are always zeroed out:
  93  */
  94 void *__genradix_ptr_alloc(struct __genradix *radix, size_t offset,
  95                            gfp_t gfp_mask)
  96 {
  97         struct genradix_root *v = READ_ONCE(radix->root);
  98         struct genradix_node *n, *new_node = NULL;
  99         unsigned level;
 100
 101         /* Increase tree depth if necessary: */
 102         while (1) {
 103                 struct genradix_root *r = v, *new_root;
 104
 105                 n       = genradix_root_to_node(r);
 106                 level   = genradix_root_to_depth(r);
 107
 108                 if (n && ilog2(offset) < genradix_depth_shift(level))
 109                         break;
 110
 111                 if (!new_node) {
 112                         new_node = genradix_alloc_node(gfp_mask);
 113                         if (!new_node)
 114                                 return NULL;
 115                 }
 116
 117                 new_node->children[0] = n;
 118                 new_root = ((struct genradix_root *)
 119                             ((unsigned long) new_node | (n ? level + 1 : 0)));
 120
 121                 if ((v = cmpxchg_release(&radix->root, r, new_root)) == r) {
 122                         v = new_root;
 123                         new_node = NULL;
 124                 }
 125         }
 126
 127         while (level--) {
 128                 struct genradix_node **p =
 129                         &n->children[offset >> genradix_depth_shift(level)];
 130                 offset &= genradix_depth_size(level) - 1;
 131
 132                 n = READ_ONCE(*p);
 133                 if (!n) {
 134                         if (!new_node) {
 135                                 new_node = genradix_alloc_node(gfp_mask);
 136                                 if (!new_node)
 137                                         return NULL;
 138                         }
 139
 140                         if (!(n = cmpxchg_release(p, NULL, new_node)))
 141                                 swap(n, new_node);
 142                 }
 143         }
 144
 145         if (new_node)
 146                 genradix_free_node(new_node);
 147
 148         return &n->data[offset];
 149 }
 150 EXPORT_SYMBOL(__genradix_ptr_alloc);
 151
 152 void *__genradix_iter_peek(struct genradix_iter *iter,
 153                            struct __genradix *radix,
 154                            size_t objs_per_page)
 155 {
 156         struct genradix_root *r;
 157         struct genradix_node *n;
 158         unsigned level, i;
 159
 160         if (iter->offset == SIZE_MAX)
 161                 return NULL;
 162
 163 restart:
 164         r = READ_ONCE(radix->root);
 165         if (!r)
 166                 return NULL;
 167
 168         n       = genradix_root_to_node(r);
 169         level   = genradix_root_to_depth(r);
 170
 171         if (ilog2(iter->offset) >= genradix_depth_shift(level))
 172                 return NULL;
 173
 174         while (level) {
 175                 level--;
 176
 177                 i = (iter->offset >> genradix_depth_shift(level)) &
 178                         (GENRADIX_ARY - 1);
 179
 180                 while (!n->children[i]) {
 181                         size_t objs_per_ptr = genradix_depth_size(level);
 182
 183                         if (iter->offset + objs_per_ptr < iter->offset) {
 184                                 iter->offset    = SIZE_MAX;
 185                                 iter->pos       = SIZE_MAX;
 186                                 return NULL;
 187                         }
 188
 189                         i++;
 190                         iter->offset = round_down(iter->offset + objs_per_ptr,
 191                                                   objs_per_ptr);
 192                         iter->pos = (iter->offset >> GENRADIX_NODE_SHIFT) *
 193                                 objs_per_page;
 194                         if (i == GENRADIX_ARY)
 195                                 goto restart;
 196                 }
 197
 198                 n = n->children[i];
 199         }
 200
 201         return &n->data[iter->offset & (GENRADIX_NODE_SIZE - 1)];
 202 }
 203 EXPORT_SYMBOL(__genradix_iter_peek);
 204
 205 void *__genradix_iter_peek_prev(struct genradix_iter *iter,
 206                                 struct __genradix *radix,
 207                                 size_t objs_per_page,
 208                                 size_t obj_size_plus_page_remainder)
 209 {
 210         struct genradix_root *r;
 211         struct genradix_node *n;
 212         unsigned level, i;
 213
 214         if (iter->offset == SIZE_MAX)
 215                 return NULL;
 216
 217 restart:
 218         r = READ_ONCE(radix->root);
 219         if (!r)
 220                 return NULL;
 221
 222         n       = genradix_root_to_node(r);
 223         level   = genradix_root_to_depth(r);
 224
 225         if (ilog2(iter->offset) >= genradix_depth_shift(level)) {
 226                 iter->offset = genradix_depth_size(level);
 227                 iter->pos = (iter->offset >> GENRADIX_NODE_SHIFT) * objs_per_page;
 228
 229                 iter->offset -= obj_size_plus_page_remainder;
 230                 iter->pos--;
 231         }
 232
 233         while (level) {
 234                 level--;
 235
 236                 i = (iter->offset >> genradix_depth_shift(level)) &
 237                         (GENRADIX_ARY - 1);
 238
 239                 while (!n->children[i]) {
 240                         size_t objs_per_ptr = genradix_depth_size(level);
 241
 242                         iter->offset = round_down(iter->offset, objs_per_ptr);
 243                         iter->pos = (iter->offset >> GENRADIX_NODE_SHIFT) * objs_per_page;
 244
 245                         if (!iter->offset)
 246                                 return NULL;
 247
 248                         iter->offset -= obj_size_plus_page_remainder;
 249                         iter->pos--;
 250
 251                         if (!i)
 252                                 goto restart;
 253                         --i;
 254                 }
 255
 256                 n = n->children[i];
 257         }
 258
 259         return &n->data[iter->offset & (GENRADIX_NODE_SIZE - 1)];
 260 }
 261 EXPORT_SYMBOL(__genradix_iter_peek_prev);
 262
 263 static void genradix_free_recurse(struct genradix_node *n, unsigned level)
 264 {
 265         if (level) {
 266                 unsigned i;
 267
 268                 for (i = 0; i < GENRADIX_ARY; i++)
 269                         if (n->children[i])
 270                                 genradix_free_recurse(n->children[i], level - 1);
 271         }
 272
 273         genradix_free_node(n);
 274 }
 275
 276 int __genradix_prealloc(struct __genradix *radix, size_t size,
 277                         gfp_t gfp_mask)
 278 {
 279         size_t offset;
 280
 281         for (offset = 0; offset < size; offset += GENRADIX_NODE_SIZE)
 282                 if (!__genradix_ptr_alloc(radix, offset, gfp_mask))
 283                         return -ENOMEM;
 284
 285         return 0;
 286 }
 287 EXPORT_SYMBOL(__genradix_prealloc);
 288
 289 void __genradix_free(struct __genradix *radix)
 290 {
 291         struct genradix_root *r = xchg(&radix->root, NULL);
 292
 293         genradix_free_recurse(genradix_root_to_node(r),
 294                               genradix_root_to_depth(r));
 295 }
 296 EXPORT_SYMBOL(__genradix_free);