Rax updated to latest antirez/rax commit.

2 files changed, 52 insertions, 25 deletions

diff --git a/src/rax.c b/src/rax.c
index 78ead7a27..4e45fd2c3 100644
--- a/src/rax.c
+++ b/src/rax.c
@@ -131,7 +131,7 @@ static inline void raxStackFree(raxStack *ts) {
 }
 
 /* ----------------------------------------------------------------------------
- * Radis tree implementation
+ * Radix tree implementation
  * --------------------------------------------------------------------------*/
 
 /* Allocate a new non compressed node with the specified number of children.
@@ -186,10 +186,10 @@ raxNode *raxReallocForData(raxNode *n, void *data) {
 void raxSetData(raxNode *n, void *data) {
     n->iskey = 1;
     if (data != NULL) {
+        n->isnull = 0;
         void **ndata = (void**)
             ((char*)n+raxNodeCurrentLength(n)-sizeof(void*));
         memcpy(ndata,&data,sizeof(data));
-        n->isnull = 0;
     } else {
         n->isnull = 1;
     }
@@ -396,6 +396,7 @@ static inline size_t raxLowWalk(rax *rax, unsigned char *s, size_t len, raxNode
                   position to 0 to signal this node represents
                   the searched key. */
     }
+    debugnode("Lookup stop node is",h);
     if (stopnode) *stopnode = h;
     if (plink) *plink = parentlink;
     if (splitpos && h->iscompr) *splitpos = j;
