Correct the source range returned from preprocessor callbacks.

This adjusts the source range passed in to the preprocessor callbacks to only include the condition range itself, rather than all of the conditionally skipped tokens.

git-svn-id: https://llvm.org/svn/llvm-project/cfe/trunk@350891 91177308-0d34-0410-b5e6-96231b3b80d8

4 files changed, 170 insertions, 38 deletions

diff --git a/include/clang/Lex/Preprocessor.h b/include/clang/Lex/Preprocessor.h
index 64ddb5307f..36cb718570 100644
--- a/include/clang/Lex/Preprocessor.h
+++ b/include/clang/Lex/Preprocessor.h
@@ -1816,8 +1816,8 @@ public:
   void CheckEndOfDirective(const char *DirType, bool EnableMacros = false);
 
   /// Read and discard all tokens remaining on the current line until
-  /// the tok::eod token is found.
-  void DiscardUntilEndOfDirective();
+  /// the tok::eod token is found. Returns the range of the skipped tokens.
+  SourceRange DiscardUntilEndOfDirective();
 
   /// Returns true if the preprocessor has seen a use of
   /// __DATE__ or __TIME__ in the file so far.
@@ -1982,6 +1982,9 @@ private:
 
     /// True if the expression contained identifiers that were undefined.
     bool IncludedUndefinedIds;
+
+    /// The source range for the expression.
+    SourceRange ExprRange;
   };
 
   /// Evaluate an integer constant expression that may occur after a
