From 4e9418ab289b141c70a8b71c9d1bcdc671db04d4 Mon Sep 17 00:00:00 2001
From: Zhongxing Xu <xuzhongxing@gmail.com>
Date: Tue, 2 Feb 2010 05:23:23 +0000
Subject: [PATCH] Add a lookup method to the IntervalMap. The difference from
 the original lookup is that if the lookup key is contained in the key, we
 return the data.

llvm-svn: 95070
---
 llvm/include/llvm/ADT/ImmutableIntervalMap.h | 37 +++++++++++++++++++-
 1 file changed, 36 insertions(+), 1 deletion(-)

diff --git a/llvm/include/llvm/ADT/ImmutableIntervalMap.h b/llvm/include/llvm/ADT/ImmutableIntervalMap.h
index 4754e446a258..74ae5f355387 100644
--- a/llvm/include/llvm/ADT/ImmutableIntervalMap.h
+++ b/llvm/include/llvm/ADT/ImmutableIntervalMap.h
@@ -65,6 +65,13 @@ struct ImutIntervalInfo {
     }
   }
 
+  static bool isContainedIn(key_type_ref K, key_type_ref L) {
+    if (K.getStart() >= L.getStart() && K.getEnd() <= L.getEnd())
+      return true;
+    else
+      return false;
+  }
+
   static void Profile(FoldingSetNodeID &ID, value_type_ref V) {
     ID.AddInteger(V.first.getStart());
     ID.AddInteger(V.first.getEnd());
@@ -85,12 +92,26 @@ class ImutIntervalAVLFactory : public ImutAVLFactory<ImutInfo> {
   typedef typename ImutInfo::data_type_ref  data_type_ref;
 
 public:
-  TreeTy *Add(TreeTy* T, value_type_ref V) {
+  TreeTy *Add(TreeTy *T, value_type_ref V) {
     T = Add_internal(V,T);
     MarkImmutable(T);
     return T;
   }
 
+  TreeTy *Find(TreeTy *T, key_type_ref K) {
+    if (!T)
+      return NULL;
+
+    key_type_ref CurrentKey = ImutInfo::KeyOfValue(Value(T));
+
+    if (ImutInfo::isContainedIn(K, CurrentKey))
+      return T;
+    else if (ImutInfo::isLess(K, CurrentKey))
+      return Find(Left(T), K);
+    else
+      return Find(Right(T), K);
+  }
+
 private:
   TreeTy *Add_internal(value_type_ref V, TreeTy *T) {
     key_type_ref K = ImutInfo::KeyOfValue(V);
@@ -198,7 +219,21 @@ public:
       TreeTy *T = F.Remove(Old.Root, K);
       return ImmutableIntervalMap(F.GetCanonicalTree(T));
     }
+
+    data_type *Lookup(ImmutableIntervalMap M, key_type_ref K) {
+      TreeTy *T = F.Find(M.getRoot(), K);
+      if (T)
+        return &T->getValue().second;
+      else
+        return 0;
+    }
+
   };
+
+private:
+  // For ImmutableIntervalMap, the lookup operation has to be done by the 
+  // factory.
+  data_type* lookup(key_type_ref K) const;
 };
 
 } // end namespace llvm