~mcf/samurai

samurai/tree.c -rw-r--r-- 2.4 KiB
4fac369eAman Verma manual: Use Pq macro to parenthesize link. 20 days ago
                                                                                
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
/* Based on musl's src/search/tsearch.c, by Szabolcs Nagy.
 * See LICENSE file for copyright details. */
#include <stdlib.h>
#include <string.h>
#include "tree.h"
#include "util.h"

#define MAXH (sizeof(void *) * 8 * 3 / 2)

void
deltree(struct treenode *n, void delkey(void *), void delval(void *))
{
	if (!n)
		return;
	if (delkey)
		delkey(n->key);
	if (delval)
		delval(n->value);
	deltree(n->child[0], delkey, delval);
	deltree(n->child[1], delkey, delval);
	free(n);
}

static inline int
height(struct treenode *n)
{
	return n ? n->height : 0;
}

static int
rot(struct treenode **p, struct treenode *x, int dir /* deeper side */)
{
	struct treenode *y = x->child[dir];
	struct treenode *z = y->child[!dir];
	int hx = x->height;
	int hz = height(z);

	if (hz > height(y->child[dir])) {
		/*
		 *   x
		 *  / \ dir          z
		 * A   y            / \
		 *    / \   -->    x   y
		 *   z   D        /|   |\
		 *  / \          A B   C D
		 * B   C
		 */
		x->child[dir] = z->child[!dir];
		y->child[!dir] = z->child[dir];
		z->child[!dir] = x;
		z->child[dir] = y;
		x->height = hz;
		y->height = hz;
		z->height = hz + 1;
	} else {
		/*
		 *   x               y
		 *  / \             / \
		 * A   y    -->    x   D
		 *    / \         / \
		 *   z   D       A   z
		 */
		x->child[dir] = z;
		y->child[!dir] = x;
		x->height = hz + 1;
		y->height = hz + 2;
		z = y;
	}
	*p = z;
	return z->height - hx;
}

static int
balance(struct treenode **p)
{
	struct treenode *n = *p;
	int h0 = height(n->child[0]);
	int h1 = height(n->child[1]);

	if (h0 - h1 + 1u < 3u) {
		int old = n->height;
		n->height = h0 < h1 ? h1 + 1 : h0 + 1;
		return n->height - old;
	}
	return rot(p, n, h0 < h1);
}

struct treenode *
treefind(struct treenode *n, const char *key)
{
	int c;

	while (n) {
		c = strcmp(key, n->key);
		if (c == 0)
			return n;
		n = n->child[c > 0];
	}
	return NULL;
}

void *
treeinsert(struct treenode **rootp, char *key, void *value)
{
	struct treenode **a[MAXH], *n = *rootp, *r;
	void *old;
	int i = 0, c;

	a[i++] = rootp;
	while (n) {
		c = strcmp(key, n->key);
		if (c == 0) {
			old = n->value;
			n->value = value;
			return old;
		}
		a[i++] = &n->child[c > 0];
		n = n->child[c > 0];
	}
	r = xmalloc(sizeof(*r));
	r->key = key;
	r->value = value;
	r->child[0] = r->child[1] = NULL;
	r->height = 1;
	/* insert new node, rebalance ancestors.  */
	*a[--i] = r;
	while (i && balance(a[--i]))
		;
	return NULL;
}