@@ -424,18 +425,21 @@ int raxInsert(rax *rax, unsigned char *s, size_t len, void *data, void **old) {
      * our key. We have just to reallocate the node and make space for the
      * data pointer. */
     if (i == len && (!h->iscompr || j == 0 /* not in the middle if j is 0 */)) {
+        debugf("### Insert: node representing key exists\n");
+        if (!h->iskey || h->isnull) {
+            h = raxReallocForData(h,data);
+            if (h) memcpy(parentlink,&h,sizeof(h));
+        }
+        if (h == NULL) {
+            errno = ENOMEM;
+            return 0;
+        }
         if (h->iskey) {
             if (old) *old = raxGetData(h);
             raxSetData(h,data);
             errno = 0;
             return 0; /* Element already exists. */
         }
-        h = raxReallocForData(h,data);
-        if (h == NULL) {
-            errno = ENOMEM;
-            return 0;
-        }
-        memcpy(parentlink,&h,sizeof(h));
         raxSetData(h,data);
         rax->numele++;
         return 1; /* Element inserted. */
@@ -734,9 +738,7 @@ int raxInsert(rax *rax, unsigned char *s, size_t len, void *data, void **old) {
     }
 
     /* We walked the radix tree as far as we could, but still there are left
-     * chars in our string. We need to insert the missing nodes.
-     * Note: while loop never entered if the node was split by ALGO2,
-     * since i == len. */
+     * chars in our string. We need to insert the missing nodes. */
     while(i < len) {
         raxNode *child;
 
@@ -871,7 +873,8 @@ raxNode *raxRemoveChild(raxNode *parent, raxNode *child) {
     memmove(((char*)cp)-1,cp,(parent->size-taillen-1)*sizeof(raxNode**));
 
     /* Move the remaining "tail" pointer at the right position as well. */
-    memmove(((char*)c)-1,c+1,taillen*sizeof(raxNode**)+parent->iskey*sizeof(void*));
+    size_t valuelen = (parent->iskey && !parent->isnull) ? sizeof(void*) : 0;
+    memmove(((char*)c)-1,c+1,taillen*sizeof(raxNode**)+valuelen);
 
     /* 4. Update size. */
     parent->size--;
@@ -1090,27 +1093,36 @@ int raxRemove(rax *rax, unsigned char *s, size_t len, void **old) {
 
 /* This is the core of raxFree(): performs a depth-first scan of the
  * tree and releases all the nodes found. */
-void raxRecursiveFree(rax *rax, raxNode *n) {
+void raxRecursiveFree(rax *rax, raxNode *n, void (*free_callback)(void*)) {
+    debugnode("free traversing",n);
     int numchildren = n->iscompr ? 1 : n->size;
     raxNode **cp = raxNodeLastChildPtr(n);
     while(numchildren--) {
         raxNode *child;
         memcpy(&child,cp,sizeof(child));
-        raxRecursiveFree(rax,child);
+        raxRecursiveFree(rax,child,free_callback);
         cp--;
     }
     debugnode("free depth-first",n);
+    if (free_callback && n->iskey && !n->isnull)
+        free_callback(raxGetData(n));
     rax_free(n);
     rax->numnodes--;
 }
 
-/* Free a whole radix tree. */
-void raxFree(rax *rax) {
-    raxRecursiveFree(rax,rax->head);
+/* Free a whole radix tree, calling the specified callback in order to
+ * free the auxiliary data. */
+void raxFreeWithCallback(rax *rax, void (*free_callback)(void*)) {
+    raxRecursiveFree(rax,rax->head,free_callback);
     assert(rax->numnodes == 0);
     rax_free(rax);
 }
 
+/* Free a whole radix tree. */
+void raxFree(rax *rax) {
+    raxFreeWithCallback(rax,NULL);
+}
+
 /* ------------------------------- Iterator --------------------------------- */
 
 /* Initialize a Rax iterator. This call should be performed a single time
@@ -1172,7 +1184,7 @@ void raxIteratorDelChars(raxIterator *it, size_t count) {
  * The function returns 1 on success or 0 on out of memory. */
 int raxIteratorNextStep(raxIterator *it, int noup) {
     if (it->flags & RAX_ITER_EOF) {
-        return 0;
+        return 1;
     } else if (it->flags & RAX_ITER_JUST_SEEKED) {
         it->flags &= ~RAX_ITER_JUST_SEEKED;
         return 1;
@@ -1184,10 +1196,6 @@ int raxIteratorNextStep(raxIterator *it, int noup) {
     size_t orig_stack_items = it->stack.items;
     raxNode *orig_node = it->node;
 
-    /* Clear the EOF flag: it will be set again if the EOF condition
-     * is still valid. */
-    it->flags &= ~RAX_ITER_EOF;
-
     while(1) {
         int children = it->node->iscompr ? 1 : it->node->size;
         if (!noup && children) {
@@ -1288,7 +1296,7 @@ int raxSeekGreatest(raxIterator *it) {
  * effect to the one of raxIteratorPrevSte(). */
 int raxIteratorPrevStep(raxIterator *it, int noup) {
     if (it->flags & RAX_ITER_EOF) {
-        return 0;
+        return 1;
     } else if (it->flags & RAX_ITER_JUST_SEEKED) {
         it->flags &= ~RAX_ITER_JUST_SEEKED;
         return 1;
@@ -1409,6 +1417,7 @@ int raxSeek(raxIterator *it, const char *op, unsigned char *ele, size_t len) {
         it->node = it->rt->head;
         if (!raxSeekGreatest(it)) return 0;
         assert(it->node->iskey);
+        it->data = raxGetData(it->node);
         return 1;
     }
 
@@ -1427,6 +1436,7 @@ int raxSeek(raxIterator *it, const char *op, unsigned char *ele, size_t len) {
         /* We found our node, since the key matches and we have an
          * "equal" condition. */
         if (!raxIteratorAddChars(it,ele,len)) return 0; /* OOM. */
+        it->data = raxGetData(it->node);
     } else if (lt || gt) {
         /* Exact key not found or eq flag not set. We have to set as current
          * key the one represented by the node we stopped at, and perform
@@ -1499,6 +1509,7 @@ int raxSeek(raxIterator *it, const char *op, unsigned char *ele, size_t len) {
                  * the previous sub-tree. */
                 if (nodechar < keychar) {
                     if (!raxSeekGreatest(it)) return 0;
+                    it->data = raxGetData(it->node);
                 } else {
                     if (!raxIteratorAddChars(it,it->node->data,it->node->size))
                         return 0;
@@ -1615,8 +1626,8 @@ int raxCompare(raxIterator *iter, const char *op, unsigned char *key, size_t key
     int eq = 0, lt = 0, gt = 0;
 
     if (op[0] == '=' || op[1] == '=') eq = 1;
-    if (op[1] == '>') gt = 1;
-    else if (op[1] == '<') lt = 1;
+    if (op[0] == '>') gt = 1;
+    else if (op[0] == '<') lt = 1;
     else if (op[1] != '=') return 0; /* Syntax error. */
 
     size_t minlen = key_len < iter->key_len ? key_len : iter->key_len;
@@ -1644,6 +1655,19 @@ void raxStop(raxIterator *it) {
     raxStackFree(&it->stack);
 }
 
+/* Return if the iterator is in an EOF state. This happens when raxSeek()
+ * failed to seek an appropriate element, so that raxNext() or raxPrev()
+ * will return zero, or when an EOF condition was reached while iterating
+ * with raxNext() and raxPrev(). */
+int raxEOF(raxIterator *it) {
+    return it->flags & RAX_ITER_EOF;
+}
+
+/* Return the number of elements inside the radix tree. */
+uint64_t raxSize(rax *rax) {
+    return rax->numele;
+}
+
 /* ----------------------------- Introspection ------------------------------ */
 
 /* This function is mostly used for debugging and learning purposes.
diff --git a/src/rax.h b/src/rax.h
index 6f91f4c1b..b4e2fd91e 100644
--- a/src/rax.h
+++ b/src/rax.h
@@ -148,6 +148,7 @@ int raxInsert(rax *rax, unsigned char *s, size_t len, void *data, void **old);
 int raxRemove(rax *rax, unsigned char *s, size_t len, void **old);
 void *raxFind(rax *rax, unsigned char *s, size_t len);
 void raxFree(rax *rax);
+void raxFreeWithCallback(rax *rax, void (*free_callback)(void*));
 void raxStart(raxIterator *it, rax *rt);
 int raxSeek(raxIterator *it, const char *op, unsigned char *ele, size_t len);
 int raxNext(raxIterator *it);
@@ -155,6 +156,8 @@ int raxPrev(raxIterator *it);
 int raxRandomWalk(raxIterator *it, size_t steps);
 int raxCompare(raxIterator *iter, const char *op, unsigned char *key, size_t key_len);
 void raxStop(raxIterator *it);
+int raxEOF(raxIterator *it);
 void raxShow(rax *rax);
+uint64_t raxSize(rax *rax);
 
 #endif