[clang-repl] Fix printing preprocessed tokens and macros (#104964)

1 files changed, 4 insertions, 6 deletions

diff --git a/clang/lib/Frontend/PrintPreprocessedOutput.cpp b/clang/lib/Frontend/PrintPreprocessedOutput.cpp
index 135dca0..383d435 100644
--- a/clang/lib/Frontend/PrintPreprocessedOutput.cpp
+++ b/clang/lib/Frontend/PrintPreprocessedOutput.cpp
@@ -916,8 +916,7 @@ static void PrintPreprocessedTokens(Preprocessor &PP, Token &Tok,
       PP.Lex(Tok);
       continue;
     } else if (Tok.is(tok::annot_repl_input_end)) {
-      PP.Lex(Tok);
-      continue;
+      // Fall through to exit the loop.
     } else if (Tok.is(tok::eod)) {
       // Don't print end of directive tokens, since they are typically newlines
       // that mess up our line tracking. These come from unknown pre-processor
@@ -1025,7 +1024,8 @@ static void PrintPreprocessedTokens(Preprocessor &PP, Token &Tok,
     Callbacks->setEmittedTokensOnThisLine();
     IsStartOfLine = false;
 
-    if (Tok.is(tok::eof)) break;
+    if (Tok.is(tok::eof) || Tok.is(tok::annot_repl_input_end))
+      break;
 
     PP.Lex(Tok);
     // If lexing that token causes us to need to skip future tokens, do so now.
@@ -1048,9 +1048,7 @@ static void DoPrintMacros(Preprocessor &PP, raw_ostream *OS) {
   // the macro table at the end.
   PP.EnterMainSourceFile();
 
-  Token Tok;
-  do PP.Lex(Tok);
-  while (Tok.isNot(tok::eof));
+  PP.LexTokensUntilEOF();
 
   SmallVector<id_macro_pair, 128> MacrosByID;
   for (Preprocessor::macro_iterator I = PP.macro_begin(), E = PP.macro_end();