Mercurial > hg > ucx / changeset

--- a/ucx/avl.c	Sun May 17 18:32:41 2015 +0200
+++ b/ucx/avl.c	Mon May 18 12:54:18 2015 +0200
@@ -28,6 +28,80 @@
 
 #include "avl.h"
 
+#define ptrcast(ptr) ((void*)(ptr))
+
+static void ucx_avl_connect(UcxAVLTree *tree,
+        UcxAVLNode *node, UcxAVLNode *child, intptr_t nullkey) {
+    if (child) {
+        child->parent = node;
+    }
+    // if child is NULL, nullkey decides if left or right pointer is cleared
+    if (tree->cmpfunc(
+        ptrcast(child ? child->key : nullkey),
+        ptrcast(node->key), tree->userdata) > 0) {
+      node->right = child;
+    } else {
+      node->left = child;
+    }
+    size_t lh = node->left ? node->left->height : 0;
+    size_t rh = node->right ? node->right->height : 0;
+    node->height = 1 + (lh > rh ? lh : rh);
+}
+
+#define avlheight(node) ((node) ? (node)->height : 0)
+
+static UcxAVLNode* avl_rotright(UcxAVLTree *tree, UcxAVLNode *l0) {
+    UcxAVLNode *p = l0->parent;
+    UcxAVLNode *l1 = l0->left;
+    if (p) {
+        ucx_avl_connect(tree, p, l1, 0);
+    } else {
+        l1->parent = NULL;
+    }
+    ucx_avl_connect(tree, l0, l1->right, l1->key);
+    ucx_avl_connect(tree, l1, l0, 0);
+    return l1;
+}
+
+static UcxAVLNode* avl_rotleft(UcxAVLTree *tree, UcxAVLNode *l0) {
+    UcxAVLNode *p = l0->parent;
+    UcxAVLNode *l1 = l0->right;
+    if (p) {
+        ucx_avl_connect(tree, p, l1, 0);
+    } else {
+        l1->parent = NULL;
+    }
+    ucx_avl_connect(tree, l0, l1->left, l1->key);
+    ucx_avl_connect(tree, l1, l0, 0);
+    return l1;
+}
+
+static void ucx_avl_balance(UcxAVLTree *tree, UcxAVLNode *n) {
+    int lh = avlheight(n->left);
+    int rh = avlheight(n->right);
+    n->height = 1 + (lh > rh ? lh : rh);
+    
+    if (lh - rh == 2) {
+      UcxAVLNode *c = n->left;
+      if (avlheight(c->right) - avlheight(c->left) == 1) {
+        avl_rotleft(tree, c);
+      }
+      n = avl_rotright(tree, n);
+    } else if (rh - lh == 2) {  
+      UcxAVLNode *c = n->right;
+      if (avlheight(c->left) - avlheight(c->right) == 1) {
+        avl_rotright(tree, c);
+      }
+      n = avl_rotleft(tree, n);
+    }
+
+    if (n->parent) {
+      ucx_avl_balance(tree, n->parent);
+    } else {
+      tree->root = n;
+    }
+}
+
 UcxAVLTree *ucx_avl_new(cmp_func cmpfunc) {
     return ucx_avl_new_a(cmpfunc, ucx_default_allocator());
 }
@@ -45,14 +119,83 @@
 }
 
 void *ucx_avl_get(UcxAVLTree *tree, intptr_t key) {
-    return NULL;
+    UcxAVLNode *n = tree->root;
+    int cmpresult;
+    while (n && (cmpresult = tree->cmpfunc(
+            ptrcast(key), ptrcast(n->key), tree->userdata))) {
+        n = cmpresult > 0 ? n->right : n->left;
+    }
+    return n ? n->value : NULL;
 }
 
 void* ucx_avl_put(UcxAVLTree *tree, intptr_t key, void *value) {
-    return NULL;
+    if (tree->root) {
+        UcxAVLNode *n = tree->root;
+        int cmpresult;
+        while (cmpresult = tree->cmpfunc(
+                ptrcast(key), ptrcast(n->key), tree->userdata)) {
+            UcxAVLNode *m = cmpresult > 0 ? n->right : n->left;
+            if (m) {
+                n = m;
+            } else {
+                break;
+            }
+        }
+
+        if (cmpresult) {
+            UcxAVLNode *e = malloc(sizeof(UcxAVLNode));
+            e->key = key; e->value = value; e->height = 1;
+            e->parent = e->left = e->right = NULL;
+            ucx_avl_connect(tree, n, e, 0);
+            ucx_avl_balance(tree, n);
+            return NULL;
+        } else {
+            void *prevval = n->value;
+            n->value = value;
+            return prevval;
+        }
+    } else {
+        tree->root = malloc(sizeof(UcxAVLNode));
+        tree->root->key = key; tree->root->value = value;
+        tree->root->height = 1;
+        tree->root->parent = tree->root->left = tree->root->right = NULL;
+        return NULL;
+    }
 }
 
 void* ucx_avl_remove(UcxAVLTree *tree, intptr_t key) {
-    return NULL;
+    UcxAVLNode *n = tree->root;
+    int cmpresult;
+    while (n && (cmpresult = tree->cmpfunc(
+            ptrcast(key), ptrcast(n->key), tree->userdata))) {
+        n = cmpresult > 0 ? n->right : n->left;
+    }
+    if (n) {
+        void *prevval = n->value;
+        UcxAVLNode *p = n->parent;
+        if (n->left && n->right) {
+            UcxAVLNode *s = n->right;
+            while (s->left) {
+                s = s->left;
+            }
+            ucx_avl_connect(tree, s->parent, s->right, s->key);
+            n->key = s->key; n->value = s->value;
+            p = s->parent;
+        } else {
+            if (p) {
+                ucx_avl_connect(tree, p, n->right ? n->right:n->left, n->key);
+            } else {
+                tree->root = n->right ? n->right : n->left;
+            }
+        }
+        // TODO: free the correct memory
+
+        if (p) {
+            ucx_avl_balance(tree, p);
+        }
+        return prevval;
+    } else {
+        return NULL;
+    }
 }