2 files changed, 22 insertions, 1 deletions
diff --git a/lib/Lex/Lexer.cpp b/lib/Lex/Lexer.cpp
index 23ba6e1ca7..351b63f6bb 100644
--- a/lib/Lex/Lexer.cpp
+++ b/lib/Lex/Lexer.cpp
@@ -70,7 +70,8 @@ void Lexer::InitLexer(const char *BufStart, const char *BufPtr,
          " to simplify lexing!");
 
   Is_PragmaLexer = false;
-
+  IsEofCodeCompletion = false;
+  
   // Start of the file is a start of line.
   IsAtStartOfLine = true;
 
@@ -1309,6 +1310,18 @@ bool Lexer::LexEndOfFile(Token &Result, const char *CurPtr) {
     return true;  // Have a token.
   }
 
+  if (IsEofCodeCompletion) {
+    // We're at the end of the file, but we've been asked to conside the
+    // end of the file to be a code-completion token. Return the
+    // code-completion token.
+    Result.startToken();
+    FormTokenWithChars(Result, CurPtr, tok::code_completion);
+
+    // Only do the eof -> code_completion translation once.
+    IsEofCodeCompletion = false;
+    return true;
+  }
+  
   // If we are in raw mode, return this event as an EOF token.  Let the caller
   // that put us in raw mode handle the event.
   if (isLexingRawMode()) {
diff --git a/lib/Lex/Preprocessor.cpp b/lib/Lex/Preprocessor.cpp
index bfa090a09e..4e522cbb8a 100644
--- a/lib/Lex/Preprocessor.cpp
+++ b/lib/Lex/Preprocessor.cpp
@@ -71,6 +71,7 @@ Preprocessor::Preprocessor(Diagnostic &diags, const LangOptions &opts,
   // Macro expansion is enabled.
   DisableMacroExpansion = false;
   InMacroArgs = false;
+  IsMainFileEofCodeCompletion = false;
   NumCachedTokenLexers = 0;
 
   CachedLexPos = 0;
@@ -368,6 +369,13 @@ void Preprocessor::EnterMainSourceFile() {
   // Enter the main file source buffer.
   EnterSourceFile(MainFileID, 0);
 
+  if (IsMainFileEofCodeCompletion) {
+    // Tell our newly-created lexer that it should treat its end-of-file as
+    // a code-completion token.
+    IsMainFileEofCodeCompletion = false;
+    static_cast<Lexer *>(getCurrentFileLexer())->SetEofIsCodeCompletion();
+  }
+  
   // Tell the header info that the main file was entered.  If the file is later
   // #imported, it won't be re-entered.
   if (const FileEntry *FE = SourceMgr.getFileEntryForID(MainFileID))