fs/bcachefs/eytzinger.h

   1 /* SPDX-License-Identifier: GPL-2.0 */
   2 #ifndef _EYTZINGER_H
   3 #define _EYTZINGER_H
   4
   5 #include <linux/bitops.h>
   6 #include <linux/log2.h>
   7
   8 #ifdef EYTZINGER_DEBUG
   9 #define EYTZINGER_BUG_ON(cond)          BUG_ON(cond)
  10 #else
  11 #define EYTZINGER_BUG_ON(cond)
  12 #endif
  13
  14 /*
  15  * Traversal for trees in eytzinger layout - a full binary tree layed out in an
  16  * array.
  17  *
  18  * Consider using an eytzinger tree any time you would otherwise be doing binary
  19  * search over an array. Binary search is a worst case scenario for branch
  20  * prediction and prefetching, but in an eytzinger tree every node's children
  21  * are adjacent in memory, thus we can prefetch children before knowing the
  22  * result of the comparison, assuming multiple nodes fit on a cacheline.
  23  *
  24  * Two variants are provided, for one based indexing and zero based indexing.
  25  *
  26  * Zero based indexing is more convenient, but one based indexing has better
  27  * alignment and thus better performance because each new level of the tree
  28  * starts at a power of two, and thus if element 0 was cacheline aligned, each
  29  * new level will be as well.
  30  */
  31
  32 static inline unsigned eytzinger1_child(unsigned i, unsigned child)
  33 {
  34         EYTZINGER_BUG_ON(child > 1);
  35
  36         return (i << 1) + child;
  37 }
  38
  39 static inline unsigned eytzinger1_left_child(unsigned i)
  40 {
  41         return eytzinger1_child(i, 0);
  42 }
  43
  44 static inline unsigned eytzinger1_right_child(unsigned i)
  45 {
  46         return eytzinger1_child(i, 1);
  47 }
  48
  49 static inline unsigned eytzinger1_first(unsigned size)
  50 {
  51         return size ? rounddown_pow_of_two(size) : 0;
  52 }
  53
  54 static inline unsigned eytzinger1_last(unsigned size)
  55 {
  56         return rounddown_pow_of_two(size + 1) - 1;
  57 }
  58
  59 /*
  60  * eytzinger1_next() and eytzinger1_prev() have the nice properties that
  61  *
  62  * eytzinger1_next(0) == eytzinger1_first())
  63  * eytzinger1_prev(0) == eytzinger1_last())
  64  *
  65  * eytzinger1_prev(eytzinger1_first()) == 0
  66  * eytzinger1_next(eytzinger1_last()) == 0
  67  */
  68
  69 static inline unsigned eytzinger1_next(unsigned i, unsigned size)
  70 {
  71         EYTZINGER_BUG_ON(i > size);
  72
  73         if (eytzinger1_right_child(i) <= size) {
  74                 i = eytzinger1_right_child(i);
  75
  76                 i <<= __fls(size + 1) - __fls(i);
  77                 i >>= i > size;
  78         } else {
  79                 i >>= ffz(i) + 1;
  80         }
  81
  82         return i;
  83 }
  84
  85 static inline unsigned eytzinger1_prev(unsigned i, unsigned size)
  86 {
  87         EYTZINGER_BUG_ON(i > size);
  88
  89         if (eytzinger1_left_child(i) <= size) {
  90                 i = eytzinger1_left_child(i) + 1;
  91
  92                 i <<= __fls(size + 1) - __fls(i);
  93                 i -= 1;
  94                 i >>= i > size;
  95         } else {
  96                 i >>= __ffs(i) + 1;
  97         }
  98
  99         return i;
 100 }
 101
 102 static inline unsigned eytzinger1_extra(unsigned size)
 103 {
 104         return size
 105                 ? (size + 1 - rounddown_pow_of_two(size)) << 1
 106                 : 0;
 107 }
 108
 109 static inline unsigned __eytzinger1_to_inorder(unsigned i, unsigned size,
 110                                               unsigned extra)
 111 {
 112         unsigned b = __fls(i);
 113         unsigned shift = __fls(size) - b;
 114         int s;
 115
 116         EYTZINGER_BUG_ON(!i || i > size);
 117
 118         i  ^= 1U << b;
 119         i <<= 1;
 120         i  |= 1;
 121         i <<= shift;
 122
 123         /*
 124          * sign bit trick:
 125          *
 126          * if (i > extra)
 127          *      i -= (i - extra) >> 1;
 128          */
 129         s = extra - i;
 130         i += (s >> 1) & (s >> 31);
 131
 132         return i;
 133 }
 134
 135 static inline unsigned __inorder_to_eytzinger1(unsigned i, unsigned size,
 136                                                unsigned extra)
 137 {
 138         unsigned shift;
 139         int s;
 140
 141         EYTZINGER_BUG_ON(!i || i > size);
 142
 143         /*
 144          * sign bit trick:
 145          *
 146          * if (i > extra)
 147          *      i += i - extra;
 148          */
 149         s = extra - i;
 150         i -= s & (s >> 31);
 151
 152         shift = __ffs(i);
 153
 154         i >>= shift + 1;
 155         i  |= 1U << (__fls(size) - shift);
 156
 157         return i;
 158 }
 159
 160 static inline unsigned eytzinger1_to_inorder(unsigned i, unsigned size)
 161 {
 162         return __eytzinger1_to_inorder(i, size, eytzinger1_extra(size));
 163 }
 164
 165 static inline unsigned inorder_to_eytzinger1(unsigned i, unsigned size)
 166 {
 167         return __inorder_to_eytzinger1(i, size, eytzinger1_extra(size));
 168 }
 169
 170 #define eytzinger1_for_each(_i, _size)                  \
 171         for (unsigned (_i) = eytzinger1_first((_size)); \
 172              (_i) != 0;                                 \
 173              (_i) = eytzinger1_next((_i), (_size)))
 174
 175 /* Zero based indexing version: */
 176
 177 static inline unsigned eytzinger0_child(unsigned i, unsigned child)
 178 {
 179         EYTZINGER_BUG_ON(child > 1);
 180
 181         return (i << 1) + 1 + child;
 182 }
 183
 184 static inline unsigned eytzinger0_left_child(unsigned i)
 185 {
 186         return eytzinger0_child(i, 0);
 187 }
 188
 189 static inline unsigned eytzinger0_right_child(unsigned i)
 190 {
 191         return eytzinger0_child(i, 1);
 192 }
 193
 194 static inline unsigned eytzinger0_first(unsigned size)
 195 {
 196         return eytzinger1_first(size) - 1;
 197 }
 198
 199 static inline unsigned eytzinger0_last(unsigned size)
 200 {
 201         return eytzinger1_last(size) - 1;
 202 }
 203
 204 static inline unsigned eytzinger0_next(unsigned i, unsigned size)
 205 {
 206         return eytzinger1_next(i + 1, size) - 1;
 207 }
 208
 209 static inline unsigned eytzinger0_prev(unsigned i, unsigned size)
 210 {
 211         return eytzinger1_prev(i + 1, size) - 1;
 212 }
 213
 214 static inline unsigned eytzinger0_extra(unsigned size)
 215 {
 216         return eytzinger1_extra(size);
 217 }
 218
 219 static inline unsigned __eytzinger0_to_inorder(unsigned i, unsigned size,
 220                                                unsigned extra)
 221 {
 222         return __eytzinger1_to_inorder(i + 1, size, extra) - 1;
 223 }
 224
 225 static inline unsigned __inorder_to_eytzinger0(unsigned i, unsigned size,
 226                                                unsigned extra)
 227 {
 228         return __inorder_to_eytzinger1(i + 1, size, extra) - 1;
 229 }
 230
 231 static inline unsigned eytzinger0_to_inorder(unsigned i, unsigned size)
 232 {
 233         return __eytzinger0_to_inorder(i, size, eytzinger0_extra(size));
 234 }
 235
 236 static inline unsigned inorder_to_eytzinger0(unsigned i, unsigned size)
 237 {
 238         return __inorder_to_eytzinger0(i, size, eytzinger0_extra(size));
 239 }
 240
 241 #define eytzinger0_for_each(_i, _size)                  \
 242         for (unsigned (_i) = eytzinger0_first((_size)); \
 243              (_i) != -1;                                \
 244              (_i) = eytzinger0_next((_i), (_size)))
 245
 246 /* return greatest node <= @search, or -1 if not found */
 247 static inline int eytzinger0_find_le(void *base, size_t nr, size_t size,
 248                                      cmp_func_t cmp, const void *search)
 249 {
 250         unsigned i, n = 0;
 251
 252         if (!nr)
 253                 return -1;
 254
 255         do {
 256                 i = n;
 257                 n = eytzinger0_child(i, cmp(base + i * size, search) <= 0);
 258         } while (n < nr);
 259
 260         if (n & 1) {
 261                 /*
 262                  * @i was greater than @search, return previous node:
 263                  *
 264                  * if @i was leftmost/smallest element,
 265                  * eytzinger0_prev(eytzinger0_first())) returns -1, as expected
 266                  */
 267                 return eytzinger0_prev(i, nr);
 268         } else {
 269                 return i;
 270         }
 271 }
 272
 273 static inline int eytzinger0_find_gt(void *base, size_t nr, size_t size,
 274                                      cmp_func_t cmp, const void *search)
 275 {
 276         ssize_t idx = eytzinger0_find_le(base, nr, size, cmp, search);
 277
 278         /*
 279          * if eytitzinger0_find_le() returned -1 - no element was <= search - we
 280          * want to return the first element; next/prev identities mean this work
 281          * as expected
 282          *
 283          * similarly if find_le() returns last element, we should return -1;
 284          * identities mean this all works out:
 285          */
 286         return eytzinger0_next(idx, nr);
 287 }
 288
 289 static inline int eytzinger0_find_ge(void *base, size_t nr, size_t size,
 290                                      cmp_func_t cmp, const void *search)
 291 {
 292         ssize_t idx = eytzinger0_find_le(base, nr, size, cmp, search);
 293
 294         if (idx < nr && !cmp(base + idx * size, search))
 295                 return idx;
 296
 297         return eytzinger0_next(idx, nr);
 298 }
 299
 300 #define eytzinger0_find(base, nr, size, _cmp, search)                   \
 301 ({                                                                      \
 302         void *_base             = (base);                               \
 303         const void *_search     = (search);                             \
 304         size_t _nr              = (nr);                                 \
 305         size_t _size            = (size);                               \
 306         size_t _i               = 0;                                    \
 307         int _res;                                                       \
 308                                                                         \
 309         while (_i < _nr &&                                              \
 310                (_res = _cmp(_search, _base + _i * _size)))              \
 311                 _i = eytzinger0_child(_i, _res > 0);                    \
 312         _i;                                                             \
 313 })
 314
 315 void eytzinger0_sort_r(void *, size_t, size_t,
 316                        cmp_r_func_t, swap_r_func_t, const void *);
 317 void eytzinger0_sort(void *, size_t, size_t, cmp_func_t, swap_func_t);
 318
 319 #endif /* _EYTZINGER_H */