Support code-completion for C++ inline methods and ObjC buffering methods.

Previously we would cut off the source file buffer at the code-completion
point; this impeded code-completion inside C++ inline methods and,
recently, with buffering ObjC methods.

Have the code-completion inserted into the source buffer so that it can
be buffered along with a method body. When we actually hit the code-completion
point the cut-off lexing or parsing.

Fixes rdar://10056932&8319466

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

7 files changed, 130 insertions, 89 deletions

diff --git a/lib/Lex/Lexer.cpp b/lib/Lex/Lexer.cpp
index e9db93ee43..4de6ce7ad5 100644
--- a/lib/Lex/Lexer.cpp
+++ b/lib/Lex/Lexer.cpp
@@ -1384,16 +1384,21 @@ void Lexer::LexStringLiteral(Token &Result, const char *CurPtr,
     
     if (C == '\n' || C == '\r' ||             // Newline.
         (C == 0 && CurPtr-1 == BufferEnd)) {  // End of file.
-      if (C == 0 && PP && PP->isCodeCompletionFile(FileLoc))
-        PP->CodeCompleteNaturalLanguage();
-      else if (!isLexingRawMode() && !Features.AsmPreprocessor)
+      if (!isLexingRawMode() && !Features.AsmPreprocessor)
         Diag(BufferPtr, diag::warn_unterminated_string);
       FormTokenWithChars(Result, CurPtr-1, tok::unknown);
       return;
     }
     
-    if (C == 0)
+    if (C == 0) {
+      if (isCodeCompletionPoint(CurPtr-1)) {
+        PP->CodeCompleteNaturalLanguage();
+        FormTokenWithChars(Result, CurPtr-1, tok::unknown);
+        return cutOffLexing();
+      }
+
       NulCharacter = CurPtr-1;
+    }
     C = getAndAdvanceChar(CurPtr, Result);
   }
 
@@ -1491,7 +1496,8 @@ void Lexer::LexAngledStringLiteral(Token &Result, const char *CurPtr) {
       // Skip the escaped character.
       C = getAndAdvanceChar(CurPtr, Result);
     } else if (C == '\n' || C == '\r' ||             // Newline.
-               (C == 0 && CurPtr-1 == BufferEnd)) {  // End of file.
+               (C == 0 && (CurPtr-1 == BufferEnd ||  // End of file.
+                           isCodeCompletionPoint(CurPtr-1)))) {
       // If the filename is unterminated, then it must just be a lone <
       // character.  Return this as such.
       FormTokenWithChars(Result, AfterLessPos, tok::less);
@@ -1535,13 +1541,17 @@ void Lexer::LexCharConstant(Token &Result, const char *CurPtr,
       C = getAndAdvanceChar(CurPtr, Result);
     } else if (C == '\n' || C == '\r' ||             // Newline.
                (C == 0 && CurPtr-1 == BufferEnd)) {  // End of file.
-      if (C == 0 && PP && PP->isCodeCompletionFile(FileLoc))
-        PP->CodeCompleteNaturalLanguage();
-      else if (!isLexingRawMode() && !Features.AsmPreprocessor)
+      if (!isLexingRawMode() && !Features.AsmPreprocessor)
         Diag(BufferPtr, diag::warn_unterminated_char);
       FormTokenWithChars(Result, CurPtr-1, tok::unknown);
       return;
     } else if (C == 0) {
+      if (isCodeCompletionPoint(CurPtr-1)) {
+        PP->CodeCompleteNaturalLanguage();
+        FormTokenWithChars(Result, CurPtr-1, tok::unknown);
+        return cutOffLexing();
+      }
+
       NulCharacter = CurPtr-1;
     }
     C = getAndAdvanceChar(CurPtr, Result);
@@ -1682,12 +1692,16 @@ bool Lexer::SkipBCPLComment(Token &Result, const char *CurPtr) {
     }
 
     if (CurPtr == BufferEnd+1) { 
-      if (PP && PP->isCodeCompletionFile(FileLoc))
-        PP->CodeCompleteNaturalLanguage();
-
       --CurPtr; 
       break; 
     }
+
+    if (C == '\0' && isCodeCompletionPoint(CurPtr-1)) {
+      PP->CodeCompleteNaturalLanguage();
+      cutOffLexing();
+      return false;
+    }
+
   } while (C != '\n' && C != '\r');
 
   // Found but did not consume the newline.  Notify comment handlers about the
