1 // SPDX-License-Identifier: GPL-2.0 2 /* 3 * Copyright (C) 2011 STRATO AG 4 * 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 u64 13 * values. The only operations it supports is adding to the list and 14 * 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 without 18 * 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 the 32 * usage for tree enumeration in btrfs, where the logical addresses are 33 * 64 bit. 34 * 35 * It is also useful for tree enumeration which could be done elegantly 36 * recursively, but is not possible due to kernel stack limitations. The 37 * loop would be similar to the above. 38 */ 39 40 /* 41 * Freshly initialize a ulist. 42 * 43 * @ulist: the ulist to initialize 44 * 45 * Note: don't use this function to init an already used ulist, use 46 * ulist_reinit instead. 47 */ 48 void 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 memory 60 * 61 * This is useful in cases where the base 'struct ulist' has been statically 62 * allocated. 63 */ 64 void 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 reused 82 * 83 * Free up all additional memory allocated for the list elements and reinit 84 * the ulist. 85 */ 86 void 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 allocation 96 * 97 * The allocated ulist will be returned in an initialized state. 98 */ 99 struct 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 111 void 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 free 121 * 122 * It is not necessary to call ulist_release before. 123 */ 124 void ulist_free(struct ulist *ulist) 125 { 126 if (!ulist) 127 return; 128 ulist_release(ulist); 129 kfree(ulist); 130 } 131 132 static 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 else 144 return u; 145 } 146 return NULL; 147 } 148 149 static 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 158 static 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 else 173 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 to 184 * @val: value to add to ulist 185 * @aux: auxiliary value to store along with val 186 * @gfp_mask: flags to use for allocation 187 * 188 * Note: locking must be provided by the caller. In case of rwlocks write 189 * locking is needed 190 * 191 * Add an element to a ulist. The @val will only be added if it doesn't 192 * already exist. If it is added, the auxiliary value @aux is stored along with 193 * it. In case @val already exists in the ulist, @aux is ignored, even if 194 * it differs from the already stored value. 195 * 196 * ulist_add returns 0 if @val already exists in ulist and 1 if @val has been 197 * inserted. 198 * In case of allocation failure -ENOMEM is returned and the ulist stays 199 * unaltered. 200 */ 201 int 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 206 int 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 from 243 * @val: value to delete 244 * @aux: aux to delete 245 * 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 */ 250 int 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 iterate 271 * @uiter: iterator variable, initialized with ULIST_ITER_INIT(&iterator) 272 * 273 * Note: locking must be provided by the caller. In case of rwlocks only read 274 * locking is needed 275 * 276 * This function is used to iterate an ulist. 277 * It returns the next element from the ulist or %NULL when the 278 * end is reached. No guarantee is made with respect to the order in which 279 * the elements are returned. They might neither be returned in order of 280 * addition nor in ascending order. 281 * It is allowed to call ulist_add during an enumeration. Newly added items 282 * are guaranteed to show up in the running enumeration. 283 */ 284 struct 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