ulist.c 7.2 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278
  1. /*
  2. * Copyright (C) 2011 STRATO AG
  3. * written by Arne Jansen <sensille@gmx.net>
  4. * Distributed under the GNU GPL license version 2.
  5. */
  6. #include <linux/slab.h>
  7. #include <linux/export.h>
  8. #include "ulist.h"
  9. /*
  10. * ulist is a generic data structure to hold a collection of unique u64
  11. * values. The only operations it supports is adding to the list and
  12. * enumerating it.
  13. * It is possible to store an auxiliary value along with the key.
  14. *
  15. * The implementation is preliminary and can probably be sped up
  16. * significantly. A first step would be to store the values in an rbtree
  17. * as soon as ULIST_SIZE is exceeded.
  18. *
  19. * A sample usage for ulists is the enumeration of directed graphs without
  20. * visiting a node twice. The pseudo-code could look like this:
  21. *
  22. * ulist = ulist_alloc();
  23. * ulist_add(ulist, root);
  24. * ULIST_ITER_INIT(&uiter);
  25. *
  26. * while ((elem = ulist_next(ulist, &uiter)) {
  27. * for (all child nodes n in elem)
  28. * ulist_add(ulist, n);
  29. * do something useful with the node;
  30. * }
  31. * ulist_free(ulist);
  32. *
  33. * This assumes the graph nodes are adressable by u64. This stems from the
  34. * usage for tree enumeration in btrfs, where the logical addresses are
  35. * 64 bit.
  36. *
  37. * It is also useful for tree enumeration which could be done elegantly
  38. * recursively, but is not possible due to kernel stack limitations. The
  39. * loop would be similar to the above.
  40. */
  41. /**
  42. * ulist_init - freshly initialize a ulist
  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. ulist->nnodes = 0;
  51. ulist->nodes = ulist->int_nodes;
  52. ulist->nodes_alloced = ULIST_SIZE;
  53. ulist->root = RB_ROOT;
  54. }
  55. EXPORT_SYMBOL(ulist_init);
  56. /**
  57. * ulist_fini - free up additionally allocated memory for the ulist
  58. * @ulist: the ulist from which to free the additional memory
  59. *
  60. * This is useful in cases where the base 'struct ulist' has been statically
  61. * allocated.
  62. */
  63. void ulist_fini(struct ulist *ulist)
  64. {
  65. /*
  66. * The first ULIST_SIZE elements are stored inline in struct ulist.
  67. * Only if more elements are alocated they need to be freed.
  68. */
  69. if (ulist->nodes_alloced > ULIST_SIZE)
  70. kfree(ulist->nodes);
  71. ulist->nodes_alloced = 0; /* in case ulist_fini is called twice */
  72. ulist->root = RB_ROOT;
  73. }
  74. EXPORT_SYMBOL(ulist_fini);
  75. /**
  76. * ulist_reinit - prepare a ulist for reuse
  77. * @ulist: ulist to be reused
  78. *
  79. * Free up all additional memory allocated for the list elements and reinit
  80. * the ulist.
  81. */
  82. void ulist_reinit(struct ulist *ulist)
  83. {
  84. ulist_fini(ulist);
  85. ulist_init(ulist);
  86. }
  87. EXPORT_SYMBOL(ulist_reinit);
  88. /**
  89. * ulist_alloc - dynamically allocate a ulist
  90. * @gfp_mask: allocation flags to for base allocation
  91. *
  92. * The allocated ulist will be returned in an initialized state.
  93. */
  94. struct ulist *ulist_alloc(gfp_t gfp_mask)
  95. {
  96. struct ulist *ulist = kmalloc(sizeof(*ulist), gfp_mask);
  97. if (!ulist)
  98. return NULL;
  99. ulist_init(ulist);
  100. return ulist;
  101. }
  102. EXPORT_SYMBOL(ulist_alloc);
  103. /**
  104. * ulist_free - free dynamically allocated ulist
  105. * @ulist: ulist to free
  106. *
  107. * It is not necessary to call ulist_fini before.
  108. */
  109. void ulist_free(struct ulist *ulist)
  110. {
  111. if (!ulist)
  112. return;
  113. ulist_fini(ulist);
  114. kfree(ulist);
  115. }
  116. EXPORT_SYMBOL(ulist_free);
  117. static struct ulist_node *ulist_rbtree_search(struct ulist *ulist, u64 val)
  118. {
  119. struct rb_node *n = ulist->root.rb_node;
  120. struct ulist_node *u = NULL;
  121. while (n) {
  122. u = rb_entry(n, struct ulist_node, rb_node);
  123. if (u->val < val)
  124. n = n->rb_right;
  125. else if (u->val > val)
  126. n = n->rb_left;
  127. else
  128. return u;
  129. }
  130. return NULL;
  131. }
  132. static int ulist_rbtree_insert(struct ulist *ulist, struct ulist_node *ins)
  133. {
  134. struct rb_node **p = &ulist->root.rb_node;
  135. struct rb_node *parent = NULL;
  136. struct ulist_node *cur = NULL;
  137. while (*p) {
  138. parent = *p;
  139. cur = rb_entry(parent, struct ulist_node, rb_node);
  140. if (cur->val < ins->val)
  141. p = &(*p)->rb_right;
  142. else if (cur->val > ins->val)
  143. p = &(*p)->rb_left;
  144. else
  145. return -EEXIST;
  146. }
  147. rb_link_node(&ins->rb_node, parent, p);
  148. rb_insert_color(&ins->rb_node, &ulist->root);
  149. return 0;
  150. }
  151. /**
  152. * ulist_add - add an element to the ulist
  153. * @ulist: ulist to add the element to
  154. * @val: value to add to ulist
  155. * @aux: auxiliary value to store along with val
  156. * @gfp_mask: flags to use for allocation
  157. *
  158. * Note: locking must be provided by the caller. In case of rwlocks write
  159. * locking is needed
  160. *
  161. * Add an element to a ulist. The @val will only be added if it doesn't
  162. * already exist. If it is added, the auxiliary value @aux is stored along with
  163. * it. In case @val already exists in the ulist, @aux is ignored, even if
  164. * it differs from the already stored value.
  165. *
  166. * ulist_add returns 0 if @val already exists in ulist and 1 if @val has been
  167. * inserted.
  168. * In case of allocation failure -ENOMEM is returned and the ulist stays
  169. * unaltered.
  170. */
  171. int ulist_add(struct ulist *ulist, u64 val, u64 aux, gfp_t gfp_mask)
  172. {
  173. return ulist_add_merge(ulist, val, aux, NULL, gfp_mask);
  174. }
  175. int ulist_add_merge(struct ulist *ulist, u64 val, u64 aux,
  176. u64 *old_aux, gfp_t gfp_mask)
  177. {
  178. int ret = 0;
  179. struct ulist_node *node = NULL;
  180. node = ulist_rbtree_search(ulist, val);
  181. if (node) {
  182. if (old_aux)
  183. *old_aux = node->aux;
  184. return 0;
  185. }
  186. if (ulist->nnodes >= ulist->nodes_alloced) {
  187. u64 new_alloced = ulist->nodes_alloced + 128;
  188. struct ulist_node *new_nodes;
  189. void *old = NULL;
  190. int i;
  191. for (i = 0; i < ulist->nnodes; i++)
  192. rb_erase(&ulist->nodes[i].rb_node, &ulist->root);
  193. /*
  194. * if nodes_alloced == ULIST_SIZE no memory has been allocated
  195. * yet, so pass NULL to krealloc
  196. */
  197. if (ulist->nodes_alloced > ULIST_SIZE)
  198. old = ulist->nodes;
  199. new_nodes = krealloc(old, sizeof(*new_nodes) * new_alloced,
  200. gfp_mask);
  201. if (!new_nodes)
  202. return -ENOMEM;
  203. if (!old)
  204. memcpy(new_nodes, ulist->int_nodes,
  205. sizeof(ulist->int_nodes));
  206. ulist->nodes = new_nodes;
  207. ulist->nodes_alloced = new_alloced;
  208. /*
  209. * krealloc actually uses memcpy, which does not copy rb_node
  210. * pointers, so we have to do it ourselves. Otherwise we may
  211. * be bitten by crashes.
  212. */
  213. for (i = 0; i < ulist->nnodes; i++) {
  214. ret = ulist_rbtree_insert(ulist, &ulist->nodes[i]);
  215. if (ret < 0)
  216. return ret;
  217. }
  218. }
  219. ulist->nodes[ulist->nnodes].val = val;
  220. ulist->nodes[ulist->nnodes].aux = aux;
  221. ret = ulist_rbtree_insert(ulist, &ulist->nodes[ulist->nnodes]);
  222. BUG_ON(ret);
  223. ++ulist->nnodes;
  224. return 1;
  225. }
  226. EXPORT_SYMBOL(ulist_add);
  227. /**
  228. * ulist_next - iterate ulist
  229. * @ulist: ulist to iterate
  230. * @uiter: iterator variable, initialized with ULIST_ITER_INIT(&iterator)
  231. *
  232. * Note: locking must be provided by the caller. In case of rwlocks only read
  233. * locking is needed
  234. *
  235. * This function is used to iterate an ulist.
  236. * It returns the next element from the ulist or %NULL when the
  237. * end is reached. No guarantee is made with respect to the order in which
  238. * the elements are returned. They might neither be returned in order of
  239. * addition nor in ascending order.
  240. * It is allowed to call ulist_add during an enumeration. Newly added items
  241. * are guaranteed to show up in the running enumeration.
  242. */
  243. struct ulist_node *ulist_next(struct ulist *ulist, struct ulist_iterator *uiter)
  244. {
  245. if (ulist->nnodes == 0)
  246. return NULL;
  247. if (uiter->i < 0 || uiter->i >= ulist->nnodes)
  248. return NULL;
  249. return &ulist->nodes[uiter->i++];
  250. }
  251. EXPORT_SYMBOL(ulist_next);