1 files changed, 51 insertions, 9 deletions

diff --git a/fs/btrfs/ulist.c b/fs/btrfs/ulist.c
index 99be4c138db6..7b417e20efe2 100644
--- a/fs/btrfs/ulist.c
+++ b/fs/btrfs/ulist.c
@@ -5,7 +5,7 @@
  */
 
 #include <linux/slab.h>
-#include <linux/module.h>
+#include <linux/export.h>
 #include "ulist.h"
 
 /*
@@ -53,6 +53,7 @@ void ulist_init(struct ulist *ulist)
         ulist->nnodes = 0;
         ulist->nodes = ulist->int_nodes;
         ulist->nodes_alloced = ULIST_SIZE;
+        ulist->root = RB_ROOT;
 }
 EXPORT_SYMBOL(ulist_init);
 
@@ -72,6 +73,7 @@ void ulist_fini(struct ulist *ulist)
         if (ulist->nodes_alloced > ULIST_SIZE)
                 kfree(ulist->nodes);
         ulist->nodes_alloced = 0;       /* in case ulist_fini is called twice */
+        ulist->root = RB_ROOT;
 }
 EXPORT_SYMBOL(ulist_fini);
 
@@ -123,6 +125,45 @@ void ulist_free(struct ulist *ulist)
 }
 EXPORT_SYMBOL(ulist_free);
 
+static struct ulist_node *ulist_rbtree_search(struct ulist *ulist, u64 val)
+{
+        struct rb_node *n = ulist->root.rb_node;
+        struct ulist_node *u = NULL;
+
+        while (n) {
+                u = rb_entry(n, struct ulist_node, rb_node);
+                if (u->val < val)
+                        n = n->rb_right;
+                else if (u->val > val)
+                        n = n->rb_left;
+                else
+                        return u;
+        }
+        return NULL;
+}
+
+static int ulist_rbtree_insert(struct ulist *ulist, struct ulist_node *ins)
+{
+        struct rb_node **p = &ulist->root.rb_node;
+        struct rb_node *parent = NULL;
+        struct ulist_node *cur = NULL;
+
+        while (*p) {
+                parent = *p;
+                cur = rb_entry(parent, struct ulist_node, rb_node);
+
+                if (cur->val < ins->val)
+                        p = &(*p)->rb_right;
+                else if (cur->val > ins->val)
+                        p = &(*p)->rb_left;
+                else
+                        return -EEXIST;
+        }
+        rb_link_node(&ins->rb_node, parent, p);
+        rb_insert_color(&ins->rb_node, &ulist->root);
+        return 0;
+}
+
 /**
  * ulist_add - add an element to the ulist
  * @ulist:      ulist to add the element to
@@ -151,14 +192,13 @@ int ulist_add(struct ulist *ulist, u64 val, u64 aux, gfp_t gfp_mask)
 int ulist_add_merge(struct ulist *ulist, u64 val, u64 aux,
                     u64 *old_aux, gfp_t gfp_mask)
 {
-        int i;
-
-        for (i = 0; i < ulist->nnodes; ++i) {
-                if (ulist->nodes[i].val == val) {
-                        if (old_aux)
-                                *old_aux = ulist->nodes[i].aux;
-                        return 0;
-                }
+        int ret = 0;
+        struct ulist_node *node = NULL;
+        node = ulist_rbtree_search(ulist, val);
+        if (node) {
+                if (old_aux)
+                        *old_aux = node->aux;
+                return 0;
         }
 
         if (ulist->nnodes >= ulist->nodes_alloced) {
@@ -187,6 +227,8 @@ int ulist_add_merge(struct ulist *ulist, u64 val, u64 aux,
         }
         ulist->nodes[ulist->nnodes].val = val;
         ulist->nodes[ulist->nnodes].aux = aux;
+        ret = ulist_rbtree_insert(ulist, &ulist->nodes[ulist->nnodes]);
+        BUG_ON(ret);
         ++ulist->nnodes;
 
         return 1;