@@ -1842,8 +1856,7 @@ bool Lexer::SkipBlockComment(Token &Result, const char *CurPtr) {
   unsigned char C = getCharAndSize(CurPtr, CharSize);
   CurPtr += CharSize;
   if (C == 0 && CurPtr == BufferEnd+1) {
-    if (!isLexingRawMode() &&
-        !PP->isCodeCompletionFile(FileLoc))
+    if (!isLexingRawMode())
       Diag(BufferPtr, diag::err_unterminated_block_comment);
     --CurPtr;
 
@@ -1866,7 +1879,10 @@ bool Lexer::SkipBlockComment(Token &Result, const char *CurPtr) {
   while (1) {
     // Skip over all non-interesting characters until we find end of buffer or a
     // (probably ending) '/' character.
-    if (CurPtr + 24 < BufferEnd) {
+    if (CurPtr + 24 < BufferEnd &&
+        // If there is a code-completion point avoid the fast scan because it
+        // doesn't check for '\0'.
+        !(PP && PP->getCodeCompletionFileLoc() == FileLoc)) {
       // While not aligned to a 16-byte boundary.
       while (C != '/' && ((intptr_t)CurPtr & 0x0F) != 0)
         C = *CurPtr++;
@@ -1926,9 +1942,7 @@ bool Lexer::SkipBlockComment(Token &Result, const char *CurPtr) {
           Diag(CurPtr-1, diag::warn_nested_block_comment);
       }
     } else if (C == 0 && CurPtr == BufferEnd+1) {
-      if (PP && PP->isCodeCompletionFile(FileLoc))
-        PP->CodeCompleteNaturalLanguage();
-      else if (!isLexingRawMode())
+      if (!isLexingRawMode())
         Diag(BufferPtr, diag::err_unterminated_block_comment);
       // Note: the user probably forgot a */.  We could continue immediately
       // after the /*, but this would involve lexing a lot of what really is the
@@ -1944,7 +1958,12 @@ bool Lexer::SkipBlockComment(Token &Result, const char *CurPtr) {
 
       BufferPtr = CurPtr;
       return false;
+    } else if (C == '\0' && isCodeCompletionPoint(CurPtr-1)) {
+      PP->CodeCompleteNaturalLanguage();
+      cutOffLexing();
+      return false;
     }
+
     C = *CurPtr++;
   }
 
@@ -2001,6 +2020,12 @@ std::string Lexer::ReadToEndOfLine() {
     case 0:  // Null.
       // Found end of file?
       if (CurPtr-1 != BufferEnd) {
+        if (isCodeCompletionPoint(CurPtr-1)) {
+          PP->CodeCompleteNaturalLanguage();
+          cutOffLexing();
+          return Result;
+        }
+
         // Nope, normal character, continue.
         Result += Char;
         break;
@@ -2015,8 +2040,8 @@ std::string Lexer::ReadToEndOfLine() {
       // Next, lex the character, which should handle the EOD transition.
       Lex(Tmp);
       if (Tmp.is(tok::code_completion)) {
-        if (PP && PP->getCodeCompletionHandler())
-          PP->getCodeCompletionHandler()->CodeCompleteNaturalLanguage();
+        if (PP)
+          PP->CodeCompleteNaturalLanguage();
         Lex(Tmp);
       }
       assert(Tmp.is(tok::eod) && "Unexpected token!");
@@ -2032,22 +2057,6 @@ std::string Lexer::ReadToEndOfLine() {
 /// This returns true if Result contains a token, false if PP.Lex should be
 /// called again.
 bool Lexer::LexEndOfFile(Token &Result, const char *CurPtr) {
-  // Check if we are performing code completion.
-  if (PP && PP->isCodeCompletionFile(FileLoc)) {
-    // We're at the end of the file, but we've been asked to consider 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.
-    PP->SetCodeCompletionPoint(0, 0, 0);
-    
-    // Silence any diagnostics that occur once we hit the code-completion point.
-    PP->getDiagnostics().setSuppressAllDiagnostics(true);
-    return true;
-  }
-
   // If we hit the end of the file while parsing a preprocessor directive,
   // end the preprocessor directive first.  The next token returned will
   // then be the end of file.
@@ -2075,7 +2084,7 @@ bool Lexer::LexEndOfFile(Token &Result, const char *CurPtr) {
 
   // If we are in a #if directive, emit an error.
   while (!ConditionalStack.empty()) {
-    if (!PP->isCodeCompletionFile(FileLoc))
+    if (PP->getCodeCompletionFileLoc() != FileLoc)
       PP->Diag(ConditionalStack.back().IfLoc,
                diag::err_pp_unterminated_conditional);
     ConditionalStack.pop_back();
@@ -2225,6 +2234,15 @@ bool Lexer::HandleEndOfConflictMarker(const char *CurPtr) {
   return false;
 }
 
+bool Lexer::isCodeCompletionPoint(const char *CurPtr) const {
+  if (PP && PP->isCodeCompletionEnabled()) {
+    SourceLocation Loc = FileLoc.getFileLocWithOffset(CurPtr-BufferStart);
+    return Loc == PP->getCodeCompletionLoc();
+  }
+
+  return false;
+}
+
 
 /// LexTokenInternal - This implements a simple C family lexer.  It is an
 /// extremely performance critical piece of code.  This assumes that the buffer
@@ -2277,6 +2295,14 @@ LexNextToken:
       return PPCache->Lex(Result);
     }
 
+    // Check if we are performing code completion.
+    if (isCodeCompletionPoint(CurPtr-1)) {
+      // Return the code-completion token.
+      Result.startToken();
+      FormTokenWithChars(Result, CurPtr, tok::code_completion);
+      return;
+    }
+
     if (!isLexingRawMode())
       Diag(CurPtr-1, diag::null_in_file);
     Result.setFlag(Token::LeadingSpace);
diff --git a/lib/Lex/PPDirectives.cpp b/lib/Lex/PPDirectives.cpp
index c8fcc3e37d..be70ceccfa 100644
--- a/lib/Lex/PPDirectives.cpp
+++ b/lib/Lex/PPDirectives.cpp
@@ -102,8 +102,8 @@ void Preprocessor::ReadMacroName(Token &MacroNameTok, char isDefineUndef) {
   if (MacroNameTok.is(tok::code_completion)) {
     if (CodeComplete)
       CodeComplete->CodeCompleteMacroName(isDefineUndef == 1);
+    setCodeCompletionReached();
     LexUnexpandedToken(MacroNameTok);
-    return;
   }
   
   // Missing macro name?
@@ -214,6 +214,7 @@ void Preprocessor::SkipExcludedConditionalBlock(SourceLocation IfTokenLoc,
     if (Tok.is(tok::code_completion)) {
       if (CodeComplete)
         CodeComplete->CodeCompleteInConditionalExclusion();
+      setCodeCompletionReached();
       continue;
     }
     
@@ -222,7 +223,7 @@ void Preprocessor::SkipExcludedConditionalBlock(SourceLocation IfTokenLoc,
       // Emit errors for each unterminated conditional on the stack, including
       // the current one.
       while (!CurPPLexer->ConditionalStack.empty()) {
-        if (!isCodeCompletionFile(Tok.getLocation()))
+        if (CurLexer->getFileLoc() != CodeCompletionFileLoc)
           Diag(CurPPLexer->ConditionalStack.back().IfLoc,
                diag::err_pp_unterminated_conditional);
         CurPPLexer->ConditionalStack.pop_back();
@@ -581,6 +582,7 @@ TryAgain:
     if (CodeComplete)
       CodeComplete->CodeCompleteDirective(
                                     CurPPLexer->getConditionalStackDepth() > 0);
+    setCodeCompletionReached();
     return;
   case tok::numeric_constant:  // # 7  GNU line marker directive.
     if (getLangOptions().AsmPreprocessor)
@@ -1104,6 +1106,7 @@ bool Preprocessor::ConcatenateIncludeName(
     
     // FIXME: Provide code completion for #includes.
     if (CurTok.is(tok::code_completion)) {
+      setCodeCompletionReached();
       Lex(CurTok);
       continue;
     }
diff --git a/lib/Lex/PPExpressions.cpp b/lib/Lex/PPExpressions.cpp
index 7f00e6ee96..2f810973b1 100644
--- a/lib/Lex/PPExpressions.cpp
+++ b/lib/Lex/PPExpressions.cpp
@@ -96,6 +96,7 @@ static bool EvaluateDefined(PPValue &Result, Token &PeekTok, DefinedTracker &DT,
   if (PeekTok.is(tok::code_completion)) {
     if (PP.getCodeCompletionHandler())
       PP.getCodeCompletionHandler()->CodeCompleteMacroName(false);
+    PP.setCodeCompletionReached();
     PP.LexUnexpandedNonComment(PeekTok);
   }
   
@@ -156,6 +157,7 @@ static bool EvaluateValue(PPValue &Result, Token &PeekTok, DefinedTracker &DT,
   if (PeekTok.is(tok::code_completion)) {
     if (PP.getCodeCompletionHandler())
       PP.getCodeCompletionHandler()->CodeCompletePreprocessorExpression();
+    PP.setCodeCompletionReached();
     PP.LexNonComment(PeekTok);
   }
       
diff --git a/lib/Lex/PPLexerChange.cpp b/lib/Lex/PPLexerChange.cpp
index cb7b279358..758bfc4bb0 100644
--- a/lib/Lex/PPLexerChange.cpp
+++ b/lib/Lex/PPLexerChange.cpp
@@ -89,7 +89,14 @@ void Preprocessor::EnterSourceFile(FileID FID, const DirectoryLookup *CurDir,
       << std::string(SourceMgr.getBufferName(FileStart)) << "";
     return;
   }
-  
+
+  if (isCodeCompletionEnabled() &&
+      SourceMgr.getFileEntryForID(FID) == CodeCompletionFile) {
+    CodeCompletionFileLoc = SourceMgr.getLocForStartOfFile(FID);
+    CodeCompletionLoc =
+        CodeCompletionFileLoc.getFileLocWithOffset(CodeCompletionOffset);
+  }
+
   EnterSourceFileWithLexer(new Lexer(FID, InputFile, *this), CurDir);
   return;
 }
@@ -204,6 +211,25 @@ bool Preprocessor::HandleEndOfFile(Token &Result, bool isEndOfMacro) {
   // If this is a #include'd file, pop it off the include stack and continue
   // lexing the #includer file.
   if (!IncludeMacroStack.empty()) {
+
+    // If we lexed the code-completion file, act as if we reached EOF.
+    if (isCodeCompletionEnabled() && CurPPLexer &&
+        SourceMgr.getLocForStartOfFile(CurPPLexer->getFileID()) ==
+            CodeCompletionFileLoc) {
+      if (CurLexer) {
+        Result.startToken();
+        CurLexer->FormTokenWithChars(Result, CurLexer->BufferEnd, tok::eof);
+        CurLexer.reset();
+      } else {
+        assert(CurPTHLexer && "Got EOF but no current lexer set!");
+        CurPTHLexer->getEOF(Result);
+        CurPTHLexer.reset();
+      }
+
+      CurPPLexer = 0;
+      return true;
+    }
+
     if (!isEndOfMacro && CurPPLexer &&
         SourceMgr.getIncludeLoc(CurPPLexer->getFileID()).isValid()) {
       // Notify SourceManager to record the number of FileIDs that were created
diff --git a/lib/Lex/PPMacroExpansion.cpp b/lib/Lex/PPMacroExpansion.cpp
index 96b729af1a..8bc1351f6e 100644
--- a/lib/Lex/PPMacroExpansion.cpp
+++ b/lib/Lex/PPMacroExpansion.cpp
@@ -353,29 +353,6 @@ MacroArgs *Preprocessor::ReadFunctionLikeMacroArgs(Token &MacroName,
       // an argument value in a macro could expand to ',' or '(' or ')'.
       LexUnexpandedToken(Tok);
 
-      if (Tok.is(tok::code_completion)) {
-        if (CodeComplete)
-          CodeComplete->CodeCompleteMacroArgument(MacroName.getIdentifierInfo(),
-                                                  MI, NumActuals);
-
-        // Add the code-completion token and finish the lexing normally so that
-        // normal code-completion occurs again with the expanded tokens.
-        ArgTokens.push_back(Tok);
-        // Add a marker EOF token to the end of the token list.
-        Token EOFTok;
-        EOFTok.startToken();
-        EOFTok.setKind(tok::eof);
-        EOFTok.setLocation(Tok.getLocation());
-        EOFTok.setLength(0);
-        ArgTokens.push_back(EOFTok);
-        ++NumActuals;
-        // "Fill out" the other arguments.
-        for (; NumActuals < MI->getNumArgs(); ++NumActuals)
-          ArgTokens.push_back(EOFTok);
-        return MacroArgs::create(MI, ArgTokens.data(), ArgTokens.size(),
-                                 /*isVarargsElided=*/false, *this);
-      }
-      
       if (Tok.is(tok::eof) || Tok.is(tok::eod)) { // "#if f(<eof>" & "#if f(\n"
         Diag(MacroName, diag::err_unterm_macro_invoc);
         // Do not lose the EOF/EOD.  Return it to the client.
@@ -410,7 +387,15 @@ MacroArgs *Preprocessor::ReadFunctionLikeMacroArgs(Token &MacroName,
         if (MacroInfo *MI = getMacroInfo(Tok.getIdentifierInfo()))
           if (!MI->isEnabled())
             Tok.setFlag(Token::DisableExpand);
+      } else if (Tok.is(tok::code_completion)) {
+        if (CodeComplete)
+          CodeComplete->CodeCompleteMacroArgument(MacroName.getIdentifierInfo(),
+                                                  MI, NumActuals);
+        // Don't mark that we reached the code-completion point because the
+        // parser is going to handle the token and there will be another
+        // code-completion callback.
       }
+
       ArgTokens.push_back(Tok);
     }
 
diff --git a/lib/Lex/PTHLexer.cpp b/lib/Lex/PTHLexer.cpp
index b61b8f35f8..02af1ec27e 100644
--- a/lib/Lex/PTHLexer.cpp
+++ b/lib/Lex/PTHLexer.cpp
@@ -147,7 +147,7 @@ bool PTHLexer::LexEndOfFile(Token &Result) {
 
   // If we are in a #if directive, emit an error.
   while (!ConditionalStack.empty()) {
-    if (!PP->isCodeCompletionFile(FileStartLoc))
+    if (PP->getCodeCompletionFileLoc() != FileStartLoc)
       PP->Diag(ConditionalStack.back().IfLoc,
                diag::err_pp_unterminated_conditional);
     ConditionalStack.pop_back();
diff --git a/lib/Lex/Preprocessor.cpp b/lib/Lex/Preprocessor.cpp
index e189a9da81..3175d0450e 100644
--- a/lib/Lex/Preprocessor.cpp
+++ b/lib/Lex/Preprocessor.cpp
@@ -59,7 +59,8 @@ Preprocessor::Preprocessor(Diagnostic &diags, LangOptions &opts,
     SourceMgr(SM), HeaderInfo(Headers), TheModuleLoader(TheModuleLoader),
     ExternalSource(0), 
     Identifiers(opts, IILookup), CodeComplete(0),
-    CodeCompletionFile(0), SkipMainFilePreamble(0, true), CurPPLexer(0), 
+    CodeCompletionFile(0), CodeCompletionOffset(0), CodeCompletionReached(0),
+    SkipMainFilePreamble(0, true), CurPPLexer(0), 
     CurDirLookup(0), Callbacks(0), MacroArgCache(0), Record(0), MIChainHead(0),
     MICache(0) 
 {
@@ -73,7 +74,8 @@ Preprocessor::Preprocessor(Diagnostic &diags, LangOptions &opts,
 
 Preprocessor::~Preprocessor() {
   assert(BacktrackPositions.empty() && "EnableBacktrack/Backtrack imbalance!");
-  assert(MacroExpandingLexersStack.empty() && MacroExpandedTokens.empty() &&
+  assert(((MacroExpandingLexersStack.empty() && MacroExpandedTokens.empty()) ||
+          isCodeCompletionReached()) &&
          "Preprocessor::HandleEndOfTokenLexer should have cleared those");
 
   while (!IncludeMacroStack.empty()) {
@@ -270,15 +272,13 @@ Preprocessor::macro_end(bool IncludeExternalMacros) const {
 }
 
 bool Preprocessor::SetCodeCompletionPoint(const FileEntry *File,
-                                          unsigned TruncateAtLine,
-                                          unsigned TruncateAtColumn) {
-  using llvm::MemoryBuffer;
-
-  CodeCompletionFile = File;
+                                          unsigned CompleteLine,
+                                          unsigned CompleteColumn) {
+  assert(File);
+  assert(CompleteLine && CompleteColumn && "Starts from 1:1");
+  assert(!CodeCompletionFile && "Already set");
 
-  // Okay to clear out the code-completion point by passing NULL.
-  if (!CodeCompletionFile)
-    return false;
+  using llvm::MemoryBuffer;
 
   // Load the actual file's contents.
   bool Invalid = false;
@@ -288,7 +288,7 @@ bool Preprocessor::SetCodeCompletionPoint(const FileEntry *File,
 
   // Find the byte position of the truncation point.
   const char *Position = Buffer->getBufferStart();
-  for (unsigned Line = 1; Line < TruncateAtLine; ++Line) {
+  for (unsigned Line = 1; Line < CompleteLine; ++Line) {
     for (; *Position; ++Position) {
       if (*Position != '\r' && *Position != '\n')
         continue;
@@ -302,31 +302,30 @@ bool Preprocessor::SetCodeCompletionPoint(const FileEntry *File,
     }
   }
 
-  Position += TruncateAtColumn - 1;
+  Position += CompleteColumn - 1;
 
-  // Truncate the buffer.
+  // Insert '\0' at the code-completion point.
   if (Position < Buffer->getBufferEnd()) {
-    StringRef Data(Buffer->getBufferStart(),
-                         Position-Buffer->getBufferStart());
-    MemoryBuffer *TruncatedBuffer
-      = MemoryBuffer::getMemBufferCopy(Data, Buffer->getBufferIdentifier());
-    SourceMgr.overrideFileContents(File, TruncatedBuffer);
+    CodeCompletionFile = File;
+    CodeCompletionOffset = Position - Buffer->getBufferStart();
+
+    MemoryBuffer *NewBuffer =
+        MemoryBuffer::getNewUninitMemBuffer(Buffer->getBufferSize() + 1,
+                                            Buffer->getBufferIdentifier());
+    char *NewBuf = (char*)NewBuffer->getBufferStart();
+    char *NewPos = std::copy(Buffer->getBufferStart(), Position, NewBuf);
+    *NewPos = '\0';
+    std::copy(Position, Buffer->getBufferEnd(), NewPos+1);
+    SourceMgr.overrideFileContents(File, NewBuffer);
   }
 
   return false;
 }
 
-bool Preprocessor::isCodeCompletionFile(SourceLocation FileLoc) const {
-  return CodeCompletionFile && FileLoc.isFileID() &&
-    SourceMgr.getFileEntryForID(SourceMgr.getFileID(FileLoc))
-      == CodeCompletionFile;
-}
-
 void Preprocessor::CodeCompleteNaturalLanguage() {
-  SetCodeCompletionPoint(0, 0, 0);
-  getDiagnostics().setSuppressAllDiagnostics(true);
   if (CodeComplete)
     CodeComplete->CodeCompleteNaturalLanguage();
+  setCodeCompletionReached();
 }
 
 /// getSpelling - This method is used to get the spelling of a token into a