add some methods, fix a major bug in getLeader() that was causing things to
[oota-llvm.git] / include / llvm / ADT / EquivalenceClasses.h
1 //===-- llvm/ADT/EquivalenceClasses.h - Generic Equiv. Classes --*- C++ -*-===//
2 // 
3 //                     The LLVM Compiler Infrastructure
4 //
5 // This file was developed by the LLVM research group and is distributed under
6 // the University of Illinois Open Source License. See LICENSE.TXT for details.
7 // 
8 //===----------------------------------------------------------------------===//
9 // 
10 // Generic implementation of equivalence classes through the use Tarjan's
11 // efficient union-find algorithm.
12 // 
13 //===----------------------------------------------------------------------===//
14
15 #ifndef LLVM_ADT_EQUIVALENCECLASSES_H
16 #define LLVM_ADT_EQUIVALENCECLASSES_H
17
18 #include "llvm/ADT/iterator"
19 #include <set>
20
21 namespace llvm {
22
23 /// EquivalenceClasses - This represents a collection of equivalence classes and
24 /// supports three efficient operations: insert an element into a class of its
25 /// own, union two classes, and find the class for a given element.  In
26 /// addition to these modification methods, it is possible to iterate over all
27 /// of the equivalence classes and all of the elements in a class.
28 ///
29 /// This implementation is an efficient implementation that only stores one copy
30 /// of the element being indexed per entry in the set, and allows any arbitrary
31 /// type to be indexed (as long as it can be ordered with operator<).
32 ///
33 /// Here is a simple example using integers:
34 ///
35 ///  EquivalenceClasses<int> EC;
36 ///  EC.unionSets(1, 2);                // insert 1, 2 into the same set
37 ///  EC.insert(4); EC.insert(5);        // insert 4, 5 into own sets
38 ///  EC.unionSets(5, 1);                // merge the set for 1 with 5's set.
39 ///
40 ///  for (EquivalenceClasses<int>::iterator I = EC.begin(), E = EC.end();
41 ///       I != E; ++I) {           // Iterate over all of the equivalence sets.
42 ///    if (!I->isLeader()) continue;   // Ignore non-leader sets.
43 ///    for (EquivalenceClasses<int>::member_iterator MI = EC.member_begin(I);
44 ///         MI != EC.member_end(); ++MI)   // Loop over members in this set.
45 ///      std::cerr << *MI << " ";  // Print member.
46 ///    std::cerr << "\n";   // Finish set.
47 ///  }
48 ///
49 /// This example prints:
50 ///   4
51 ///   5 1 2
52 ///
53 template <class ElemTy>
54 class EquivalenceClasses {
55   /// ECValue - The EquivalenceClasses data structure is just a set of these.
56   /// Each of these represents a relation for a value.  First it stores the
57   /// value itself, which provides the ordering that the set queries.  Next, it
58   /// provides a "next pointer", which is used to enumerate all of the elements
59   /// in the unioned set.  Finally, it defines either a "end of list pointer" or
60   /// "leader pointer" depending on whether the value itself is a leader.  A
61   /// "leader pointer" points to the node that is the leader for this element,
62   /// if the node is not a leader.  A "end of list pointer" points to the last
63   /// node in the list of members of this list.  Whether or not a node is a
64   /// leader is determined by a bit stolen from one of the pointers.
65   class ECValue {
66     friend class EquivalenceClasses;
67     mutable const ECValue *Leader, *Next;
68     ElemTy Data;
69     // ECValue ctor - Start out with EndOfList pointing to this node, Next is
70     // Null, isLeader = true.
71     ECValue(const ElemTy &Elt)
72       : Leader(this), Next((ECValue*)(intptr_t)1), Data(Elt) {}
73
74     const ECValue *getLeader() const {
75       if (isLeader()) return this;
76       if (Leader->isLeader()) return Leader;
77       // Path compression.
78       return Leader = Leader->getLeader();
79     }
80     const ECValue *getEndOfList() const {
81       assert(isLeader() && "Cannot get the end of a list for a non-leader!");
82       return Leader;
83     }
84
85     void setNext(const ECValue *NewNext) const {
86       assert(getNext() == 0 && "Already has a next pointer!");
87       bool isL = isLeader();
88       Next = (const ECValue*)((intptr_t)NewNext | isLeader());
89     }
90   public:
91     ECValue(const ECValue &RHS) : Leader(this), Next((ECValue*)(intptr_t)1),
92                                   Data(RHS.Data) {
93       // Only support copying of singleton nodes.
94       assert(RHS.isLeader() && RHS.getNext() == 0 && "Not a singleton!");
95     }
96
97     bool operator<(const ECValue &UFN) const { return Data < UFN.Data; }
98
99     bool isLeader() const { return (intptr_t)Next & 1; }
100     const ElemTy &getData() const { return Data; }
101
102     const ECValue *getNext() const {
103       return (ECValue*)((intptr_t)Next & ~(intptr_t)1);
104     }
105
106     template<typename T>
107     bool operator<(const T &Val) const { return Data < Val; }
108   };
109
110   /// TheMapping - This implicitly provides a mapping from ElemTy values to the
111   /// ECValues, it just keeps the key as part of the value.
112   std::set<ECValue> TheMapping;
113
114 public:
115   EquivalenceClasses() {}
116   EquivalenceClasses(const EquivalenceClasses &RHS) {
117     operator=(RHS);
118   }
119
120   const EquivalenceClasses &operator=(const EquivalenceClasses &RHS) {
121     for (iterator I = RHS.begin(), E = RHS.end(); I != E; ++I)
122       if (I->isLeader())
123         insert(I->getData());
124       else
125         unionSets(I->getData(), *RHS.findLeader(I));
126     return *this;
127   }
128   
129   //===--------------------------------------------------------------------===//
130   // Inspection methods
131   //
132
133   /// iterator* - Provides a way to iterate over all values in the set.
134   typedef typename std::set<ECValue>::const_iterator iterator;
135   iterator begin() const { return TheMapping.begin(); }
136   iterator end() const { return TheMapping.end(); }
137
138   /// member_* Iterate over the members of an equivalence class.
139   ///
140   class member_iterator;
141   member_iterator member_begin(iterator I) const {
142     // Only leaders provide anything to iterate over.
143     return member_iterator(I->isLeader() ? &*I : 0);
144   }
145   member_iterator member_end() const {
146     return member_iterator(0);
147   }
148
149   /// findValue - Return an iterator to the specified value.  If it does not
150   /// exist, end() is returned.
151   iterator findValue(const ElemTy &V) const {
152     return TheMapping.find(V);
153   }
154
155   /// getLeaderValue - Return the leader for the specified value that is in the
156   /// set.  It is an error to call this method for a value that is not yet in
157   /// the set.  For that, call getOrInsertLeaderValue(V).
158   const ElemTy &getLeaderValue(const ElemTy &V) const {
159     member_iterator MI = findLeader(V);
160     assert(MI != member_end() && "Value is not in the set!");
161     return *MI;
162   }
163
164   /// getOrInsertLeaderValue - Return the leader for the specified value that is
165   /// in the set.  If the member is not in the set, it is inserted, then
166   /// returned.
167   const ElemTy &getOrInsertLeaderValue(const ElemTy &V) const {
168     member_iterator MI = findLeader(insert(V));
169     assert(MI != member_end() && "Value is not in the set!");
170     return *MI;
171   }
172
173   /// getNumClasses - Return the number of equivalence classes in this set.
174   /// Note that this is a linear time operation.
175   unsigned getNumClasses() const {
176     unsigned NC = 0;
177     for (iterator I = begin(), E = end(); I != E; ++I)
178       if (I->isLeader()) ++NC;
179     return NC;
180   }
181
182
183   //===--------------------------------------------------------------------===//
184   // Mutation methods
185
186   /// insert - Insert a new value into the union/find set, ignoring the request
187   /// if the value already exists.
188   iterator insert(const ElemTy &Data) {
189     return TheMapping.insert(Data).first;
190   }
191
192   /// findLeader - Given a value in the set, return a member iterator for the
193   /// equivalence class it is in.  This does the path-compression part that
194   /// makes union-find "union findy".  This returns an end iterator if the value
195   /// is not in the equivalence class.
196   ///
197   member_iterator findLeader(iterator I) const {
198     if (I == TheMapping.end()) return member_end();
199     return member_iterator(I->getLeader());
200   }
201   member_iterator findLeader(const ElemTy &V) const {
202     return findLeader(TheMapping.find(V));
203   }
204
205
206   /// union - Merge the two equivalence sets for the specified values, inserting
207   /// them if they do not already exist in the equivalence set.
208   member_iterator unionSets(const ElemTy &V1, const ElemTy &V2) {
209     iterator V1I = insert(V1), V2I = insert(V2);
210     return unionSets(findLeader(V1I), findLeader(V2I));
211   }
212   member_iterator unionSets(member_iterator L1, member_iterator L2) {
213     assert(L1 != member_end() && L2 != member_end() && "Illegal inputs!");
214     if (L1 == L2) return L1;   // Unifying the same two sets, noop.
215
216     // Otherwise, this is a real union operation.  Set the end of the L1 list to
217     // point to the L2 leader node.
218     const ECValue &L1LV = *L1.Node, &L2LV = *L2.Node;
219     L1LV.getEndOfList()->setNext(&L2LV);
220     
221     // Update L1LV's end of list pointer.
222     L1LV.Leader = L2LV.getEndOfList();
223
224     // Clear L2's leader flag:
225     L2LV.Next = L2LV.getNext();
226
227     // L2's leader is now L1.
228     L2LV.Leader = &L1LV;
229     return L1;
230   }
231
232   class member_iterator : public forward_iterator<ElemTy, ptrdiff_t> {
233     typedef forward_iterator<const ElemTy, ptrdiff_t> super;
234     const ECValue *Node;
235     friend class EquivalenceClasses;
236   public:
237     typedef size_t size_type;
238     typedef typename super::pointer pointer;
239     typedef typename super::reference reference;
240
241     explicit member_iterator() {}
242     explicit member_iterator(const ECValue *N) : Node(N) {}
243     member_iterator(const member_iterator &I) : Node(I.Node) {}
244
245     reference operator*() const {
246       assert(Node != 0 && "Dereferencing end()!");
247       return Node->getData();
248     }
249     reference operator->() const { return operator*(); }
250
251     member_iterator &operator++() {
252       assert(Node != 0 && "++'d off the end of the list!");
253       Node = Node->getNext();
254       return *this;
255     }
256
257     member_iterator operator++(int) {    // postincrement operators.
258       member_iterator tmp = *this;
259       ++*this;
260       return tmp;
261     }
262
263     bool operator==(const member_iterator &RHS) const {
264       return Node == RHS.Node;
265     }
266     bool operator!=(const member_iterator &RHS) const {
267       return Node != RHS.Node;
268     }
269   };
270 };
271
272 } // End llvm namespace
273
274 #endif