1 // Copyright 2009 The Archiveopteryx Developers <info@aox.org>
2 
3 #ifndef PATRICIATREE_H
4 #define PATRICIATREE_H
5 
6 #include "global.h"
7 #include "allocator.h"
8 
9 
10 template< class T >
11 class PatriciaTree
12     : public Garbage
13 {
14 public:
PatriciaTree()15     PatriciaTree(): root( 0 ) { }
~PatriciaTree()16     virtual ~PatriciaTree() {}
17 
18     class Node
19         : public Garbage
20     {
21     public:
Node()22         Node() : zero( 0 ), one( 0 ),
23                  parent( 0 ),
24                  data( 0 ), length( 0 ) {
25             setFirstNonPointer( &length );
26         }
27 
count()28         uint count() {
29             uint c = 0;
30             if ( data )
31                 c = 1;
32             if ( zero )
33                 c += zero->count();
34             if ( one )
35                 c += one->count();
36             return c;
37         }
38 
new(size_t ownSize,uint extra)39         void * operator new( size_t ownSize, uint extra ) {
40             return Allocator::alloc( ownSize + extra );
41         }
42 
43     public: // really want private, but g++ 4.0.3 throws a fit
44         friend class PatriciaTree;
45         Node * zero;
46         Node * one;
47         Node * parent;
48         T * data;
49         uint length;
50         char key[1];
51     };
52 
find(const char * k,uint l)53     T * find( const char * k, uint l ) const {
54         Node * n = locate( k, l );
55         if ( n )
56             return n->data;
57         return 0;
58     }
59 
remove(Node * n)60     T * remove( Node * n ) {
61         if ( !n )
62             return 0;
63         T * r = n->data;
64 
65         if ( n->zero || n->one ) {
66             // this is an internal node, so we have to drop the
67             // data, then do no more.
68             n->data = 0;
69             return r;
70         }
71 
72         if ( !n->parent ) {
73             // this is the root
74             root = 0;
75             free( n );
76         }
77         else if ( n->parent->data ) {
78             // the parent has to lose this child, but has data, so
79             // it has to stay
80             if ( n->parent->zero == n )
81                 n->parent->zero = 0;
82             else
83                 n->parent->one = 0;
84             free( n );
85         }
86         else {
87             // the other child can be promoted to the parent's slot.
88             Node * p = n->parent;
89             Node * c = p->zero;
90             if ( c == n )
91                 c = p->one;
92             if ( c )
93                 c->parent = p->parent;
94             if ( !p->parent )
95                 root = c;
96             else if ( p->parent->one == p )
97                 p->parent->one = c;
98             else
99                 p->parent->zero = c;
100             free( p );
101             free( n );
102         }
103         return r;
104     }
105 
remove(const char * k,uint l)106     T * remove( const char * k, uint l ) {
107         return remove( locate( k, l ) );
108     }
109 
insert(const char * k,uint l,T * t)110     void insert( const char * k, uint l, T * t ) {
111         Node * n = root;
112         bool d = false;
113         uint b = 0;
114         while ( n && !d ) {
115             // check entire bytes until we run out of something
116             while ( !d &&
117                     b / 8 < l / 8 &&
118                     b / 8 < n->length / 8 ) {
119                 if ( k[b/8] == n->key[b/8] )
120                     b = ( b | 7 ) + 1;
121                 else
122                     d = true;
123             }
124             // if no entire byte differed, see if the last bits of n
125             // differ from k
126             if ( !d && b < n->length && n->length <= l ) {
127                 uint mask = ( 0xff00 >> (n->length%8) ) & 0xff;
128                 if ( ( k[b/8] & mask ) == ( n->key[b/8] & mask ) )
129                     b = n->length;
130                 else
131                     d = true;
132             }
133             // if we found a difference, then set b to the first
134             // differing bit
135             if ( d && b < n->length && b < l ) {
136                 uint mask = 128 >> ( b % 8 );
137                 while ( b < n->length && b < l &&
138                         ( k[b/8] & mask ) == ( n->key[b/8] & mask ) ) {
139                     b++;
140                     mask >>= 1;
141                     if ( !mask )
142                         mask = 128;
143                 }
144             }
145             // if the first differing bit is at the end of this node,
146             // then we need to go to the right child
147             if ( b == n->length ) {
148                 if ( b == l ) {
149                     // no, not to the child, n IS the right node
150                     n->data = t;
151                     return;
152                 }
153                 d = false;
154                 Node * c = 0;
155                 if ( k[b / 8] & ( 128 >> ( b % 8 ) ) )
156                     c = n->one;
157                 else
158                     c = n->zero;
159                 if ( c )
160                     n = c;
161                 else
162                     d = true;
163             }
164             if ( b == l && l < n->length )
165                 d = true;
166         }
167 
168         uint kl = (l+7) / 8;
169         Node * x = node( kl );
170         x->length = l;
171         x->data = t;
172         uint i = 0;
173         while ( i < kl ) {
174             x->key[i] = k[i];
175             i++;
176         }
177 
178         if ( !n ) {
179             // the tree is empty; x is the new root
180             root = x;
181         }
182         else if ( b == n->length ) {
183             // n's key is a prefix of k, so x must be a child of n
184             x->parent = n;
185             if ( k[b / 8] & ( 128 >> ( b % 8 ) ) )
186                 n->one = x;
187             else
188                 n->zero = x;
189         }
190         else if ( b == l ) {
191             // k is a prefix of n's key, so n must be a child of x
192             x->parent = n->parent;
193             n->parent = x;
194             if ( !x->parent )
195                 root = x;
196             else if ( x->parent->one == n )
197                 x->parent->one = x;
198             else
199                 x->parent->zero = x;
200             if ( n->key[b / 8] & ( 128 >> ( b % 8 ) ) )
201                 x->one = n;
202             else
203                 x->zero = n;
204         }
205         else {
206             // n's key and k differ, so we make a new intermediate node
207             kl = (b+7) / 8;
208             Node * p = node( kl );
209             x->parent = p;
210             p->parent = n->parent;
211             n->parent = p;
212             if ( !p->parent )
213                 root = p;
214             else if ( p->parent->one == n )
215                 p->parent->one = p;
216             else
217                 p->parent->zero = p;
218             if ( k[b / 8] & ( 128 >> ( b % 8 ) ) ) {
219                 p->zero = n;
220                 p->one = x;
221             }
222             else {
223                 p->zero = x;
224                 p->one = n;
225             }
226 
227             p->length = b;
228             i = 0;
229             while ( i < kl ) {
230                 p->key[i] = k[i];
231                 i++;
232             }
233         }
234     }
235 
isEmpty()236     bool isEmpty() {
237         if ( !root )
238             return true;
239         // is it possible for a tree to contain no data nodes? no?
240         return false;
241     }
242 
count()243     uint count() const {
244         if ( !root )
245             return 0;
246         return root->count();
247     }
248 
clear()249     void clear() {
250         root = 0;
251     }
252 
253     class Iterator
254         : public Garbage
255     {
256     public:
Iterator()257         Iterator()                   { cur = 0; }
Iterator(Node * n)258         Iterator( Node *n )          { cur = n; }
Iterator(PatriciaTree<T> * t)259         Iterator( PatriciaTree<T> * t ) {
260             if ( t )
261                 cur = t->firstNode();
262             else
263                 cur = 0;
264         }
Iterator(const PatriciaTree<T> & t)265         Iterator( const PatriciaTree<T> &t ) {
266             cur = t.firstNode();
267         }
268 
269         operator bool() { return cur && cur->data ? true : false; }
270         operator T *() { return cur ? cur->data : 0; }
271         T *operator ->() { ok(); return cur->data; }
272         Iterator &operator ++() { ok(); return next(); }
273         Iterator &operator --() { ok(); return prev(); }
274         Iterator &operator ++( int ) {
275             ok();
276             Node * p = cur; cur = next();
277             return newRef(p);
278         }
279 
280         Iterator &operator --( int ) {
281             ok();
282             Node *p = cur; cur = prev();
283             return newRef(p);
284         }
285 
286 
287         T &operator *() {
288             ok();
289             if ( !cur->data )
290                 die( Invariant );
291             return *(cur->data);
292         }
293 
294         bool operator ==( const Iterator &x ) { return cur == x.cur; }
295         bool operator !=( const Iterator &x ) { return cur != x.cur; }
296 
newRef(Node * n)297         static Iterator &newRef( Node *n ) {
298             return *( new Iterator(n) );
299         }
300 
301     private:
next()302         Iterator &next() {
303             do {
304                 if ( cur->one ) {
305                     cur = cur->one;
306                     while ( cur->zero )
307                         cur = cur->zero;
308                 }
309                 else if ( cur->parent ) {
310                     while ( cur->parent && cur->parent->one == cur )
311                         cur = cur->parent;
312                     if ( cur->parent )
313                         cur = cur->parent;
314                     else
315                         cur = 0;
316                 }
317                 else {
318                     cur = 0;
319                 }
320             } while ( cur && !cur->data );
321             return *this;
322         }
prev()323         Iterator &prev() {
324             do {
325                 if ( cur->zero ) {
326                     cur = cur->zero;
327                     while ( cur->one )
328                         cur = cur->one;
329                 }
330                 else if ( cur->parent ) {
331                     while ( cur->parent && cur->parent->zero == cur )
332                         cur = cur->parent;
333                     if ( cur->parent )
334                         cur = cur->parent;
335                     else
336                         cur = 0;
337                 }
338                 else {
339                     cur = 0;
340                 }
341             } while ( cur && !cur->data );
342             return *this;
343         }
344 
ok()345         void ok() {
346             if ( !cur )
347                 die( Invariant );
348         }
349 
350         Node *cur;
351     };
352 
remove(Iterator & i)353     T * remove( Iterator & i ) {
354         return remove( i.cur );
355     }
356 
first()357     T * first() {
358         Node * n = firstNode();
359         if ( n )
360             return n->data;
361         return 0;
362     }
363 
last()364     T * last() {
365         Node * n = lastNode();
366         if ( n )
367             return n->data;
368         return 0;
369     }
370 
371 private:
node(uint x)372     virtual Node * node( uint x ) {
373         Node * n = new ( x ) Node;
374         return n;
375     }
free(Node * n)376     virtual void free( Node * n ) {
377         Allocator::dealloc( n );
378     }
379 
best(const char * k,uint l)380     Node * best( const char * k, uint l ) const {
381         Node * n = root;
382         Node * p = n;
383         while ( n && n->length < l ) {
384             p = n;
385             if ( k[n->length / 8] & ( 128 >> ( n->length % 8 ) ) )
386                 n = n->one;
387             else
388                 n = n->zero;
389         }
390         if ( n )
391             return n;
392         return p;
393     }
394 
ifMatch(Node * n,const char * k,uint l)395     Node * ifMatch( Node * n, const char * k, uint l ) const {
396         if ( !n )
397             return 0;
398         if ( n->length != l )
399             return 0;
400         uint i = 0;
401         while ( i < l / 8 ) {
402             if ( n->key[i] != k[i] )
403                 return 0;
404             i++;
405         }
406         return n;
407     }
408 
locate(const char * k,uint l)409     Node * locate( const char * k, uint l ) const {
410         return ifMatch( best( k, l ), k, l );
411     }
412 
firstNode()413     Node * firstNode() const {
414         Node * n = root;
415         while ( n && n->zero )
416             n = n->zero;
417         return n;
418     }
419 
lastNode()420     Node * lastNode() const {
421         Node * n = root;
422         while ( n && n->one )
423             n = n->one;
424         return n;
425     }
426 
427 private:
428     Node * root;
429 };
430 
431 
432 #endif
433