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#include "util.h"
  9
 10/*
 11 * Traversal for trees in eytzinger layout - a full binary tree layed out in an
 12 * array
 13 */
 14
 15/*
 16 * One based indexing version:
 17 *
 18 * With one based indexing each level of the tree starts at a power of two -
 19 * good for cacheline alignment:
 20 */
 21
 22static inline unsigned eytzinger1_child(unsigned i, unsigned child)
 23{
 24	EBUG_ON(child > 1);
 25
 26	return (i << 1) + child;
 27}
 28
 29static inline unsigned eytzinger1_left_child(unsigned i)
 30{
 31	return eytzinger1_child(i, 0);
 32}
 33
 34static inline unsigned eytzinger1_right_child(unsigned i)
 35{
 36	return eytzinger1_child(i, 1);
 37}
 38
 39static inline unsigned eytzinger1_first(unsigned size)
 40{
 41	return rounddown_pow_of_two(size);
 42}
 43
 44static inline unsigned eytzinger1_last(unsigned size)
 45{
 46	return rounddown_pow_of_two(size + 1) - 1;
 47}
 48
 49/*
 50 * eytzinger1_next() and eytzinger1_prev() have the nice properties that
 51 *
 52 * eytzinger1_next(0) == eytzinger1_first())
 53 * eytzinger1_prev(0) == eytzinger1_last())
 54 *
 55 * eytzinger1_prev(eytzinger1_first()) == 0
 56 * eytzinger1_next(eytzinger1_last()) == 0
 57 */
 58
 59static inline unsigned eytzinger1_next(unsigned i, unsigned size)
 60{
 61	EBUG_ON(i > size);
 62
 63	if (eytzinger1_right_child(i) <= size) {
 64		i = eytzinger1_right_child(i);
 65
 66		i <<= __fls(size + 1) - __fls(i);
 67		i >>= i > size;
 68	} else {
 69		i >>= ffz(i) + 1;
 70	}
 71
 72	return i;
 73}
 74
 75static inline unsigned eytzinger1_prev(unsigned i, unsigned size)
 76{
 77	EBUG_ON(i > size);
 78
 79	if (eytzinger1_left_child(i) <= size) {
 80		i = eytzinger1_left_child(i) + 1;
 81
 82		i <<= __fls(size + 1) - __fls(i);
 83		i -= 1;
 84		i >>= i > size;
 85	} else {
 86		i >>= __ffs(i) + 1;
 87	}
 88
 89	return i;
 90}
 91
 92static inline unsigned eytzinger1_extra(unsigned size)
 93{
 94	return (size + 1 - rounddown_pow_of_two(size)) << 1;
 95}
 96
 97static inline unsigned __eytzinger1_to_inorder(unsigned i, unsigned size,
 98					      unsigned extra)
 99{
100	unsigned b = __fls(i);
101	unsigned shift = __fls(size) - b;
102	int s;
103
104	EBUG_ON(!i || i > size);
105
106	i  ^= 1U << b;
107	i <<= 1;
108	i  |= 1;
109	i <<= shift;
110
111	/*
112	 * sign bit trick:
113	 *
114	 * if (i > extra)
115	 *	i -= (i - extra) >> 1;
116	 */
117	s = extra - i;
118	i += (s >> 1) & (s >> 31);
119
120	return i;
121}
122
123static inline unsigned __inorder_to_eytzinger1(unsigned i, unsigned size,
124					       unsigned extra)
125{
126	unsigned shift;
127	int s;
128
129	EBUG_ON(!i || i > size);
130
131	/*
132	 * sign bit trick:
133	 *
134	 * if (i > extra)
135	 *	i += i - extra;
136	 */
137	s = extra - i;
138	i -= s & (s >> 31);
139
140	shift = __ffs(i);
141
142	i >>= shift + 1;
143	i  |= 1U << (__fls(size) - shift);
144
145	return i;
146}
147
148static inline unsigned eytzinger1_to_inorder(unsigned i, unsigned size)
149{
150	return __eytzinger1_to_inorder(i, size, eytzinger1_extra(size));
151}
152
153static inline unsigned inorder_to_eytzinger1(unsigned i, unsigned size)
154{
155	return __inorder_to_eytzinger1(i, size, eytzinger1_extra(size));
156}
157
158#define eytzinger1_for_each(_i, _size)			\
159	for (unsigned (_i) = eytzinger1_first((_size));	\
160	     (_i) != 0;					\
161	     (_i) = eytzinger1_next((_i), (_size)))
162
163/* Zero based indexing version: */
164
165static inline unsigned eytzinger0_child(unsigned i, unsigned child)
166{
167	EBUG_ON(child > 1);
168
169	return (i << 1) + 1 + child;
170}
171
172static inline unsigned eytzinger0_left_child(unsigned i)
173{
174	return eytzinger0_child(i, 0);
175}
176
177static inline unsigned eytzinger0_right_child(unsigned i)
178{
179	return eytzinger0_child(i, 1);
180}
181
182static inline unsigned eytzinger0_first(unsigned size)
183{
184	return eytzinger1_first(size) - 1;
185}
186
187static inline unsigned eytzinger0_last(unsigned size)
188{
189	return eytzinger1_last(size) - 1;
190}
191
192static inline unsigned eytzinger0_next(unsigned i, unsigned size)
193{
194	return eytzinger1_next(i + 1, size) - 1;
195}
196
197static inline unsigned eytzinger0_prev(unsigned i, unsigned size)
198{
199	return eytzinger1_prev(i + 1, size) - 1;
200}
201
202static inline unsigned eytzinger0_extra(unsigned size)
203{
204	return eytzinger1_extra(size);
205}
206
207static inline unsigned __eytzinger0_to_inorder(unsigned i, unsigned size,
208					       unsigned extra)
209{
210	return __eytzinger1_to_inorder(i + 1, size, extra) - 1;
211}
212
213static inline unsigned __inorder_to_eytzinger0(unsigned i, unsigned size,
214					       unsigned extra)
215{
216	return __inorder_to_eytzinger1(i + 1, size, extra) - 1;
217}
218
219static inline unsigned eytzinger0_to_inorder(unsigned i, unsigned size)
220{
221	return __eytzinger0_to_inorder(i, size, eytzinger0_extra(size));
222}
223
224static inline unsigned inorder_to_eytzinger0(unsigned i, unsigned size)
225{
226	return __inorder_to_eytzinger0(i, size, eytzinger0_extra(size));
227}
228
229#define eytzinger0_for_each(_i, _size)			\
230	for (unsigned (_i) = eytzinger0_first((_size));	\
231	     (_i) != -1;				\
232	     (_i) = eytzinger0_next((_i), (_size)))
233
234typedef int (*eytzinger_cmp_fn)(const void *l, const void *r, size_t size);
235
236/* return greatest node <= @search, or -1 if not found */
237static inline ssize_t eytzinger0_find_le(void *base, size_t nr, size_t size,
238					 eytzinger_cmp_fn cmp, const void *search)
239{
240	unsigned i, n = 0;
241
242	if (!nr)
243		return -1;
244
245	do {
246		i = n;
247		n = eytzinger0_child(i, cmp(search, base + i * size, size) >= 0);
248	} while (n < nr);
249
250	if (n & 1) {
251		/* @i was greater than @search, return previous node: */
252
253		if (i == eytzinger0_first(nr))
254			return -1;
255
256		return eytzinger0_prev(i, nr);
257	} else {
258		return i;
259	}
260}
261
262#define eytzinger0_find(base, nr, size, _cmp, search)			\
263({									\
264	void *_base		= (base);				\
265	const void *_search	= (search);				\
266	size_t _nr		= (nr);					\
267	size_t _size		= (size);				\
268	size_t _i		= 0;					\
269	int _res;							\
270									\
271	while (_i < _nr &&						\
272	       (_res = _cmp(_search, _base + _i * _size, _size)))	\
273		_i = eytzinger0_child(_i, _res > 0);			\
274	_i;								\
275})
276
277void eytzinger0_sort(void *, size_t, size_t,
278		    int (*cmp_func)(const void *, const void *, size_t),
279		    void (*swap_func)(void *, void *, size_t));
280
281#endif /* _EYTZINGER_H */