improve enumerate so arbitrary ranges can be enumerated quickly
[ffmpeg.git] / libavutil / tree.c
1 /*
2  * copyright (c) 2006 Michael Niedermayer <michaelni@gmx.at>
3  *
4  * This file is part of FFmpeg.
5  *
6  * FFmpeg is free software; you can redistribute it and/or
7  * modify it under the terms of the GNU Lesser General Public
8  * License as published by the Free Software Foundation; either
9  * version 2.1 of the License, or (at your option) any later version.
10  *
11  * FFmpeg is distributed in the hope that it will be useful,
12  * but WITHOUT ANY WARRANTY; without even the implied warranty of
13  * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
14  * Lesser General Public License for more details.
15  *
16  * You should have received a copy of the GNU Lesser General Public
17  * License along with FFmpeg; if not, write to the Free Software
18  * Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301 USA
19  */
20
21 #include "common.h"
22 #include "log.h"
23 #include "tree.h"
24
25 typedef struct AVTreeNode{
26     struct AVTreeNode *child[2];
27     void *elem;
28     int state;
29 }AVTreeNode;
30
31 void *av_tree_find(const AVTreeNode *t, void *key, int (*cmp)(void *key, const void *b), void *next[2]){
32     if(t){
33         unsigned int v= cmp(t->elem, key);
34         if(v){
35             if(next) next[(v>>31)^1]= t->elem;
36             return av_tree_find(t->child[v>>31], key, cmp, next);
37         }else{
38             return t->elem;
39         }
40     }
41     return NULL;
42 }
43
44 void *av_tree_insert(AVTreeNode **tp, void *key, int (*cmp)(void *key, const void *b)){
45     AVTreeNode *t= *tp;
46     if(t){
47         unsigned int v= cmp(t->elem, key);
48         if(v){
49             int i= v>>31;
50             AVTreeNode **child= &t->child[i];
51             void *ret= av_tree_insert(child, key, cmp);
52             if(!ret){
53                 t->state -= ((int)v>>31)|1;
54                 if(!(t->state&1)){
55                     if(t->state){
56                         if((*child)->state*2 == t->state){
57                             *tp= *child;
58                             *child= (*child)->child[i^1];
59                             (*tp)->child[i^1]= t;
60                             t->state= 0;
61                         }else{
62                             *tp= (*child)->child[i^1];
63                             (*child)->child[i^1]= (*tp)->child[i];
64                             (*tp)->child[i]= *child;
65                             *child= (*tp)->child[i^1];
66                             (*tp)->child[i^1]= t;
67
68                             i= (*tp)->state > 0;
69                             (*tp)->child[i  ]->state= 0;
70                             (*tp)->child[i^1]->state= -(*tp)->state;
71                         }
72                         (*tp)->state=0;
73                     }
74                     return key;
75                 }
76             }
77             return ret;
78         }else{
79             return t->elem;
80         }
81     }else{
82         *tp= av_mallocz(sizeof(AVTreeNode));
83         (*tp)->elem= key;
84         return NULL;
85     }
86 }
87
88 void av_tree_destroy(AVTreeNode *t){
89     av_tree_destroy(t->child[0]);
90     av_tree_destroy(t->child[1]);
91     av_free(t);
92 }
93
94 #if 0
95 void av_tree_enumerate(AVTreeNode *t, void *opaque, int (*f)(void *opaque, void *elem)){
96     int v= f(opaque, t->elem);
97     if(v>=0) av_tree_enumerate(t->child[0], opaque, f);
98     if(v<=0) av_tree_enumerate(t->child[1], opaque, f);
99 }
100 #endif
101
102 #ifdef TEST
103
104 static int check(AVTreeNode *t){
105     if(t){
106         int left= check(t->child[0]);
107         int right= check(t->child[1]);
108
109         if(left>999 || right>999)
110             return 1000;
111         if(right - left != t->state)
112             return 1000;
113         if(t->state>1 || t->state<-1)
114             return 1000;
115         return FFMAX(left, right)+1;
116     }
117     return 0;
118 }
119
120 static void print(AVTreeNode *t, int depth){
121     int i;
122     for(i=0; i<depth*4; i++) av_log(NULL, AV_LOG_ERROR, " ");
123     if(t){
124         av_log(NULL, AV_LOG_ERROR, "Node %p %2d %4d\n", t, t->state, t->elem);
125         print(t->child[0], depth+1);
126         print(t->child[1], depth+1);
127     }else
128         av_log(NULL, AV_LOG_ERROR, "NULL\n");
129 }
130
131 int cmp(const void *a, const void *b){
132     return a-b;
133 }
134
135 int main(){
136     int i,j,k;
137     AVTreeNode *root= av_mallocz(sizeof(AVTreeNode));
138
139     for(i=0; i<10000; i++){
140         int j= (random()%863294);
141         if(check(root) > 999){
142             av_log(NULL, AV_LOG_ERROR, "FATAL error %d\n", i);
143         print(root, 0);
144             return -1;
145         }
146         av_log(NULL, AV_LOG_ERROR, "inserting %4d\n", j);
147         av_tree_insert(&root, (void*)(j+1), cmp);
148     }
149     return 0;
150 }
151 #endif