avl.c 6.3 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414
  1. #include <u.h>
  2. #include <libc.h>
  3. #include <bio.h>
  4. #include <avl.h>
  5. /*
  6. * In-memory database stored as self-balancing AVL tree.
  7. * See Lewis & Denenberg, Data Structures and Their Algorithms.
  8. */
  9. static void
  10. singleleft(Avl **tp, Avl *p)
  11. {
  12. int l, r2;
  13. Avl *a, *c;
  14. a = *tp;
  15. c = a->n[1];
  16. r2 = c->bal;
  17. l = (r2 > 0? r2: 0)+1 - a->bal;
  18. if((a->n[1] = c->n[0]) != nil)
  19. a->n[1]->p = a;
  20. if((c->n[0] = a) != nil)
  21. c->n[0]->p = c;
  22. if((*tp = c) != nil)
  23. (*tp)->p = p;
  24. a->bal = -l;
  25. c->bal = r2 - ((l > 0? l: 0)+1);
  26. }
  27. static void
  28. singleright(Avl **tp, Avl *p)
  29. {
  30. int l2, r;
  31. Avl *a, *c;
  32. a = *tp;
  33. c = a->n[0];
  34. l2 = - c->bal;
  35. r = a->bal + ((l2 > 0? l2: 0)+1);
  36. if((a->n[0] = c->n[1]) != nil)
  37. a->n[0]->p = a;
  38. if((c->n[1] = a) != nil)
  39. c->n[1]->p = c;
  40. if((*tp = c) != nil)
  41. (*tp)->p = p;
  42. a->bal = r;
  43. c->bal = ((r > 0? r: 0)+1) - l2;
  44. }
  45. static void
  46. doublerightleft(Avl **tp, Avl *p)
  47. {
  48. singleright(&(*tp)->n[1], *tp);
  49. singleleft(tp, p);
  50. }
  51. static void
  52. doubleleftright(Avl **tp, Avl *p)
  53. {
  54. singleleft(&(*tp)->n[0], *tp);
  55. singleright(tp, p);
  56. }
  57. static void
  58. balance(Avl **tp, Avl *p)
  59. {
  60. switch((*tp)->bal){
  61. case -2:
  62. if((*tp)->n[0]->bal <= 0)
  63. singleright(tp, p);
  64. else if((*tp)->n[0]->bal == 1)
  65. doubleleftright(tp, p);
  66. else
  67. assert(0);
  68. break;
  69. case 2:
  70. if((*tp)->n[1]->bal >= 0)
  71. singleleft(tp, p);
  72. else if((*tp)->n[1]->bal == -1)
  73. doublerightleft(tp, p);
  74. else
  75. assert(0);
  76. break;
  77. }
  78. }
  79. static int
  80. _insertavl(Avl **tp, Avl *p, Avl *r, int (*cmp)(Avl*,Avl*), Avl **rfree)
  81. {
  82. int i, ob;
  83. if(*tp == nil){
  84. r->bal = 0;
  85. r->n[0] = nil;
  86. r->n[1] = nil;
  87. r->p = p;
  88. *tp = r;
  89. return 1;
  90. }
  91. ob = (*tp)->bal;
  92. if((i = cmp(r, *tp)) != 0){
  93. (*tp)->bal += i * _insertavl(&(*tp)->n[(i+1)/2], *tp, r, cmp,
  94. rfree);
  95. balance(tp, p);
  96. return ob == 0 && (*tp)->bal != 0;
  97. }
  98. /* install new entry */
  99. *rfree = *tp; /* save old node for freeing */
  100. *tp = r; /* insert new node */
  101. **tp = **rfree; /* copy old node's Avl contents */
  102. if(r->n[0]) /* fix node's children's parent pointers */
  103. r->n[0]->p = r;
  104. if(r->n[1])
  105. r->n[1]->p = r;
  106. return 0;
  107. }
  108. static Avl*
  109. _lookupavl(Avl *t, Avl *r, int (*cmp)(Avl*,Avl*))
  110. {
  111. int i;
  112. Avl *p;
  113. p = nil;
  114. while(t != nil){
  115. assert(t->p == p);
  116. if((i = cmp(r, t)) == 0)
  117. return t;
  118. p = t;
  119. t = t->n[(i+1)/2];
  120. }
  121. return nil;
  122. }
  123. static int
  124. successor(Avl **tp, Avl *p, Avl **r)
  125. {
  126. int ob;
  127. if((*tp)->n[0] == nil){
  128. *r = *tp;
  129. *tp = (*r)->n[1];
  130. if(*tp)
  131. (*tp)->p = p;
  132. return -1;
  133. }
  134. ob = (*tp)->bal;
  135. (*tp)->bal -= successor(&(*tp)->n[0], *tp, r);
  136. balance(tp, p);
  137. return -(ob != 0 && (*tp)->bal == 0);
  138. }
  139. static int
  140. _deleteavl(Avl **tp, Avl *p, Avl *rx, int(*cmp)(Avl*,Avl*), Avl **del,
  141. void (*predel)(Avl*, void*), void *arg)
  142. {
  143. int i, ob;
  144. Avl *r, *or;
  145. if(*tp == nil)
  146. return 0;
  147. ob = (*tp)->bal;
  148. if((i=cmp(rx, *tp)) != 0){
  149. (*tp)->bal += i * _deleteavl(&(*tp)->n[(i+1)/2], *tp, rx, cmp,
  150. del, predel, arg);
  151. balance(tp, p);
  152. return -(ob != 0 && (*tp)->bal == 0);
  153. }
  154. if(predel)
  155. (*predel)(*tp, arg);
  156. or = *tp;
  157. if(or->n[i=0] == nil || or->n[i=1] == nil){
  158. *tp = or->n[1-i];
  159. if(*tp)
  160. (*tp)->p = p;
  161. *del = or;
  162. return -1;
  163. }
  164. /* deleting node with two kids, find successor */
  165. or->bal += successor(&or->n[1], or, &r);
  166. r->bal = or->bal;
  167. r->n[0] = or->n[0];
  168. r->n[1] = or->n[1];
  169. *tp = r;
  170. (*tp)->p = p;
  171. /* node has changed; fix children's parent pointers */
  172. if(r->n[0])
  173. r->n[0]->p = r;
  174. if(r->n[1])
  175. r->n[1]->p = r;
  176. *del = or;
  177. balance(tp, p);
  178. return -(ob != 0 && (*tp)->bal == 0);
  179. }
  180. static void
  181. checkparents(Avl *a, Avl *p)
  182. {
  183. if(a == nil)
  184. return;
  185. if(a->p != p)
  186. print("bad parent\n");
  187. checkparents(a->n[0], a);
  188. checkparents(a->n[1], a);
  189. }
  190. struct Avltree
  191. {
  192. Avl *root;
  193. int (*cmp)(Avl*, Avl*);
  194. Avlwalk *walks;
  195. };
  196. struct Avlwalk
  197. {
  198. int started;
  199. int moved;
  200. Avlwalk *next;
  201. Avltree *tree;
  202. Avl *node;
  203. };
  204. Avltree*
  205. mkavltree(int (*cmp)(Avl*, Avl*))
  206. {
  207. Avltree *t;
  208. t = malloc(sizeof *t);
  209. if(t == nil)
  210. return nil;
  211. memset(t, 0, sizeof *t);
  212. t->cmp = cmp;
  213. return t;
  214. }
  215. void
  216. insertavl(Avltree *t, Avl *new, Avl **oldp)
  217. {
  218. *oldp = nil;
  219. _insertavl(&t->root, nil, new, t->cmp, oldp);
  220. }
  221. Avl*
  222. lookupavl(Avltree *t, Avl *key)
  223. {
  224. return _lookupavl(t->root, key, t->cmp);
  225. }
  226. static Avl*
  227. findpredecessor(Avl *a)
  228. {
  229. if(a == nil)
  230. return nil;
  231. if(a->n[0] != nil){
  232. /* predecessor is rightmost descendant of left child */
  233. for(a = a->n[0]; a->n[1]; a = a->n[1])
  234. ;
  235. return a;
  236. }else{
  237. /* we're at a leaf, successor is a parent we enter from the right */
  238. while(a->p && a->p->n[0] == a)
  239. a = a->p;
  240. return a->p;
  241. }
  242. }
  243. static Avl*
  244. findsuccessor(Avl *a)
  245. {
  246. if(a == nil)
  247. return nil;
  248. if(a->n[1] != nil){
  249. /* successor is leftmost descendant of right child */
  250. for(a = a->n[1]; a->n[0]; a = a->n[0])
  251. ;
  252. return a;
  253. }else{
  254. /* we're at a leaf, successor is a parent we enter from the left going up */
  255. while(a->p && a->p->n[1] == a)
  256. a = a->p;
  257. return a->p;
  258. }
  259. }
  260. static void
  261. walkdel(Avl *a, void *v)
  262. {
  263. Avl *p;
  264. Avlwalk *w;
  265. Avltree *t;
  266. if(a == nil)
  267. return;
  268. p = findpredecessor(a);
  269. t = v;
  270. for(w = t->walks; w; w = w->next){
  271. if(w->node == a){
  272. /* back pointer to predecessor; not perfect but adequate */
  273. w->moved = 1;
  274. w->node = p;
  275. if(p == nil)
  276. w->started = 0;
  277. }
  278. }
  279. }
  280. void
  281. deleteavl(Avltree *t, Avl *key, Avl **oldp)
  282. {
  283. *oldp = nil;
  284. _deleteavl(&t->root, nil, key, t->cmp, oldp, walkdel, t);
  285. }
  286. Avlwalk*
  287. avlwalk(Avltree *t)
  288. {
  289. Avlwalk *w;
  290. w = malloc(sizeof *w);
  291. if(w == nil)
  292. return nil;
  293. memset(w, 0, sizeof *w);
  294. w->tree = t;
  295. w->next = t->walks;
  296. t->walks = w;
  297. return w;
  298. }
  299. Avl*
  300. avlnext(Avlwalk *w)
  301. {
  302. Avl *a;
  303. if(w->started==0){
  304. for(a = w->tree->root; a && a->n[0]; a = a->n[0])
  305. ;
  306. w->node = a;
  307. w->started = 1;
  308. }else{
  309. a = findsuccessor(w->node);
  310. if(a == w->node)
  311. abort();
  312. w->node = a;
  313. }
  314. return w->node;
  315. }
  316. Avl*
  317. avlprev(Avlwalk *w)
  318. {
  319. Avl *a;
  320. if(w->started == 0){
  321. for(a = w->tree->root; a && a->n[1]; a = a->n[1])
  322. ;
  323. w->node = a;
  324. w->started = 1;
  325. }else if(w->moved){
  326. w->moved = 0;
  327. return w->node;
  328. }else{
  329. a = findpredecessor(w->node);
  330. if(a == w->node)
  331. abort();
  332. w->node = a;
  333. }
  334. return w->node;
  335. }
  336. void
  337. endwalk(Avlwalk *w)
  338. {
  339. Avltree *t;
  340. Avlwalk **l;
  341. t = w->tree;
  342. for(l = &t->walks; *l; l = &(*l)->next){
  343. if(*l == w){
  344. *l = w->next;
  345. break;
  346. }
  347. }
  348. free(w);
  349. }
  350. static void
  351. walkavl(Avl *t, void (*f)(Avl*, void*), void *v)
  352. {
  353. if(t == nil)
  354. return;
  355. walkavl(t->n[0], f, v);
  356. f(t, v);
  357. walkavl(t->n[1], f, v);
  358. }