Fix Lexer::ComputePreamble when MaxLines parameter is non-zero.

The function was only counting lines that included tokens and not empty lines,
but MaxLines (mainly initiated to the line where the code-completion point resides)
is a count of overall lines (even empty ones).

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

2 files changed, 20 insertions, 4 deletions

diff --git a/lib/Frontend/ASTUnit.cpp b/lib/Frontend/ASTUnit.cpp
index b8551e9a42..190bc49a85 100644
--- a/lib/Frontend/ASTUnit.cpp
+++ b/lib/Frontend/ASTUnit.cpp
@@ -2254,7 +2254,8 @@ void ASTUnit::CodeComplete(StringRef File, unsigned Line, unsigned Column,
     llvm::sys::PathWithStatus MainPath(OriginalSourceFile);
     if (const FileStatus *CompleteFileStatus = CompleteFilePath.getFileStatus())
       if (const FileStatus *MainStatus = MainPath.getFileStatus())
-        if (CompleteFileStatus->getUniqueID() == MainStatus->getUniqueID())
+        if (CompleteFileStatus->getUniqueID() == MainStatus->getUniqueID() &&
+            Line > 1)
           OverrideMainBuffer
             = getMainBufferWithPrecompiledPreamble(*CCInvocation, false, 
                                                    Line - 1);
diff --git a/lib/Lex/Lexer.cpp b/lib/Lex/Lexer.cpp
index 64b8744922..e9db93ee43 100644
--- a/lib/Lex/Lexer.cpp
+++ b/lib/Lex/Lexer.cpp
@@ -519,7 +519,22 @@ Lexer::ComputePreamble(const llvm::MemoryBuffer *Buffer,
   Token TheTok;
   Token IfStartTok;
   unsigned IfCount = 0;
-  unsigned Line = 0;
+
+  unsigned MaxLineOffset = 0;
+  if (MaxLines) {
+    const char *CurPtr = Buffer->getBufferStart();
+    unsigned CurLine = 0;
+    while (CurPtr != Buffer->getBufferEnd()) {
+      char ch = *CurPtr++;
+      if (ch == '\n') {
+        ++CurLine;
+        if (CurLine == MaxLines)
+          break;
+      }
+    }
+    if (CurPtr != Buffer->getBufferEnd())
+      MaxLineOffset = CurPtr - Buffer->getBufferStart();
+  }
 
   do {
     TheLexer.LexFromRawLexer(TheTok);
@@ -543,11 +558,11 @@ Lexer::ComputePreamble(const llvm::MemoryBuffer *Buffer,
     
     // Keep track of the # of lines in the preamble.
     if (TheTok.isAtStartOfLine()) {
-      ++Line;
+      unsigned TokOffset = TheTok.getLocation().getRawEncoding() - StartOffset;
 
       // If we were asked to limit the number of lines in the preamble,
       // and we're about to exceed that limit, we're done.
-      if (MaxLines && Line >= MaxLines)
+      if (MaxLineOffset && TokOffset >= MaxLineOffset)
         break;
     }