2 files changed, 83 insertions, 23 deletions

diff --git a/include/clang/Basic/SourceManager.h b/include/clang/Basic/SourceManager.h
index 930fb52873..6a4be46094 100644
--- a/include/clang/Basic/SourceManager.h
+++ b/include/clang/Basic/SourceManager.h
@@ -280,6 +280,57 @@ public:
   /// \brief Read the source location entry with index ID.
   virtual void ReadSLocEntry(unsigned ID) = 0;
 };
+  
+
+/// IsBeforeInTranslationUnitCache - This class holds the cache used by
+/// isBeforeInTranslationUnit.  The cache structure is complex enough to be
+/// worth breaking out of SourceManager.
+class IsBeforeInTranslationUnitCache {
+  /// L/R QueryFID - These are the FID's of the cached query.  If these match up
+  /// with a subsequent query, the result can be reused.
+  FileID LQueryFID, RQueryFID;
+  
+  /// CommonFID - This is the file found in common between the two #include
+  /// traces.  It is the nearest common ancestor of the #include tree.
+  FileID CommonFID;
+  
+  /// L/R CommonOffset - This is the offset of the previous query in CommonFID.
+  /// Usually, this represents the location of the #include for QueryFID, but if
+  /// LQueryFID is a parent of RQueryFID (or vise versa) then these can be a
+  /// random token in the parent.
+  unsigned LCommonOffset, RCommonOffset;
+public:
+  
+  /// isCacheValid - Return true if the currently cached values match up with
+  /// the specified LHS/RHS query.  If not, we can't use the cache.
+  bool isCacheValid(FileID LHS, FileID RHS) const {
+    return LQueryFID == LHS && RQueryFID == RHS;
+  }
+  
+  /// getCachedResult - If the cache is valid, compute the result given the
+  /// specified offsets in the LHS/RHS FID's.
+  bool getCachedResult(unsigned LOffset, unsigned ROffset) const {
+    // If one of the query files is the common file, use the offset.  Otherwise,
+    // use the #include loc in the common file.
+    if (LQueryFID != CommonFID) LOffset = LCommonOffset;
+    if (RQueryFID != CommonFID) ROffset = RCommonOffset;
+    return LOffset < ROffset;
+  }
+  
+  // Set up a new query.
+  void setQueryFIDs(FileID LHS, FileID RHS) {
+    LQueryFID = LHS;
+    RQueryFID = RHS;
+  }
+  
+  void setCommonLoc(FileID commonFID, unsigned lCommonOffset,
+                    unsigned rCommonOffset) {
+    CommonFID = commonFID;
+    LCommonOffset = lCommonOffset;
+    RCommonOffset = rCommonOffset;
+  }
+  
+};
 
 /// SourceManager - This file handles loading and caching of source files into
 /// memory.  This object owns the MemoryBuffer objects for all of the loaded
