Simplify code. We can compare TNew with T in one batch.
[oota-llvm.git] / include / llvm / ADT / ImmutableSet.h
1 //===--- ImmutableSet.h - Immutable (functional) set interface --*- C++ -*-===//
2 //
3 //                     The LLVM Compiler Infrastructure
4 //
5 // This file is distributed under the University of Illinois Open Source
6 // License. See LICENSE.TXT for details.
7 //
8 //===----------------------------------------------------------------------===//
9 //
10 // This file defines the ImutAVLTree and ImmutableSet classes.
11 //
12 //===----------------------------------------------------------------------===//
13
14 #ifndef LLVM_ADT_IMSET_H
15 #define LLVM_ADT_IMSET_H
16
17 #include "llvm/Support/Allocator.h"
18 #include "llvm/ADT/FoldingSet.h"
19 #include "llvm/System/DataTypes.h"
20 #include <cassert>
21 #include <functional>
22
23 namespace llvm {
24
25 //===----------------------------------------------------------------------===//
26 // Immutable AVL-Tree Definition.
27 //===----------------------------------------------------------------------===//
28
29 template <typename ImutInfo> class ImutAVLFactory;
30 template <typename ImutInfo> class ImutAVLTreeInOrderIterator;
31 template <typename ImutInfo> class ImutAVLTreeGenericIterator;
32
33 template <typename ImutInfo >
34 class ImutAVLTree : public FoldingSetNode {
35 public:
36   typedef typename ImutInfo::key_type_ref   key_type_ref;
37   typedef typename ImutInfo::value_type     value_type;
38   typedef typename ImutInfo::value_type_ref value_type_ref;
39
40   typedef ImutAVLFactory<ImutInfo>          Factory;
41   friend class ImutAVLFactory<ImutInfo>;
42
43   friend class ImutAVLTreeGenericIterator<ImutInfo>;
44   friend class FoldingSet<ImutAVLTree>;
45
46   typedef ImutAVLTreeInOrderIterator<ImutInfo>  iterator;
47
48   //===----------------------------------------------------===//
49   // Public Interface.
50   //===----------------------------------------------------===//
51
52   /// getLeft - Returns a pointer to the left subtree.  This value
53   ///  is NULL if there is no left subtree.
54   ImutAVLTree *getLeft() const { return Left; }
55
56   /// getRight - Returns a pointer to the right subtree.  This value is
57   ///  NULL if there is no right subtree.
58   ImutAVLTree *getRight() const { return Right; }
59
60   /// getHeight - Returns the height of the tree.  A tree with no subtrees
61   ///  has a height of 1.
62   unsigned getHeight() const { return Height; }
63
64   /// getValue - Returns the data value associated with the tree node.
65   const value_type& getValue() const { return Value; }
66
67   /// find - Finds the subtree associated with the specified key value.
68   ///  This method returns NULL if no matching subtree is found.
69   ImutAVLTree* find(key_type_ref K) {
70     ImutAVLTree *T = this;
71
72     while (T) {
73       key_type_ref CurrentKey = ImutInfo::KeyOfValue(T->getValue());
74
75       if (ImutInfo::isEqual(K,CurrentKey))
76         return T;
77       else if (ImutInfo::isLess(K,CurrentKey))
78         T = T->getLeft();
79       else
80         T = T->getRight();
81     }
82
83     return NULL;
84   }
85   
86   /// getMaxElement - Find the subtree associated with the highest ranged
87   ///  key value.
88   ImutAVLTree* getMaxElement() {
89     ImutAVLTree *T = this;
90     ImutAVLTree *Right = T->getRight();    
91     while (Right) { T = Right; Right = T->getRight(); }
92     return T;
93   }
94
95   /// size - Returns the number of nodes in the tree, which includes
96   ///  both leaves and non-leaf nodes.
97   unsigned size() const {
98     unsigned n = 1;
99
100     if (const ImutAVLTree* L = getLeft())  n += L->size();
101     if (const ImutAVLTree* R = getRight()) n += R->size();
102
103     return n;
104   }
105
106   /// begin - Returns an iterator that iterates over the nodes of the tree
107   ///  in an inorder traversal.  The returned iterator thus refers to the
108   ///  the tree node with the minimum data element.
109   iterator begin() const { return iterator(this); }
110
111   /// end - Returns an iterator for the tree that denotes the end of an
112   ///  inorder traversal.
113   iterator end() const { return iterator(); }
114
115   bool ElementEqual(value_type_ref V) const {
116     // Compare the keys.
117     if (!ImutInfo::isEqual(ImutInfo::KeyOfValue(getValue()),
118                            ImutInfo::KeyOfValue(V)))
119       return false;
120
121     // Also compare the data values.
122     if (!ImutInfo::isDataEqual(ImutInfo::DataOfValue(getValue()),
123                                ImutInfo::DataOfValue(V)))
124       return false;
125
126     return true;
127   }
128
129   bool ElementEqual(const ImutAVLTree* RHS) const {
130     return ElementEqual(RHS->getValue());
131   }
132
133   /// isEqual - Compares two trees for structural equality and returns true
134   ///   if they are equal.  This worst case performance of this operation is
135   //    linear in the sizes of the trees.
136   bool isEqual(const ImutAVLTree& RHS) const {
137     if (&RHS == this)
138       return true;
139
140     iterator LItr = begin(), LEnd = end();
141     iterator RItr = RHS.begin(), REnd = RHS.end();
142
143     while (LItr != LEnd && RItr != REnd) {
144       if (*LItr == *RItr) {
145         LItr.SkipSubTree();
146         RItr.SkipSubTree();
147         continue;
148       }
149
150       if (!LItr->ElementEqual(*RItr))
151         return false;
152
153       ++LItr;
154       ++RItr;
155     }
156
157     return LItr == LEnd && RItr == REnd;
158   }
159
160   /// isNotEqual - Compares two trees for structural inequality.  Performance
161   ///  is the same is isEqual.
162   bool isNotEqual(const ImutAVLTree& RHS) const { return !isEqual(RHS); }
163
164   /// contains - Returns true if this tree contains a subtree (node) that
165   ///  has an data element that matches the specified key.  Complexity
166   ///  is logarithmic in the size of the tree.
167   bool contains(key_type_ref K) { return (bool) find(K); }
168
169   /// foreach - A member template the accepts invokes operator() on a functor
170   ///  object (specifed by Callback) for every node/subtree in the tree.
171   ///  Nodes are visited using an inorder traversal.
172   template <typename Callback>
173   void foreach(Callback& C) {
174     if (ImutAVLTree* L = getLeft()) L->foreach(C);
175
176     C(Value);
177
178     if (ImutAVLTree* R = getRight()) R->foreach(C);
179   }
180
181   /// verify - A utility method that checks that the balancing and
182   ///  ordering invariants of the tree are satisifed.  It is a recursive
183   ///  method that returns the height of the tree, which is then consumed
184   ///  by the enclosing verify call.  External callers should ignore the
185   ///  return value.  An invalid tree will cause an assertion to fire in
186   ///  a debug build.
187   unsigned verify() const {
188     unsigned HL = getLeft() ? getLeft()->verify() : 0;
189     unsigned HR = getRight() ? getRight()->verify() : 0;
190
191     assert(getHeight() == ( HL > HR ? HL : HR ) + 1
192             && "Height calculation wrong");
193
194     assert((HL > HR ? HL-HR : HR-HL) <= 2
195            && "Balancing invariant violated");
196
197     assert(!getLeft()
198            || ImutInfo::isLess(ImutInfo::KeyOfValue(getLeft()->getValue()),
199                                ImutInfo::KeyOfValue(getValue()))
200            && "Value in left child is not less that current value");
201
202
203     assert(!getRight()
204            || ImutInfo::isLess(ImutInfo::KeyOfValue(getValue()),
205                                ImutInfo::KeyOfValue(getRight()->getValue()))
206            && "Current value is not less that value of right child");
207
208     return getHeight();
209   }
210
211   /// Profile - Profiling for ImutAVLTree.
212   void Profile(llvm::FoldingSetNodeID& ID) {
213     ID.AddInteger(ComputeDigest());
214   }
215
216   //===----------------------------------------------------===//
217   // Internal Values.
218   //===----------------------------------------------------===//
219
220 private:
221   ImutAVLTree*     Left;
222   ImutAVLTree*     Right;
223   unsigned         Height       : 28;
224   unsigned         Mutable      : 1;
225   unsigned         CachedDigest : 1;
226   value_type       Value;
227   uint32_t         Digest;
228
229   //===----------------------------------------------------===//
230   // Internal methods (node manipulation; used by Factory).
231   //===----------------------------------------------------===//
232
233 private:
234   /// ImutAVLTree - Internal constructor that is only called by
235   ///   ImutAVLFactory.
236   ImutAVLTree(ImutAVLTree* l, ImutAVLTree* r, value_type_ref v,
237               unsigned height)
238     : Left(l), Right(r), Height(height), Mutable(true), CachedDigest(false),
239       Value(v), Digest(0) {}
240
241   /// isMutable - Returns true if the left and right subtree references
242   ///  (as well as height) can be changed.  If this method returns false,
243   ///  the tree is truly immutable.  Trees returned from an ImutAVLFactory
244   ///  object should always have this method return true.  Further, if this
245   ///  method returns false for an instance of ImutAVLTree, all subtrees
246   ///  will also have this method return false.  The converse is not true.
247   bool isMutable() const { return Mutable; }
248   
249   /// hasCachedDigest - Returns true if the digest for this tree is cached.
250   ///  This can only be true if the tree is immutable.
251   bool hasCachedDigest() const { return CachedDigest; }
252
253   //===----------------------------------------------------===//
254   // Mutating operations.  A tree root can be manipulated as
255   // long as its reference has not "escaped" from internal
256   // methods of a factory object (see below).  When a tree
257   // pointer is externally viewable by client code, the
258   // internal "mutable bit" is cleared to mark the tree
259   // immutable.  Note that a tree that still has its mutable
260   // bit set may have children (subtrees) that are themselves
261   // immutable.
262   //===----------------------------------------------------===//
263
264   /// MarkImmutable - Clears the mutable flag for a tree.  After this happens,
265   ///   it is an error to call setLeft(), setRight(), and setHeight().
266   void MarkImmutable() {
267     assert(isMutable() && "Mutable flag already removed.");
268     Mutable = false;
269   }
270   
271   /// MarkedCachedDigest - Clears the NoCachedDigest flag for a tree.
272   void MarkedCachedDigest() {
273     assert(!hasCachedDigest() && "NoCachedDigest flag already removed.");
274     CachedDigest = true;
275   }
276
277   /// setLeft - Changes the reference of the left subtree.  Used internally
278   ///   by ImutAVLFactory.
279   void setLeft(ImutAVLTree* NewLeft) {
280     assert(isMutable() &&
281            "Only a mutable tree can have its left subtree changed.");
282     Left = NewLeft;
283     CachedDigest = false;
284   }
285
286   /// setRight - Changes the reference of the right subtree.  Used internally
287   ///  by ImutAVLFactory.
288   void setRight(ImutAVLTree* NewRight) {
289     assert(isMutable() &&
290            "Only a mutable tree can have its right subtree changed.");
291
292     Right = NewRight;
293     CachedDigest = false;
294   }
295
296   /// setHeight - Changes the height of the tree.  Used internally by
297   ///  ImutAVLFactory.
298   void setHeight(unsigned h) {
299     assert(isMutable() && "Only a mutable tree can have its height changed.");
300     Height = h;
301   }
302
303   static inline
304   uint32_t ComputeDigest(ImutAVLTree* L, ImutAVLTree* R, value_type_ref V) {
305     uint32_t digest = 0;
306
307     if (L)
308       digest += L->ComputeDigest();
309
310     // Compute digest of stored data.
311     FoldingSetNodeID ID;
312     ImutInfo::Profile(ID,V);
313     digest += ID.ComputeHash();
314
315     if (R)
316       digest += R->ComputeDigest();
317
318     return digest;
319   }
320
321   inline uint32_t ComputeDigest() {
322     // Check the lowest bit to determine if digest has actually been
323     // pre-computed.
324     if (hasCachedDigest())
325       return Digest;
326
327     uint32_t X = ComputeDigest(getLeft(), getRight(), getValue());
328     Digest = X;
329     MarkedCachedDigest();
330     return X;
331   }
332 };
333
334 //===----------------------------------------------------------------------===//
335 // Immutable AVL-Tree Factory class.
336 //===----------------------------------------------------------------------===//
337
338 template <typename ImutInfo >
339 class ImutAVLFactory {
340   typedef ImutAVLTree<ImutInfo> TreeTy;
341   typedef typename TreeTy::value_type_ref value_type_ref;
342   typedef typename TreeTy::key_type_ref   key_type_ref;
343
344   typedef FoldingSet<TreeTy> CacheTy;
345
346   CacheTy Cache;
347   uintptr_t Allocator;
348
349   bool ownsAllocator() const {
350     return Allocator & 0x1 ? false : true;
351   }
352
353   BumpPtrAllocator& getAllocator() const {
354     return *reinterpret_cast<BumpPtrAllocator*>(Allocator & ~0x1);
355   }
356
357   //===--------------------------------------------------===//
358   // Public interface.
359   //===--------------------------------------------------===//
360
361 public:
362   ImutAVLFactory()
363     : Allocator(reinterpret_cast<uintptr_t>(new BumpPtrAllocator())) {}
364
365   ImutAVLFactory(BumpPtrAllocator& Alloc)
366     : Allocator(reinterpret_cast<uintptr_t>(&Alloc) | 0x1) {}
367
368   ~ImutAVLFactory() {
369     if (ownsAllocator()) delete &getAllocator();
370   }
371
372   TreeTy* Add(TreeTy* T, value_type_ref V) {
373     T = Add_internal(V,T);
374     MarkImmutable(T);
375     return T;
376   }
377
378   TreeTy* Remove(TreeTy* T, key_type_ref V) {
379     T = Remove_internal(V,T);
380     MarkImmutable(T);
381     return T;
382   }
383
384   TreeTy* GetEmptyTree() const { return NULL; }
385
386   //===--------------------------------------------------===//
387   // A bunch of quick helper functions used for reasoning
388   // about the properties of trees and their children.
389   // These have succinct names so that the balancing code
390   // is as terse (and readable) as possible.
391   //===--------------------------------------------------===//
392 private:
393
394   bool           isEmpty(TreeTy* T) const { return !T; }
395   unsigned Height(TreeTy* T) const { return T ? T->getHeight() : 0; }
396   TreeTy*           Left(TreeTy* T) const { return T->getLeft(); }
397   TreeTy*          Right(TreeTy* T) const { return T->getRight(); }
398   value_type_ref   Value(TreeTy* T) const { return T->Value; }
399
400   unsigned IncrementHeight(TreeTy* L, TreeTy* R) const {
401     unsigned hl = Height(L);
402     unsigned hr = Height(R);
403     return (hl > hr ? hl : hr) + 1;
404   }
405
406   static bool CompareTreeWithSection(TreeTy* T,
407                                      typename TreeTy::iterator& TI,
408                                      typename TreeTy::iterator& TE) {
409
410     typename TreeTy::iterator I = T->begin(), E = T->end();
411
412     for ( ; I!=E ; ++I, ++TI)
413       if (TI == TE || !I->ElementEqual(*TI))
414         return false;
415
416     return true;
417   }
418
419   //===--------------------------------------------------===//
420   // "CreateNode" is used to generate new tree roots that link
421   // to other trees.  The functon may also simply move links
422   // in an existing root if that root is still marked mutable.
423   // This is necessary because otherwise our balancing code
424   // would leak memory as it would create nodes that are
425   // then discarded later before the finished tree is
426   // returned to the caller.
427   //===--------------------------------------------------===//
428
429   TreeTy* CreateNode(TreeTy* L, value_type_ref V, TreeTy* R) {   
430     BumpPtrAllocator& A = getAllocator();
431     TreeTy* T = (TreeTy*) A.Allocate<TreeTy>();
432     new (T) TreeTy(L, R, V, IncrementHeight(L,R));
433     return T;
434   }
435
436   TreeTy* CreateNode(TreeTy* L, TreeTy* OldTree, TreeTy* R) {
437     assert(!isEmpty(OldTree));
438
439     if (OldTree->isMutable()) {
440       OldTree->setLeft(L);
441       OldTree->setRight(R);
442       OldTree->setHeight(IncrementHeight(L, R));
443       return OldTree;
444     }
445     else
446       return CreateNode(L, Value(OldTree), R);
447   }
448
449   /// Balance - Used by Add_internal and Remove_internal to
450   ///  balance a newly created tree.
451   TreeTy* Balance(TreeTy* L, value_type_ref V, TreeTy* R) {
452
453     unsigned hl = Height(L);
454     unsigned hr = Height(R);
455
456     if (hl > hr + 2) {
457       assert(!isEmpty(L) && "Left tree cannot be empty to have a height >= 2");
458
459       TreeTy* LL = Left(L);
460       TreeTy* LR = Right(L);
461
462       if (Height(LL) >= Height(LR))
463         return CreateNode(LL, L, CreateNode(LR,V,R));
464
465       assert(!isEmpty(LR) && "LR cannot be empty because it has a height >= 1");
466
467       TreeTy* LRL = Left(LR);
468       TreeTy* LRR = Right(LR);
469
470       return CreateNode(CreateNode(LL,L,LRL), LR, CreateNode(LRR,V,R));
471     }
472     else if (hr > hl + 2) {
473       assert(!isEmpty(R) && "Right tree cannot be empty to have a height >= 2");
474
475       TreeTy* RL = Left(R);
476       TreeTy* RR = Right(R);
477
478       if (Height(RR) >= Height(RL))
479         return CreateNode(CreateNode(L,V,RL), R, RR);
480
481       assert(!isEmpty(RL) && "RL cannot be empty because it has a height >= 1");
482
483       TreeTy* RLL = Left(RL);
484       TreeTy* RLR = Right(RL);
485
486       return CreateNode(CreateNode(L,V,RLL), RL, CreateNode(RLR,R,RR));
487     }
488     else
489       return CreateNode(L,V,R);
490   }
491
492   /// Add_internal - Creates a new tree that includes the specified
493   ///  data and the data from the original tree.  If the original tree
494   ///  already contained the data item, the original tree is returned.
495   TreeTy* Add_internal(value_type_ref V, TreeTy* T) {
496     if (isEmpty(T))
497       return CreateNode(T, V, T);
498
499     assert(!T->isMutable());
500
501     key_type_ref K = ImutInfo::KeyOfValue(V);
502     key_type_ref KCurrent = ImutInfo::KeyOfValue(Value(T));
503
504     if (ImutInfo::isEqual(K,KCurrent))
505       return CreateNode(Left(T), V, Right(T));
506     else if (ImutInfo::isLess(K,KCurrent))
507       return Balance(Add_internal(V,Left(T)), Value(T), Right(T));
508     else
509       return Balance(Left(T), Value(T), Add_internal(V,Right(T)));
510   }
511
512   /// Remove_internal - Creates a new tree that includes all the data
513   ///  from the original tree except the specified data.  If the
514   ///  specified data did not exist in the original tree, the original
515   ///  tree is returned.
516   TreeTy* Remove_internal(key_type_ref K, TreeTy* T) {
517     if (isEmpty(T))
518       return T;
519
520     assert(!T->isMutable());
521
522     key_type_ref KCurrent = ImutInfo::KeyOfValue(Value(T));
523
524     if (ImutInfo::isEqual(K,KCurrent))
525       return CombineLeftRightTrees(Left(T),Right(T));
526     else if (ImutInfo::isLess(K,KCurrent))
527       return Balance(Remove_internal(K,Left(T)), Value(T), Right(T));
528     else
529       return Balance(Left(T), Value(T), Remove_internal(K,Right(T)));
530   }
531
532   TreeTy* CombineLeftRightTrees(TreeTy* L, TreeTy* R) {
533     if (isEmpty(L)) return R;
534     if (isEmpty(R)) return L;
535
536     TreeTy* OldNode;
537     TreeTy* NewRight = RemoveMinBinding(R,OldNode);
538     return Balance(L,Value(OldNode),NewRight);
539   }
540
541   TreeTy* RemoveMinBinding(TreeTy* T, TreeTy*& NodeRemoved) {
542     assert(!isEmpty(T));
543
544     if (isEmpty(Left(T))) {
545       NodeRemoved = T;
546       return Right(T);
547     }
548
549     return Balance(RemoveMinBinding(Left(T),NodeRemoved),Value(T),Right(T));
550   }
551
552   /// MarkImmutable - Clears the mutable bits of a root and all of its
553   ///  descendants.
554   void MarkImmutable(TreeTy* T) {
555     if (!T || !T->isMutable())
556       return;
557
558     T->MarkImmutable();
559     MarkImmutable(Left(T));
560     MarkImmutable(Right(T));
561   }
562   
563 public:
564   TreeTy *GetCanonicalTree(TreeTy *TNew) {
565     if (!TNew)
566       return NULL;    
567     
568     // Search the FoldingSet bucket for a Tree with the same digest.
569     FoldingSetNodeID ID;
570     unsigned digest = TNew->ComputeDigest();
571     ID.AddInteger(digest);
572     unsigned hash = ID.ComputeHash();
573     
574     typename CacheTy::bucket_iterator I = Cache.bucket_begin(hash);
575     typename CacheTy::bucket_iterator E = Cache.bucket_end(hash);
576     
577     for (; I != E; ++I) {
578       TreeTy *T = &*I;
579       
580       if (T->ComputeDigest() != digest)
581         continue;
582       
583       // We found a collision.  Perform a comparison of Contents('T')
584       // with Contents('TNew')
585       typename TreeTy::iterator TI = T->begin(), TE = T->end();
586       
587       if (!CompareTreeWithSection(TNew, TI, TE))
588         continue;
589       
590       if (TI != TE)
591         continue; // T has more contents than TNew.
592       
593       // Trees did match!  Return 'T'.
594       return T;
595     }
596
597     // 'TNew' is the only tree of its kind.  Return it.
598     Cache.InsertNode(TNew, (void*) &*Cache.bucket_end(hash));
599     return TNew;
600   }
601 };
602
603
604 //===----------------------------------------------------------------------===//
605 // Immutable AVL-Tree Iterators.
606 //===----------------------------------------------------------------------===//
607
608 template <typename ImutInfo>
609 class ImutAVLTreeGenericIterator {
610   SmallVector<uintptr_t,20> stack;
611 public:
612   enum VisitFlag { VisitedNone=0x0, VisitedLeft=0x1, VisitedRight=0x3,
613                    Flags=0x3 };
614
615   typedef ImutAVLTree<ImutInfo> TreeTy;
616   typedef ImutAVLTreeGenericIterator<ImutInfo> _Self;
617
618   inline ImutAVLTreeGenericIterator() {}
619   inline ImutAVLTreeGenericIterator(const TreeTy* Root) {
620     if (Root) stack.push_back(reinterpret_cast<uintptr_t>(Root));
621   }
622
623   TreeTy* operator*() const {
624     assert(!stack.empty());
625     return reinterpret_cast<TreeTy*>(stack.back() & ~Flags);
626   }
627
628   uintptr_t getVisitState() {
629     assert(!stack.empty());
630     return stack.back() & Flags;
631   }
632
633
634   bool AtEnd() const { return stack.empty(); }
635
636   bool AtBeginning() const {
637     return stack.size() == 1 && getVisitState() == VisitedNone;
638   }
639
640   void SkipToParent() {
641     assert(!stack.empty());
642     stack.pop_back();
643
644     if (stack.empty())
645       return;
646
647     switch (getVisitState()) {
648       case VisitedNone:
649         stack.back() |= VisitedLeft;
650         break;
651       case VisitedLeft:
652         stack.back() |= VisitedRight;
653         break;
654       default:
655         assert(false && "Unreachable.");
656     }
657   }
658
659   inline bool operator==(const _Self& x) const {
660     if (stack.size() != x.stack.size())
661       return false;
662
663     for (unsigned i = 0 ; i < stack.size(); i++)
664       if (stack[i] != x.stack[i])
665         return false;
666
667     return true;
668   }
669
670   inline bool operator!=(const _Self& x) const { return !operator==(x); }
671
672   _Self& operator++() {
673     assert(!stack.empty());
674
675     TreeTy* Current = reinterpret_cast<TreeTy*>(stack.back() & ~Flags);
676     assert(Current);
677
678     switch (getVisitState()) {
679       case VisitedNone:
680         if (TreeTy* L = Current->getLeft())
681           stack.push_back(reinterpret_cast<uintptr_t>(L));
682         else
683           stack.back() |= VisitedLeft;
684
685         break;
686
687       case VisitedLeft:
688         if (TreeTy* R = Current->getRight())
689           stack.push_back(reinterpret_cast<uintptr_t>(R));
690         else
691           stack.back() |= VisitedRight;
692
693         break;
694
695       case VisitedRight:
696         SkipToParent();
697         break;
698
699       default:
700         assert(false && "Unreachable.");
701     }
702
703     return *this;
704   }
705
706   _Self& operator--() {
707     assert(!stack.empty());
708
709     TreeTy* Current = reinterpret_cast<TreeTy*>(stack.back() & ~Flags);
710     assert(Current);
711
712     switch (getVisitState()) {
713       case VisitedNone:
714         stack.pop_back();
715         break;
716
717       case VisitedLeft:
718         stack.back() &= ~Flags; // Set state to "VisitedNone."
719
720         if (TreeTy* L = Current->getLeft())
721           stack.push_back(reinterpret_cast<uintptr_t>(L) | VisitedRight);
722
723         break;
724
725       case VisitedRight:
726         stack.back() &= ~Flags;
727         stack.back() |= VisitedLeft;
728
729         if (TreeTy* R = Current->getRight())
730           stack.push_back(reinterpret_cast<uintptr_t>(R) | VisitedRight);
731
732         break;
733
734       default:
735         assert(false && "Unreachable.");
736     }
737
738     return *this;
739   }
740 };
741
742 template <typename ImutInfo>
743 class ImutAVLTreeInOrderIterator {
744   typedef ImutAVLTreeGenericIterator<ImutInfo> InternalIteratorTy;
745   InternalIteratorTy InternalItr;
746
747 public:
748   typedef ImutAVLTree<ImutInfo> TreeTy;
749   typedef ImutAVLTreeInOrderIterator<ImutInfo> _Self;
750
751   ImutAVLTreeInOrderIterator(const TreeTy* Root) : InternalItr(Root) {
752     if (Root) operator++(); // Advance to first element.
753   }
754
755   ImutAVLTreeInOrderIterator() : InternalItr() {}
756
757   inline bool operator==(const _Self& x) const {
758     return InternalItr == x.InternalItr;
759   }
760
761   inline bool operator!=(const _Self& x) const { return !operator==(x); }
762
763   inline TreeTy* operator*() const { return *InternalItr; }
764   inline TreeTy* operator->() const { return *InternalItr; }
765
766   inline _Self& operator++() {
767     do ++InternalItr;
768     while (!InternalItr.AtEnd() &&
769            InternalItr.getVisitState() != InternalIteratorTy::VisitedLeft);
770
771     return *this;
772   }
773
774   inline _Self& operator--() {
775     do --InternalItr;
776     while (!InternalItr.AtBeginning() &&
777            InternalItr.getVisitState() != InternalIteratorTy::VisitedLeft);
778
779     return *this;
780   }
781
782   inline void SkipSubTree() {
783     InternalItr.SkipToParent();
784
785     while (!InternalItr.AtEnd() &&
786            InternalItr.getVisitState() != InternalIteratorTy::VisitedLeft)
787       ++InternalItr;
788   }
789 };
790
791 //===----------------------------------------------------------------------===//
792 // Trait classes for Profile information.
793 //===----------------------------------------------------------------------===//
794
795 /// Generic profile template.  The default behavior is to invoke the
796 /// profile method of an object.  Specializations for primitive integers
797 /// and generic handling of pointers is done below.
798 template <typename T>
799 struct ImutProfileInfo {
800   typedef const T  value_type;
801   typedef const T& value_type_ref;
802
803   static inline void Profile(FoldingSetNodeID& ID, value_type_ref X) {
804     FoldingSetTrait<T>::Profile(X,ID);
805   }
806 };
807
808 /// Profile traits for integers.
809 template <typename T>
810 struct ImutProfileInteger {
811   typedef const T  value_type;
812   typedef const T& value_type_ref;
813
814   static inline void Profile(FoldingSetNodeID& ID, value_type_ref X) {
815     ID.AddInteger(X);
816   }
817 };
818
819 #define PROFILE_INTEGER_INFO(X)\
820 template<> struct ImutProfileInfo<X> : ImutProfileInteger<X> {};
821
822 PROFILE_INTEGER_INFO(char)
823 PROFILE_INTEGER_INFO(unsigned char)
824 PROFILE_INTEGER_INFO(short)
825 PROFILE_INTEGER_INFO(unsigned short)
826 PROFILE_INTEGER_INFO(unsigned)
827 PROFILE_INTEGER_INFO(signed)
828 PROFILE_INTEGER_INFO(long)
829 PROFILE_INTEGER_INFO(unsigned long)
830 PROFILE_INTEGER_INFO(long long)
831 PROFILE_INTEGER_INFO(unsigned long long)
832
833 #undef PROFILE_INTEGER_INFO
834
835 /// Generic profile trait for pointer types.  We treat pointers as
836 /// references to unique objects.
837 template <typename T>
838 struct ImutProfileInfo<T*> {
839   typedef const T*   value_type;
840   typedef value_type value_type_ref;
841
842   static inline void Profile(FoldingSetNodeID &ID, value_type_ref X) {
843     ID.AddPointer(X);
844   }
845 };
846
847 //===----------------------------------------------------------------------===//
848 // Trait classes that contain element comparison operators and type
849 //  definitions used by ImutAVLTree, ImmutableSet, and ImmutableMap.  These
850 //  inherit from the profile traits (ImutProfileInfo) to include operations
851 //  for element profiling.
852 //===----------------------------------------------------------------------===//
853
854
855 /// ImutContainerInfo - Generic definition of comparison operations for
856 ///   elements of immutable containers that defaults to using
857 ///   std::equal_to<> and std::less<> to perform comparison of elements.
858 template <typename T>
859 struct ImutContainerInfo : public ImutProfileInfo<T> {
860   typedef typename ImutProfileInfo<T>::value_type      value_type;
861   typedef typename ImutProfileInfo<T>::value_type_ref  value_type_ref;
862   typedef value_type      key_type;
863   typedef value_type_ref  key_type_ref;
864   typedef bool            data_type;
865   typedef bool            data_type_ref;
866
867   static inline key_type_ref KeyOfValue(value_type_ref D) { return D; }
868   static inline data_type_ref DataOfValue(value_type_ref) { return true; }
869
870   static inline bool isEqual(key_type_ref LHS, key_type_ref RHS) {
871     return std::equal_to<key_type>()(LHS,RHS);
872   }
873
874   static inline bool isLess(key_type_ref LHS, key_type_ref RHS) {
875     return std::less<key_type>()(LHS,RHS);
876   }
877
878   static inline bool isDataEqual(data_type_ref,data_type_ref) { return true; }
879 };
880
881 /// ImutContainerInfo - Specialization for pointer values to treat pointers
882 ///  as references to unique objects.  Pointers are thus compared by
883 ///  their addresses.
884 template <typename T>
885 struct ImutContainerInfo<T*> : public ImutProfileInfo<T*> {
886   typedef typename ImutProfileInfo<T*>::value_type      value_type;
887   typedef typename ImutProfileInfo<T*>::value_type_ref  value_type_ref;
888   typedef value_type      key_type;
889   typedef value_type_ref  key_type_ref;
890   typedef bool            data_type;
891   typedef bool            data_type_ref;
892
893   static inline key_type_ref KeyOfValue(value_type_ref D) { return D; }
894   static inline data_type_ref DataOfValue(value_type_ref) { return true; }
895
896   static inline bool isEqual(key_type_ref LHS, key_type_ref RHS) {
897     return LHS == RHS;
898   }
899
900   static inline bool isLess(key_type_ref LHS, key_type_ref RHS) {
901     return LHS < RHS;
902   }
903
904   static inline bool isDataEqual(data_type_ref,data_type_ref) { return true; }
905 };
906
907 //===----------------------------------------------------------------------===//
908 // Immutable Set
909 //===----------------------------------------------------------------------===//
910
911 template <typename ValT, typename ValInfo = ImutContainerInfo<ValT> >
912 class ImmutableSet {
913 public:
914   typedef typename ValInfo::value_type      value_type;
915   typedef typename ValInfo::value_type_ref  value_type_ref;
916   typedef ImutAVLTree<ValInfo> TreeTy;
917
918 private:
919   TreeTy *Root;
920   
921 public:
922   /// Constructs a set from a pointer to a tree root.  In general one
923   /// should use a Factory object to create sets instead of directly
924   /// invoking the constructor, but there are cases where make this
925   /// constructor public is useful.
926   explicit ImmutableSet(TreeTy* R) : Root(R) {}
927
928   class Factory {
929     typename TreeTy::Factory F;
930     const bool Canonicalize;
931
932   public:
933     Factory(bool canonicalize = true)
934       : Canonicalize(canonicalize) {}
935
936     Factory(BumpPtrAllocator& Alloc, bool canonicalize = true)
937       : F(Alloc), Canonicalize(canonicalize) {}
938
939     /// GetEmptySet - Returns an immutable set that contains no elements.
940     ImmutableSet GetEmptySet() {
941       return ImmutableSet(F.GetEmptyTree());
942     }
943
944     /// Add - Creates a new immutable set that contains all of the values
945     ///  of the original set with the addition of the specified value.  If
946     ///  the original set already included the value, then the original set is
947     ///  returned and no memory is allocated.  The time and space complexity
948     ///  of this operation is logarithmic in the size of the original set.
949     ///  The memory allocated to represent the set is released when the
950     ///  factory object that created the set is destroyed.
951     ImmutableSet Add(ImmutableSet Old, value_type_ref V) {
952       TreeTy *NewT = F.Add(Old.Root, V);
953       return ImmutableSet(Canonicalize ? F.GetCanonicalTree(NewT) : NewT);
954     }
955
956     /// Remove - Creates a new immutable set that contains all of the values
957     ///  of the original set with the exception of the specified value.  If
958     ///  the original set did not contain the value, the original set is
959     ///  returned and no memory is allocated.  The time and space complexity
960     ///  of this operation is logarithmic in the size of the original set.
961     ///  The memory allocated to represent the set is released when the
962     ///  factory object that created the set is destroyed.
963     ImmutableSet Remove(ImmutableSet Old, value_type_ref V) {
964       TreeTy *NewT = F.Remove(Old.Root, V);
965       return ImmutableSet(Canonicalize ? F.GetCanonicalTree(NewT) : NewT);
966     }
967
968     BumpPtrAllocator& getAllocator() { return F.getAllocator(); }
969
970   private:
971     Factory(const Factory& RHS); // DO NOT IMPLEMENT
972     void operator=(const Factory& RHS); // DO NOT IMPLEMENT
973   };
974
975   friend class Factory;
976
977   /// contains - Returns true if the set contains the specified value.
978   bool contains(value_type_ref V) const {
979     return Root ? Root->contains(V) : false;
980   }
981
982   bool operator==(ImmutableSet RHS) const {
983     return Root && RHS.Root ? Root->isEqual(*RHS.Root) : Root == RHS.Root;
984   }
985
986   bool operator!=(ImmutableSet RHS) const {
987     return Root && RHS.Root ? Root->isNotEqual(*RHS.Root) : Root != RHS.Root;
988   }
989
990   TreeTy *getRoot() { 
991     return Root;
992   }
993
994   /// isEmpty - Return true if the set contains no elements.
995   bool isEmpty() const { return !Root; }
996
997   /// isSingleton - Return true if the set contains exactly one element.
998   ///   This method runs in constant time.
999   bool isSingleton() const { return getHeight() == 1; }
1000
1001   template <typename Callback>
1002   void foreach(Callback& C) { if (Root) Root->foreach(C); }
1003
1004   template <typename Callback>
1005   void foreach() { if (Root) { Callback C; Root->foreach(C); } }
1006
1007   //===--------------------------------------------------===//
1008   // Iterators.
1009   //===--------------------------------------------------===//
1010
1011   class iterator {
1012     typename TreeTy::iterator itr;
1013     iterator(TreeTy* t) : itr(t) {}
1014     friend class ImmutableSet<ValT,ValInfo>;
1015   public:
1016     iterator() {}
1017     inline value_type_ref operator*() const { return itr->getValue(); }
1018     inline iterator& operator++() { ++itr; return *this; }
1019     inline iterator  operator++(int) { iterator tmp(*this); ++itr; return tmp; }
1020     inline iterator& operator--() { --itr; return *this; }
1021     inline iterator  operator--(int) { iterator tmp(*this); --itr; return tmp; }
1022     inline bool operator==(const iterator& RHS) const { return RHS.itr == itr; }
1023     inline bool operator!=(const iterator& RHS) const { return RHS.itr != itr; }
1024     inline value_type *operator->() const { return &(operator*()); }
1025   };
1026
1027   iterator begin() const { return iterator(Root); }
1028   iterator end() const { return iterator(); }
1029
1030   //===--------------------------------------------------===//
1031   // Utility methods.
1032   //===--------------------------------------------------===//
1033
1034   unsigned getHeight() const { return Root ? Root->getHeight() : 0; }
1035
1036   static inline void Profile(FoldingSetNodeID& ID, const ImmutableSet& S) {
1037     ID.AddPointer(S.Root);
1038   }
1039
1040   inline void Profile(FoldingSetNodeID& ID) const {
1041     return Profile(ID,*this);
1042   }
1043
1044   //===--------------------------------------------------===//
1045   // For testing.
1046   //===--------------------------------------------------===//
1047
1048   void verify() const { if (Root) Root->verify(); }
1049 };
1050
1051 } // end namespace llvm
1052
1053 #endif