300 lines · c
1// SPDX-License-Identifier: GPL-2.02/*3 * Copyright (C) 2011 STRATO AG4 * written by Arne Jansen <sensille@gmx.net>5 */6 7#include <linux/slab.h>8#include "messages.h"9#include "ulist.h"10 11/*12 * ulist is a generic data structure to hold a collection of unique u6413 * values. The only operations it supports is adding to the list and14 * enumerating it.15 * It is possible to store an auxiliary value along with the key.16 *17 * A sample usage for ulists is the enumeration of directed graphs without18 * visiting a node twice. The pseudo-code could look like this:19 *20 * ulist = ulist_alloc();21 * ulist_add(ulist, root);22 * ULIST_ITER_INIT(&uiter);23 *24 * while ((elem = ulist_next(ulist, &uiter)) {25 * for (all child nodes n in elem)26 * ulist_add(ulist, n);27 * do something useful with the node;28 * }29 * ulist_free(ulist);30 *31 * This assumes the graph nodes are addressable by u64. This stems from the32 * usage for tree enumeration in btrfs, where the logical addresses are33 * 64 bit.34 *35 * It is also useful for tree enumeration which could be done elegantly36 * recursively, but is not possible due to kernel stack limitations. The37 * loop would be similar to the above.38 */39 40/*41 * Freshly initialize a ulist.42 *43 * @ulist: the ulist to initialize44 *45 * Note: don't use this function to init an already used ulist, use46 * ulist_reinit instead.47 */48void ulist_init(struct ulist *ulist)49{50 INIT_LIST_HEAD(&ulist->nodes);51 ulist->root = RB_ROOT;52 ulist->nnodes = 0;53 ulist->prealloc = NULL;54}55 56/*57 * Free up additionally allocated memory for the ulist.58 *59 * @ulist: the ulist from which to free the additional memory60 *61 * This is useful in cases where the base 'struct ulist' has been statically62 * allocated.63 */64void ulist_release(struct ulist *ulist)65{66 struct ulist_node *node;67 struct ulist_node *next;68 69 list_for_each_entry_safe(node, next, &ulist->nodes, list) {70 kfree(node);71 }72 kfree(ulist->prealloc);73 ulist->prealloc = NULL;74 ulist->root = RB_ROOT;75 INIT_LIST_HEAD(&ulist->nodes);76}77 78/*79 * Prepare a ulist for reuse.80 *81 * @ulist: ulist to be reused82 *83 * Free up all additional memory allocated for the list elements and reinit84 * the ulist.85 */86void ulist_reinit(struct ulist *ulist)87{88 ulist_release(ulist);89 ulist_init(ulist);90}91 92/*93 * Dynamically allocate a ulist.94 *95 * @gfp_mask: allocation flags to for base allocation96 *97 * The allocated ulist will be returned in an initialized state.98 */99struct ulist *ulist_alloc(gfp_t gfp_mask)100{101 struct ulist *ulist = kmalloc(sizeof(*ulist), gfp_mask);102 103 if (!ulist)104 return NULL;105 106 ulist_init(ulist);107 108 return ulist;109}110 111void ulist_prealloc(struct ulist *ulist, gfp_t gfp_mask)112{113 if (!ulist->prealloc)114 ulist->prealloc = kzalloc(sizeof(*ulist->prealloc), gfp_mask);115}116 117/*118 * Free dynamically allocated ulist.119 *120 * @ulist: ulist to free121 *122 * It is not necessary to call ulist_release before.123 */124void ulist_free(struct ulist *ulist)125{126 if (!ulist)127 return;128 ulist_release(ulist);129 kfree(ulist);130}131 132static struct ulist_node *ulist_rbtree_search(struct ulist *ulist, u64 val)133{134 struct rb_node *n = ulist->root.rb_node;135 struct ulist_node *u = NULL;136 137 while (n) {138 u = rb_entry(n, struct ulist_node, rb_node);139 if (u->val < val)140 n = n->rb_right;141 else if (u->val > val)142 n = n->rb_left;143 else144 return u;145 }146 return NULL;147}148 149static void ulist_rbtree_erase(struct ulist *ulist, struct ulist_node *node)150{151 rb_erase(&node->rb_node, &ulist->root);152 list_del(&node->list);153 kfree(node);154 BUG_ON(ulist->nnodes == 0);155 ulist->nnodes--;156}157 158static int ulist_rbtree_insert(struct ulist *ulist, struct ulist_node *ins)159{160 struct rb_node **p = &ulist->root.rb_node;161 struct rb_node *parent = NULL;162 struct ulist_node *cur = NULL;163 164 while (*p) {165 parent = *p;166 cur = rb_entry(parent, struct ulist_node, rb_node);167 168 if (cur->val < ins->val)169 p = &(*p)->rb_right;170 else if (cur->val > ins->val)171 p = &(*p)->rb_left;172 else173 return -EEXIST;174 }175 rb_link_node(&ins->rb_node, parent, p);176 rb_insert_color(&ins->rb_node, &ulist->root);177 return 0;178}179 180/*181 * Add an element to the ulist.182 *183 * @ulist: ulist to add the element to184 * @val: value to add to ulist185 * @aux: auxiliary value to store along with val186 * @gfp_mask: flags to use for allocation187 *188 * Note: locking must be provided by the caller. In case of rwlocks write189 * locking is needed190 *191 * Add an element to a ulist. The @val will only be added if it doesn't192 * already exist. If it is added, the auxiliary value @aux is stored along with193 * it. In case @val already exists in the ulist, @aux is ignored, even if194 * it differs from the already stored value.195 *196 * ulist_add returns 0 if @val already exists in ulist and 1 if @val has been197 * inserted.198 * In case of allocation failure -ENOMEM is returned and the ulist stays199 * unaltered.200 */201int ulist_add(struct ulist *ulist, u64 val, u64 aux, gfp_t gfp_mask)202{203 return ulist_add_merge(ulist, val, aux, NULL, gfp_mask);204}205 206int ulist_add_merge(struct ulist *ulist, u64 val, u64 aux,207 u64 *old_aux, gfp_t gfp_mask)208{209 int ret;210 struct ulist_node *node;211 212 node = ulist_rbtree_search(ulist, val);213 if (node) {214 if (old_aux)215 *old_aux = node->aux;216 return 0;217 }218 219 if (ulist->prealloc) {220 node = ulist->prealloc;221 ulist->prealloc = NULL;222 } else {223 node = kmalloc(sizeof(*node), gfp_mask);224 if (!node)225 return -ENOMEM;226 }227 228 node->val = val;229 node->aux = aux;230 231 ret = ulist_rbtree_insert(ulist, node);232 ASSERT(!ret);233 list_add_tail(&node->list, &ulist->nodes);234 ulist->nnodes++;235 236 return 1;237}238 239/*240 * Delete one node from ulist.241 *242 * @ulist: ulist to remove node from243 * @val: value to delete244 * @aux: aux to delete245 *246 * The deletion will only be done when *BOTH* val and aux matches.247 * Return 0 for successful delete.248 * Return > 0 for not found.249 */250int ulist_del(struct ulist *ulist, u64 val, u64 aux)251{252 struct ulist_node *node;253 254 node = ulist_rbtree_search(ulist, val);255 /* Not found */256 if (!node)257 return 1;258 259 if (node->aux != aux)260 return 1;261 262 /* Found and delete */263 ulist_rbtree_erase(ulist, node);264 return 0;265}266 267/*268 * Iterate ulist.269 *270 * @ulist: ulist to iterate271 * @uiter: iterator variable, initialized with ULIST_ITER_INIT(&iterator)272 *273 * Note: locking must be provided by the caller. In case of rwlocks only read274 * locking is needed275 *276 * This function is used to iterate an ulist.277 * It returns the next element from the ulist or %NULL when the278 * end is reached. No guarantee is made with respect to the order in which279 * the elements are returned. They might neither be returned in order of280 * addition nor in ascending order.281 * It is allowed to call ulist_add during an enumeration. Newly added items282 * are guaranteed to show up in the running enumeration.283 */284struct ulist_node *ulist_next(const struct ulist *ulist, struct ulist_iterator *uiter)285{286 struct ulist_node *node;287 288 if (list_empty(&ulist->nodes))289 return NULL;290 if (uiter->cur_list && uiter->cur_list->next == &ulist->nodes)291 return NULL;292 if (uiter->cur_list) {293 uiter->cur_list = uiter->cur_list->next;294 } else {295 uiter->cur_list = ulist->nodes.next;296 }297 node = list_entry(uiter->cur_list, struct ulist_node, list);298 return node;299}300