diff --git a/lib/Lex/PPDirectives.cpp b/lib/Lex/PPDirectives.cpp
index 15fc086f8a..c6baa99e90 100644
--- a/lib/Lex/PPDirectives.cpp
+++ b/lib/Lex/PPDirectives.cpp
@@ -79,12 +79,18 @@ Preprocessor::AllocateVisibilityMacroDirective(SourceLocation Loc,
 
 /// Read and discard all tokens remaining on the current line until
 /// the tok::eod token is found.
-void Preprocessor::DiscardUntilEndOfDirective() {
+SourceRange Preprocessor::DiscardUntilEndOfDirective() {
   Token Tmp;
-  do {
-    LexUnexpandedToken(Tmp);
+  SourceRange Res;
+
+  LexUnexpandedToken(Tmp);
+  Res.setBegin(Tmp.getLocation());
+  while (Tmp.isNot(tok::eod)) {
     assert(Tmp.isNot(tok::eof) && "EOF seen while discarding directive tokens");
-  } while (Tmp.isNot(tok::eod));
+    LexUnexpandedToken(Tmp);
+  }
+  Res.setEnd(Tmp.getLocation());
+  return Res;
 }
 
 /// Enumerates possible cases of #define/#undef a reserved identifier.
@@ -538,19 +544,19 @@ void Preprocessor::SkipExcludedConditionalBlock(SourceLocation HashTokenLoc,
         if (CondInfo.WasSkipping || CondInfo.FoundNonSkip) {
           DiscardUntilEndOfDirective();
         } else {
-          const SourceLocation CondBegin = CurPPLexer->getSourceLocation();
           // Restore the value of LexingRawMode so that identifiers are
           // looked up, etc, inside the #elif expression.
           assert(CurPPLexer->LexingRawMode && "We have to be skipping here!");
           CurPPLexer->LexingRawMode = false;
           IdentifierInfo *IfNDefMacro = nullptr;
-          const bool CondValue = EvaluateDirectiveExpression(IfNDefMacro).Conditional;
+          DirectiveEvalResult DER = EvaluateDirectiveExpression(IfNDefMacro);
+          const bool CondValue = DER.Conditional;
           CurPPLexer->LexingRawMode = true;
           if (Callbacks) {
-            const SourceLocation CondEnd = CurPPLexer->getSourceLocation();
-            Callbacks->Elif(Tok.getLocation(),
-                            SourceRange(CondBegin, CondEnd),
-                            (CondValue ? PPCallbacks::CVK_True : PPCallbacks::CVK_False), CondInfo.IfLoc);
+            Callbacks->Elif(
+                Tok.getLocation(), DER.ExprRange,
+                (CondValue ? PPCallbacks::CVK_True : PPCallbacks::CVK_False),
+                CondInfo.IfLoc);
           }
           // If this condition is true, enter it!
           if (CondValue) {
@@ -1116,19 +1122,24 @@ void Preprocessor::HandleLineDirective() {
     ; // ok
   else if (StrTok.isNot(tok::string_literal)) {
     Diag(StrTok, diag::err_pp_line_invalid_filename);
-    return DiscardUntilEndOfDirective();
+    DiscardUntilEndOfDirective();
+    return;
   } else if (StrTok.hasUDSuffix()) {
     Diag(StrTok, diag::err_invalid_string_udl);
-    return DiscardUntilEndOfDirective();
+    DiscardUntilEndOfDirective();
+    return;
   } else {
     // Parse and validate the string, converting it into a unique ID.
     StringLiteralParser Literal(StrTok, *this);
     assert(Literal.isAscii() && "Didn't allow wide strings in");
-    if (Literal.hadError)
-      return DiscardUntilEndOfDirective();
+    if (Literal.hadError) {
+      DiscardUntilEndOfDirective();
+      return;
+    }
     if (Literal.Pascal) {
       Diag(StrTok, diag::err_pp_linemarker_invalid_filename);
-      return DiscardUntilEndOfDirective();
+      DiscardUntilEndOfDirective();
+      return;
     }
     FilenameID = SourceMgr.getLineTableFilenameID(Literal.GetString());
 
@@ -1261,19 +1272,24 @@ void Preprocessor::HandleDigitDirective(Token &DigitTok) {
     FileKind = SourceMgr.getFileCharacteristic(DigitTok.getLocation());
   } else if (StrTok.isNot(tok::string_literal)) {
     Diag(StrTok, diag::err_pp_linemarker_invalid_filename);
-    return DiscardUntilEndOfDirective();
+    DiscardUntilEndOfDirective();
+    return;
   } else if (StrTok.hasUDSuffix()) {
     Diag(StrTok, diag::err_invalid_string_udl);
-    return DiscardUntilEndOfDirective();
+    DiscardUntilEndOfDirective();
+    return;
   } else {
     // Parse and validate the string, converting it into a unique ID.
     StringLiteralParser Literal(StrTok, *this);
     assert(Literal.isAscii() && "Didn't allow wide strings in");
-    if (Literal.hadError)
-      return DiscardUntilEndOfDirective();
+    if (Literal.hadError) {
+      DiscardUntilEndOfDirective();
+      return;
+    }
     if (Literal.Pascal) {
       Diag(StrTok, diag::err_pp_linemarker_invalid_filename);
-      return DiscardUntilEndOfDirective();
+      DiscardUntilEndOfDirective();
+      return;
     }
     FilenameID = SourceMgr.getLineTableFilenameID(Literal.GetString());
 
@@ -1343,7 +1359,8 @@ void Preprocessor::HandleIdentSCCSDirective(Token &Tok) {
 
   if (StrTok.hasUDSuffix()) {
     Diag(StrTok, diag::err_invalid_string_udl);
-    return DiscardUntilEndOfDirective();
+    DiscardUntilEndOfDirective();
+    return;
   }
 
   // Verify that there is nothing after the string, other than EOD.
@@ -2783,10 +2800,8 @@ void Preprocessor::HandleIfDirective(Token &IfToken,
 
   // Parse and evaluate the conditional expression.
   IdentifierInfo *IfNDefMacro = nullptr;
-  const SourceLocation ConditionalBegin = CurPPLexer->getSourceLocation();
   const DirectiveEvalResult DER = EvaluateDirectiveExpression(IfNDefMacro);
   const bool ConditionalTrue = DER.Conditional;
-  const SourceLocation ConditionalEnd = CurPPLexer->getSourceLocation();
 
   // If this condition is equivalent to #ifndef X, and if this is the first
   // directive seen, handle it for the multiple-include optimization.
@@ -2799,9 +2814,9 @@ void Preprocessor::HandleIfDirective(Token &IfToken,
   }
 
   if (Callbacks)
-    Callbacks->If(IfToken.getLocation(),
-                  SourceRange(ConditionalBegin, ConditionalEnd),
-                  (ConditionalTrue ? PPCallbacks::CVK_True : PPCallbacks::CVK_False));
+    Callbacks->If(
+        IfToken.getLocation(), DER.ExprRange,
+        (ConditionalTrue ? PPCallbacks::CVK_True : PPCallbacks::CVK_False));
 
   // Should we include the stuff contained by this directive?
   if (PPOpts->SingleFileParseMode && DER.IncludedUndefinedIds) {
@@ -2894,9 +2909,7 @@ void Preprocessor::HandleElifDirective(Token &ElifToken,
   // #elif directive in a non-skipping conditional... start skipping.
   // We don't care what the condition is, because we will always skip it (since
   // the block immediately before it was included).
-  const SourceLocation ConditionalBegin = CurPPLexer->getSourceLocation();
-  DiscardUntilEndOfDirective();
-  const SourceLocation ConditionalEnd = CurPPLexer->getSourceLocation();
+  SourceRange ConditionRange = DiscardUntilEndOfDirective();
 
   PPConditionalInfo CI;
   if (CurPPLexer->popConditionalLevel(CI)) {
@@ -2912,8 +2925,7 @@ void Preprocessor::HandleElifDirective(Token &ElifToken,
   if (CI.FoundElse) Diag(ElifToken, diag::pp_err_elif_after_else);
 
   if (Callbacks)
-    Callbacks->Elif(ElifToken.getLocation(),
-                    SourceRange(ConditionalBegin, ConditionalEnd),
+    Callbacks->Elif(ElifToken.getLocation(), ConditionRange,
                     PPCallbacks::CVK_NotEvaluated, CI.IfLoc);
 
   if (PPOpts->SingleFileParseMode && !CI.FoundNonSkip) {
diff --git a/lib/Lex/PPExpressions.cpp b/lib/Lex/PPExpressions.cpp
index ac01efad9b..a60b7c97fc 100644
--- a/lib/Lex/PPExpressions.cpp
+++ b/lib/Lex/PPExpressions.cpp
@@ -152,8 +152,8 @@ static bool EvaluateDefined(PPValue &Result, Token &PeekTok, DefinedTracker &DT,
       return true;
     }
     // Consume the ).
-    Result.setEnd(PeekTok.getLocation());
     PP.LexNonComment(PeekTok);
+    Result.setEnd(PeekTok.getLocation());
   } else {
     // Consume identifier.
     Result.setEnd(PeekTok.getLocation());
@@ -863,7 +863,7 @@ Preprocessor::EvaluateDirectiveExpression(IdentifierInfo *&IfNDefMacro) {
 
     // Restore 'DisableMacroExpansion'.
     DisableMacroExpansion = DisableMacroExpansionAtStartOfDirective;
-    return {ResVal.Val != 0, DT.IncludedUndefinedIds};
+    return {ResVal.Val != 0, DT.IncludedUndefinedIds, ResVal.getRange()};
   }
 
   // Otherwise, we must have a binary operator (e.g. "#if 1 < 2"), so parse the
@@ -876,7 +876,7 @@ Preprocessor::EvaluateDirectiveExpression(IdentifierInfo *&IfNDefMacro) {
 
     // Restore 'DisableMacroExpansion'.
     DisableMacroExpansion = DisableMacroExpansionAtStartOfDirective;
-    return {false, DT.IncludedUndefinedIds};
+    return {false, DT.IncludedUndefinedIds, ResVal.getRange()};
   }
 
   // If we aren't at the tok::eod token, something bad happened, like an extra
@@ -888,5 +888,5 @@ Preprocessor::EvaluateDirectiveExpression(IdentifierInfo *&IfNDefMacro) {
 
   // Restore 'DisableMacroExpansion'.
   DisableMacroExpansion = DisableMacroExpansionAtStartOfDirective;
-  return {ResVal.Val != 0, DT.IncludedUndefinedIds};
+  return {ResVal.Val != 0, DT.IncludedUndefinedIds, ResVal.getRange()};
 }
diff --git a/unittests/Lex/PPCallbacksTest.cpp b/unittests/Lex/PPCallbacksTest.cpp
index d0ff45ec7c..29d5268921 100644
--- a/unittests/Lex/PPCallbacksTest.cpp
+++ b/unittests/Lex/PPCallbacksTest.cpp
@@ -65,6 +65,29 @@ public:
   SrcMgr::CharacteristicKind FileType;
 };
 
+class CondDirectiveCallbacks : public PPCallbacks {
+public:
+  struct Result {
+    SourceRange ConditionRange;
+    ConditionValueKind ConditionValue;
+
+    Result(SourceRange R, ConditionValueKind K)
+        : ConditionRange(R), ConditionValue(K) {}
+  };
+
+  std::vector<Result> Results;
+
+  void If(SourceLocation Loc, SourceRange ConditionRange,
+          ConditionValueKind ConditionValue) override {
+    Results.emplace_back(ConditionRange, ConditionValue);
+  }
+
+  void Elif(SourceLocation Loc, SourceRange ConditionRange,
+            ConditionValueKind ConditionValue, SourceLocation IfLoc) override {
+    Results.emplace_back(ConditionRange, ConditionValue);
+  }
+};
+
 // Stub to collect data from PragmaOpenCLExtension callbacks.
 class PragmaOpenCLExtensionCallbacks : public PPCallbacks {
 public:
@@ -137,6 +160,15 @@ protected:
     return StringRef(B, E - B);
   }
 
+  StringRef GetSourceStringToEnd(CharSourceRange Range) {
+    const char *B = SourceMgr.getCharacterData(Range.getBegin());
+    const char *E = SourceMgr.getCharacterData(Range.getEnd());
+
+    return StringRef(
+        B,
+        E - B + Lexer::MeasureTokenLength(Range.getEnd(), SourceMgr, LangOpts));
+  }
+
   // Run lexer over SourceText and collect FilenameRange from
   // the InclusionDirective callback.
   CharSourceRange InclusionDirectiveFilenameRange(const char *SourceText,
@@ -199,6 +231,36 @@ protected:
     return Callbacks;
   }
 
+  std::vector<CondDirectiveCallbacks::Result>
+  DirectiveExprRange(StringRef SourceText) {
+    TrivialModuleLoader ModLoader;
+    MemoryBufferCache PCMCache;
+    std::unique_ptr<llvm::MemoryBuffer> Buf =
+        llvm::MemoryBuffer::getMemBuffer(SourceText);
+    SourceMgr.setMainFileID(SourceMgr.createFileID(std::move(Buf)));
+    HeaderSearch HeaderInfo(std::make_shared<HeaderSearchOptions>(), SourceMgr,
+                            Diags, LangOpts, Target.get());
+    Preprocessor PP(std::make_shared<PreprocessorOptions>(), Diags, LangOpts,
+                    SourceMgr, PCMCache, HeaderInfo, ModLoader,
+                    /*IILookup =*/nullptr,
+                    /*OwnsHeaderSearch =*/false);
+    PP.Initialize(*Target);
+    auto *Callbacks = new CondDirectiveCallbacks;
+    PP.addPPCallbacks(std::unique_ptr<PPCallbacks>(Callbacks));
+
+    // Lex source text.
+    PP.EnterMainSourceFile();
+
+    while (true) {
+      Token Tok;
+      PP.Lex(Tok);
+      if (Tok.is(tok::eof))
+        break;
+    }
+
+    return Callbacks->Results;
+  }
+
   PragmaOpenCLExtensionCallbacks::CallbackParameters
   PragmaOpenCLExtensionCall(const char *SourceText) {
     LangOptions OpenCLLangOpts;
@@ -368,4 +430,59 @@ TEST_F(PPCallbacksTest, OpenCLExtensionPragmaDisabled) {
   ASSERT_EQ(ExpectedState, Parameters.State);
 }
 
-} // anonoymous namespace
+TEST_F(PPCallbacksTest, DirectiveExprRanges) {
+  const auto &Results1 = DirectiveExprRange("#if FLUZZY_FLOOF\n#endif\n");
+  EXPECT_EQ(Results1.size(), 1);
+  EXPECT_EQ(
+      GetSourceStringToEnd(CharSourceRange(Results1[0].ConditionRange, false)),
+      "FLUZZY_FLOOF");
+
+  const auto &Results2 = DirectiveExprRange("#if 1 + 4 < 7\n#endif\n");
+  EXPECT_EQ(Results2.size(), 1);
+  EXPECT_EQ(
+      GetSourceStringToEnd(CharSourceRange(Results2[0].ConditionRange, false)),
+      "1 + 4 < 7");
+
+  const auto &Results3 = DirectiveExprRange("#if 1 + \\\n  2\n#endif\n");
+  EXPECT_EQ(Results3.size(), 1);
+  EXPECT_EQ(
+      GetSourceStringToEnd(CharSourceRange(Results3[0].ConditionRange, false)),
+      "1 + \\\n  2");
+
+  const auto &Results4 = DirectiveExprRange("#if 0\n#elif FLOOFY\n#endif\n");
+  EXPECT_EQ(Results4.size(), 2);
+  EXPECT_EQ(
+      GetSourceStringToEnd(CharSourceRange(Results4[0].ConditionRange, false)),
+      "0");
+  EXPECT_EQ(
+      GetSourceStringToEnd(CharSourceRange(Results4[1].ConditionRange, false)),
+      "FLOOFY");
+
+  const auto &Results5 = DirectiveExprRange("#if 1\n#elif FLOOFY\n#endif\n");
+  EXPECT_EQ(Results5.size(), 2);
+  EXPECT_EQ(
+      GetSourceStringToEnd(CharSourceRange(Results5[0].ConditionRange, false)),
+      "1");
+  EXPECT_EQ(
+      GetSourceStringToEnd(CharSourceRange(Results5[1].ConditionRange, false)),
+      "FLOOFY");
+
+  const auto &Results6 =
+      DirectiveExprRange("#if defined(FLUZZY_FLOOF)\n#endif\n");
+  EXPECT_EQ(Results6.size(), 1);
+  EXPECT_EQ(
+      GetSourceStringToEnd(CharSourceRange(Results6[0].ConditionRange, false)),
+      "defined(FLUZZY_FLOOF)");
+
+  const auto &Results7 =
+      DirectiveExprRange("#if 1\n#elif defined(FLOOFY)\n#endif\n");
+  EXPECT_EQ(Results7.size(), 2);
+  EXPECT_EQ(
+      GetSourceStringToEnd(CharSourceRange(Results7[0].ConditionRange, false)),
+      "1");
+  EXPECT_EQ(
+      GetSourceStringToEnd(CharSourceRange(Results7[1].ConditionRange, false)),
+      "defined(FLOOFY)");
+}
+
+} // namespace