@@ -347,9 +398,7 @@ class SourceManager {
   mutable unsigned NumLinearScans, NumBinaryProbes;
 
   // Cache results for the isBeforeInTranslationUnit method.
-  mutable FileID LastLFIDForBeforeTUCheck;
-  mutable FileID LastRFIDForBeforeTUCheck;
-  mutable bool   LastResForBeforeTUCheck;
+  mutable IsBeforeInTranslationUnitCache IsBeforeInTUCache;
 
   // SourceManager doesn't support copy construction.
   explicit SourceManager(const SourceManager&);
diff --git a/lib/Basic/SourceManager.cpp b/lib/Basic/SourceManager.cpp
index 3fb52a100b..15bb930366 100644
--- a/lib/Basic/SourceManager.cpp
+++ b/lib/Basic/SourceManager.cpp
@@ -1170,20 +1170,17 @@ bool SourceManager::isBeforeInTranslationUnit(SourceLocation LHS,
   if (LOffs.first == ROffs.first)
     return LOffs.second < ROffs.second;
 
-#if 0
   // If we are comparing a source location with multiple locations in the same
   // file, we get a big win by caching the result.
+  bool Cached = false;
+  bool CachedResult = false;
+  if (IsBeforeInTUCache.isCacheValid(LOffs.first, ROffs.first)) {
+    Cached = true;
+    CachedResult = IsBeforeInTUCache.getCachedResult(LOffs.second,ROffs.second);
+  }
 
-  // FIXME: This caching is wrong, but I don't know enough about this code
-  // to immediately fix it.  There are cases where passing the same input
-  // values to this method causes it to return different results.
-  if (LastLFIDForBeforeTUCheck == LOffs.first &&
-      LastRFIDForBeforeTUCheck == ROffs.first)
-    return LastResForBeforeTUCheck;
-
-  LastLFIDForBeforeTUCheck = LOffs.first;
-  LastRFIDForBeforeTUCheck = ROffs.first;
-#endif
+  // Okay, we missed in the cache, start updating the cache for this query.
+  IsBeforeInTUCache.setQueryFIDs(LOffs.first, ROffs.first);
 
   // "Traverse" the include/instantiation stacks of both locations and try to
   // find a common "ancestor".
@@ -1191,7 +1188,6 @@ bool SourceManager::isBeforeInTranslationUnit(SourceLocation LHS,
   // First we traverse the stack of the right location and check each level
   // against the level of the left location, while collecting all levels in a
   // "stack map".
-
   std::map<FileID, unsigned> ROffsMap;
   ROffsMap[ROffs.first] = ROffs.second;
 
@@ -1208,15 +1204,20 @@ bool SourceManager::isBeforeInTranslationUnit(SourceLocation LHS,
 
     ROffs = getDecomposedLoc(UpperLoc);
 
-    if (LOffs.first == ROffs.first)
-      return LastResForBeforeTUCheck = LOffs.second < ROffs.second;
+    // If we found a common file, cache and return our answer!
+    if (LOffs.first == ROffs.first) {
+      IsBeforeInTUCache.setCommonLoc(LOffs.first, LOffs.second, ROffs.second);
+      bool Result = IsBeforeInTUCache.getCachedResult(LOffs.second,
+                                                      ROffs.second);
+      assert(!Cached || CachedResult == Result);   // Validate Cache.
+      return Result;
+    }
 
     ROffsMap[ROffs.first] = ROffs.second;
   }
 
   // We didn't find a common ancestor. Now traverse the stack of the left
   // location, checking against the stack map of the right location.
-
   while (1) {
     SourceLocation UpperLoc;
     const SrcMgr::SLocEntry &Entry = getSLocEntry(LOffs.first);
@@ -1231,8 +1232,13 @@ bool SourceManager::isBeforeInTranslationUnit(SourceLocation LHS,
     LOffs = getDecomposedLoc(UpperLoc);
 
     std::map<FileID, unsigned>::iterator I = ROffsMap.find(LOffs.first);
-    if (I != ROffsMap.end())
-      return LastResForBeforeTUCheck = LOffs.second < I->second;
+    if (I != ROffsMap.end()) {
+      IsBeforeInTUCache.setCommonLoc(LOffs.first, LOffs.second, I->second);
+
+      bool Result = IsBeforeInTUCache.getCachedResult(LOffs.second, I->second);
+      assert(!Cached || CachedResult == Result);   // Validate Cache.
+      return Result;
+    }
   }
 
   // There is no common ancestor, most probably because one location is in the
@@ -1244,11 +1250,16 @@ bool SourceManager::isBeforeInTranslationUnit(SourceLocation LHS,
   // If exactly one location is a memory buffer, assume it preceeds the other.
   bool LIsMB = !getSLocEntry(LOffs.first).getFile().getContentCache()->Entry;
   bool RIsMB = !getSLocEntry(ROffs.first).getFile().getContentCache()->Entry;
-  if (LIsMB != RIsMB)
-    return LastResForBeforeTUCheck = LIsMB;
+  if (LIsMB != RIsMB) {
+    IsBeforeInTUCache.setQueryFIDs(FileID(), FileID()); // Don't try caching.
+    assert(!Cached);
+    return LIsMB;
+  }
 
   // Otherwise, just assume FileIDs were created in order.
-  return LastResForBeforeTUCheck = (LOffs.first < ROffs.first);
+  IsBeforeInTUCache.setQueryFIDs(FileID(), FileID()); // Don't try caching.
+  assert(!Cached);
+  return LOffs.first < ROffs.first;
 }
 
 /// PrintStats - Print statistics to stderr.