2 files changed, 99 insertions, 4 deletions

diff --git a/clang/lib/Lex/PreprocessingRecord.cpp b/clang/lib/Lex/PreprocessingRecord.cpp
index c87088d2dec..b8e6152466b 100644
--- a/clang/lib/Lex/PreprocessingRecord.cpp
+++ b/clang/lib/Lex/PreprocessingRecord.cpp
@@ -37,9 +37,14 @@ InclusionDirective::InclusionDirective(PreprocessingRecord &PPRec,
   this->FileName = StringRef(Memory, FileName.size());
 }
 
-PreprocessingRecord::PreprocessingRecord(SourceManager &SM)
-  : SourceMgr(SM), ExternalSource(0)
+PreprocessingRecord::PreprocessingRecord(SourceManager &SM,
+                                         bool RecordConditionalDirectives)
+  : SourceMgr(SM),
+    RecordCondDirectives(RecordConditionalDirectives), CondDirectiveNextIdx(0),
+    ExternalSource(0)
 {
+  if (RecordCondDirectives)
+    CondDirectiveStack.push_back(CondDirectiveNextIdx++);
 }
 
 /// \brief Returns a pair of [Begin, End) iterators of preprocessed entities
@@ -397,6 +402,95 @@ void PreprocessingRecord::InclusionDirective(
   addPreprocessedEntity(ID);
 }
 
+bool PreprocessingRecord::rangeIntersectsConditionalDirective(
+                                                      SourceRange Range) const {
+  if (Range.isInvalid())
+    return false;
+
+  CondDirectiveLocsTy::const_iterator
+    low = std::lower_bound(CondDirectiveLocs.begin(), CondDirectiveLocs.end(),
+                           Range.getBegin(), CondDirectiveLoc::Comp(SourceMgr));
+  if (low == CondDirectiveLocs.end())
+    return false;
+
+  if (SourceMgr.isBeforeInTranslationUnit(Range.getEnd(), low->getLoc()))
+    return false;
+
+  CondDirectiveLocsTy::const_iterator
+    upp = std::upper_bound(low, CondDirectiveLocs.end(),
+                           Range.getEnd(), CondDirectiveLoc::Comp(SourceMgr));
+  unsigned uppIdx;
+  if (upp != CondDirectiveLocs.end())
+    uppIdx = upp->getIdx();
+  else
+    uppIdx = 0;
+
+  return low->getIdx() != uppIdx;
+}
+
+unsigned PreprocessingRecord::findCondDirectiveIdx(SourceLocation Loc) const {
+  if (Loc.isInvalid())
+    return 0;
+
+  CondDirectiveLocsTy::const_iterator
+    low = std::lower_bound(CondDirectiveLocs.begin(), CondDirectiveLocs.end(),
+                           Loc, CondDirectiveLoc::Comp(SourceMgr));
+  if (low == CondDirectiveLocs.end())
+    return 0;
+  return low->getIdx();
+}
+
+void PreprocessingRecord::addCondDirectiveLoc(CondDirectiveLoc DirLoc) {
+  // Ignore directives in system headers.
+  if (SourceMgr.isInSystemHeader(DirLoc.getLoc()))
+    return;
+
+  assert(CondDirectiveLocs.empty() ||
+         SourceMgr.isBeforeInTranslationUnit(CondDirectiveLocs.back().getLoc(),
+                                             DirLoc.getLoc()));
+  CondDirectiveLocs.push_back(DirLoc);
+}
+
+void PreprocessingRecord::If(SourceLocation Loc, SourceRange ConditionRange) {
+  if (RecordCondDirectives) {
+    addCondDirectiveLoc(CondDirectiveLoc(Loc, CondDirectiveStack.back()));
+    CondDirectiveStack.push_back(CondDirectiveNextIdx++);
+  }
+}
+
+void PreprocessingRecord::Ifdef(SourceLocation Loc, const Token &MacroNameTok) {
+  if (RecordCondDirectives) {
+    addCondDirectiveLoc(CondDirectiveLoc(Loc, CondDirectiveStack.back()));
+    CondDirectiveStack.push_back(CondDirectiveNextIdx++);
+  }
+}
+
+void PreprocessingRecord::Ifndef(SourceLocation Loc,const Token &MacroNameTok) {
+  if (RecordCondDirectives) {
+    addCondDirectiveLoc(CondDirectiveLoc(Loc, CondDirectiveStack.back()));
+    CondDirectiveStack.push_back(CondDirectiveNextIdx++);
+  }
+}
+
+void PreprocessingRecord::Elif(SourceLocation Loc, SourceRange ConditionRange,
+                               SourceLocation IfLoc) {
+  if (RecordCondDirectives)
+    addCondDirectiveLoc(CondDirectiveLoc(Loc, CondDirectiveStack.back()));
+}
+
+void PreprocessingRecord::Else(SourceLocation Loc, SourceLocation IfLoc) {
+  if (RecordCondDirectives)
+    addCondDirectiveLoc(CondDirectiveLoc(Loc, CondDirectiveStack.back()));
+}
+
+void PreprocessingRecord::Endif(SourceLocation Loc, SourceLocation IfLoc) {
+  if (RecordCondDirectives) {
+    addCondDirectiveLoc(CondDirectiveLoc(Loc, CondDirectiveStack.back()));
+    assert(!CondDirectiveStack.empty());
+    CondDirectiveStack.pop_back();
+  }
+}
+
 size_t PreprocessingRecord::getTotalMemory() const {
   return BumpAlloc.getTotalMemory()
     + llvm::capacity_in_bytes(MacroDefinitions)
diff --git a/clang/lib/Lex/Preprocessor.cpp b/clang/lib/Lex/Preprocessor.cpp
index b6ea65de499..06914c7cdcc 100644
--- a/clang/lib/Lex/Preprocessor.cpp
+++ b/clang/lib/Lex/Preprocessor.cpp
@@ -654,10 +654,11 @@ CommentHandler::~CommentHandler() { }
 
 CodeCompletionHandler::~CodeCompletionHandler() { }
 
-void Preprocessor::createPreprocessingRecord() {
+void Preprocessor::createPreprocessingRecord(bool RecordConditionalDirectives) {
   if (Record)
     return;
   
-  Record = new PreprocessingRecord(getSourceManager());
+  Record = new PreprocessingRecord(getSourceManager(),
+                                   RecordConditionalDirectives);
   addPPCallbacks(Record);
 }