xref: /linux/lib/generic-radix-tree.c (revision 56d8b784c56588cd40f98e4b1d4f6e29e3cb02b8)
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 		} else {
125 			new_node->children[0] = NULL;
126 		}
127 	}
128 
129 	while (level--) {
130 		struct genradix_node **p =
131 			&n->children[offset >> genradix_depth_shift(level)];
132 		offset &= genradix_depth_size(level) - 1;
133 
134 		n = READ_ONCE(*p);
135 		if (!n) {
136 			if (!new_node) {
137 				new_node = genradix_alloc_node(gfp_mask);
138 				if (!new_node)
139 					return NULL;
140 			}
141 
142 			if (!(n = cmpxchg_release(p, NULL, new_node)))
143 				swap(n, new_node);
144 		}
145 	}
146 
147 	if (new_node)
148 		genradix_free_node(new_node);
149 
150 	return &n->data[offset];
151 }
152 EXPORT_SYMBOL(__genradix_ptr_alloc);
153 
154 void *__genradix_iter_peek(struct genradix_iter *iter,
155 			   struct __genradix *radix,
156 			   size_t objs_per_page)
157 {
158 	struct genradix_root *r;
159 	struct genradix_node *n;
160 	unsigned level, i;
161 
162 	if (iter->offset == SIZE_MAX)
163 		return NULL;
164 
165 restart:
166 	r = READ_ONCE(radix->root);
167 	if (!r)
168 		return NULL;
169 
170 	n	= genradix_root_to_node(r);
171 	level	= genradix_root_to_depth(r);
172 
173 	if (ilog2(iter->offset) >= genradix_depth_shift(level))
174 		return NULL;
175 
176 	while (level) {
177 		level--;
178 
179 		i = (iter->offset >> genradix_depth_shift(level)) &
180 			(GENRADIX_ARY - 1);
181 
182 		while (!n->children[i]) {
183 			size_t objs_per_ptr = genradix_depth_size(level);
184 
185 			if (iter->offset + objs_per_ptr < iter->offset) {
186 				iter->offset	= SIZE_MAX;
187 				iter->pos	= SIZE_MAX;
188 				return NULL;
189 			}
190 
191 			i++;
192 			iter->offset = round_down(iter->offset + objs_per_ptr,
193 						  objs_per_ptr);
194 			iter->pos = (iter->offset >> GENRADIX_NODE_SHIFT) *
195 				objs_per_page;
196 			if (i == GENRADIX_ARY)
197 				goto restart;
198 		}
199 
200 		n = n->children[i];
201 	}
202 
203 	return &n->data[iter->offset & (GENRADIX_NODE_SIZE - 1)];
204 }
205 EXPORT_SYMBOL(__genradix_iter_peek);
206 
207 void *__genradix_iter_peek_prev(struct genradix_iter *iter,
208 				struct __genradix *radix,
209 				size_t objs_per_page,
210 				size_t obj_size_plus_page_remainder)
211 {
212 	struct genradix_root *r;
213 	struct genradix_node *n;
214 	unsigned level, i;
215 
216 	if (iter->offset == SIZE_MAX)
217 		return NULL;
218 
219 restart:
220 	r = READ_ONCE(radix->root);
221 	if (!r)
222 		return NULL;
223 
224 	n	= genradix_root_to_node(r);
225 	level	= genradix_root_to_depth(r);
226 
227 	if (ilog2(iter->offset) >= genradix_depth_shift(level)) {
228 		iter->offset = genradix_depth_size(level);
229 		iter->pos = (iter->offset >> GENRADIX_NODE_SHIFT) * objs_per_page;
230 
231 		iter->offset -= obj_size_plus_page_remainder;
232 		iter->pos--;
233 	}
234 
235 	while (level) {
236 		level--;
237 
238 		i = (iter->offset >> genradix_depth_shift(level)) &
239 			(GENRADIX_ARY - 1);
240 
241 		while (!n->children[i]) {
242 			size_t objs_per_ptr = genradix_depth_size(level);
243 
244 			iter->offset = round_down(iter->offset, objs_per_ptr);
245 			iter->pos = (iter->offset >> GENRADIX_NODE_SHIFT) * objs_per_page;
246 
247 			if (!iter->offset)
248 				return NULL;
249 
250 			iter->offset -= obj_size_plus_page_remainder;
251 			iter->pos--;
252 
253 			if (!i)
254 				goto restart;
255 			--i;
256 		}
257 
258 		n = n->children[i];
259 	}
260 
261 	return &n->data[iter->offset & (GENRADIX_NODE_SIZE - 1)];
262 }
263 EXPORT_SYMBOL(__genradix_iter_peek_prev);
264 
265 static void genradix_free_recurse(struct genradix_node *n, unsigned level)
266 {
267 	if (level) {
268 		unsigned i;
269 
270 		for (i = 0; i < GENRADIX_ARY; i++)
271 			if (n->children[i])
272 				genradix_free_recurse(n->children[i], level - 1);
273 	}
274 
275 	genradix_free_node(n);
276 }
277 
278 int __genradix_prealloc(struct __genradix *radix, size_t size,
279 			gfp_t gfp_mask)
280 {
281 	size_t offset;
282 
283 	for (offset = 0; offset < size; offset += GENRADIX_NODE_SIZE)
284 		if (!__genradix_ptr_alloc(radix, offset, gfp_mask))
285 			return -ENOMEM;
286 
287 	return 0;
288 }
289 EXPORT_SYMBOL(__genradix_prealloc);
290 
291 void __genradix_free(struct __genradix *radix)
292 {
293 	struct genradix_root *r = xchg(&radix->root, NULL);
294 
295 	genradix_free_recurse(genradix_root_to_node(r),
296 			      genradix_root_to_depth(r));
297 }
298 EXPORT_SYMBOL(__genradix_free);
299