Code completion: preserve the AST for the parsed top-level code. This enables

us to find iteration variables while doing code completion in nested top-level
code.

Fixes part of rdar://15199468


Swift SVN r9343

Author: gribozavr

Diff for: include/swift/Parse/Parser.h

@@ -186,22 +186,24 @@ class Parser {
 
   void restoreParserPosition(ParserPosition PP) {
     L->restoreState(PP.LS);
-    PreviousLoc = PP.PreviousLoc;
 
     // We might be at tok::eof now, so ensure that consumeToken() does not
     // assert about lexing past eof.
     Tok.setKind(tok::unknown);
     consumeToken();
+
+    PreviousLoc = PP.PreviousLoc;
   }
 
   void backtrackToPosition(ParserPosition PP) {
     L->backtrackToState(PP.LS);
-    PreviousLoc = PP.PreviousLoc;
 
     // We might be at tok::eof now, so ensure that consumeToken() does not
     // assert about lexing past eof.
     Tok.setKind(tok::unknown);
     consumeToken();
+
+    PreviousLoc = PP.PreviousLoc;
   }
 
   /// RAII object that, when it is destructed, restores the parser and lexer to
@@ -218,7 +220,7 @@ class Parser {
       P.backtrackToPosition(PP);
     }
   };
-  
+
   /// RAII object for managing 'var' patterns. Inside a 'var' pattern, bare
   /// identifiers are parsed as new VarDecls instead of references to existing
   /// ones.
@@ -439,7 +441,8 @@ class Parser {
                  tok SeparatorK, bool OptionalSep, Diag<> ErrorDiag,
                  std::function<bool()> callback);
 
-  void consumeTopLevelDecl(ParserPosition BeginParserPosition);
+  void consumeTopLevelDecl(ParserPosition BeginParserPosition,
+                           TopLevelCodeDecl *TLCD);
 
   ParserStatus parseBraceItems(SmallVectorImpl<ExprStmtOrDecl> &Decls,
                                bool IsTopLevel,
@@ -455,7 +458,8 @@ class Parser {
   static bool isStartOfOperatorDecl(const Token &Tok, const Token &Tok2);
 
   bool parseTopLevel();
-  void consumeDecl(ParserPosition BeginParserPosition, unsigned Flags);
+  void consumeDecl(ParserPosition BeginParserPosition, unsigned Flags,
+                   bool IsTopLevel);
   ParserStatus parseDecl(SmallVectorImpl<Decl*> &Entries, unsigned Flags);
   void parseDeclDelayed();
   enum {

Diff for: include/swift/Parse/PersistentParserState.h

@@ -103,6 +103,9 @@ class PersistentParserState {
                  DeclContext *ParentContext,
                  SourceRange BodyRange, SourceLoc PreviousLoc);
 
+  void delayTopLevel(TopLevelCodeDecl *TLCD, SourceRange BodyRange,
+                     SourceLoc PreviousLoc);
+
   bool hasDelayedDecl() {
     return CodeCompletionDelayedDeclState.get() != nullptr;
   }

Diff for: include/swift/Sema/CodeCompletionTypeChecking.h

@@ -31,7 +31,7 @@ namespace swift {
   /// \brief Typecheck a declaration parsed during code completion.
   ///
   /// \returns true on success, false on error.
-  bool typeCheckCompletionDecl(ASTContext &Ctx, Decl *D);
+  bool typeCheckCompletionDecl(Decl *D);
 
   /// \brief Typecheck an expression parsed during code completion.
   ///
@@ -40,10 +40,14 @@ namespace swift {
                                       Expr *&parsedExpr);
 
   /// Partially typecheck the specified function body.
-  bool typeCheckAbstractFunctionBodyUntil(ASTContext &Ctx,
-                                          AbstractFunctionDecl *AFD,
+  bool typeCheckAbstractFunctionBodyUntil(AbstractFunctionDecl *AFD,
                                           SourceLoc EndTypeCheckLoc);
 
+  /// \brief Typecheck top-level code parsed during code completion.
+  ///
+  /// \returns true on success, false on error.
+  bool typeCheckTopLevelCodeDecl(TopLevelCodeDecl *TLCD);
+
   /// A unique_ptr for LazyResolver that can perform additional cleanup.
   using OwnedResolver = std::unique_ptr<LazyResolver, void(*)(LazyResolver*)>;
 

Diff for: lib/IDE/CodeCompletion.cpp

@@ -426,8 +426,7 @@ class CodeCompletionCallbacksImpl : public CodeCompletionCallbacks,
       typecheckContextImpl(DCToTypeCheck->getParent());
       // Then type check the function itself.
       if (auto *AFD = dyn_cast<AbstractFunctionDecl>(DCToTypeCheck))
-        return typeCheckAbstractFunctionBodyUntil(P.Context, AFD,
-                                                  EndTypeCheckLoc);
+        return typeCheckAbstractFunctionBodyUntil(AFD, EndTypeCheckLoc);
       return false;
     }
     if (DC->getContextKind() == DeclContextKind::NominalTypeDecl) {
@@ -436,7 +435,10 @@ class CodeCompletionCallbacksImpl : public CodeCompletionCallbacks,
       typecheckContextImpl(DC->getParent());
       if (NTD->hasType())
         return true;
-      return typeCheckCompletionDecl(P.Context, cast<NominalTypeDecl>(DC));
+      return typeCheckCompletionDecl(cast<NominalTypeDecl>(DC));
+    }
+    if (DC->getContextKind() == DeclContextKind::TopLevelCodeDecl) {
+      return typeCheckTopLevelCodeDecl(dyn_cast<TopLevelCodeDecl>(DC));
     }
     return true;
   }
@@ -449,7 +451,7 @@ class CodeCompletionCallbacksImpl : public CodeCompletionCallbacks,
   /// \returns true on success, false on failure.
   bool typecheckDelayedParsedDecl() {
     assert(DelayedParsedDecl && "should have a delayed parsed decl");
-    return typeCheckCompletionDecl(P.Context, DelayedParsedDecl);
+    return typeCheckCompletionDecl(DelayedParsedDecl);
   }
 
   /// \returns true on success, false on failure.

Diff for: lib/Parse/ParseDecl.cpp

@@ -453,7 +453,8 @@ bool Parser::isStartOfOperatorDecl(const Token &Tok, const Token &Tok2) {
         || Tok2.isContextualKeyword("infix"));
 }
 
-void Parser::consumeDecl(ParserPosition BeginParserPosition, unsigned Flags) {
+void Parser::consumeDecl(ParserPosition BeginParserPosition, unsigned Flags,
+                         bool IsTopLevel) {
   backtrackToPosition(BeginParserPosition);
   SourceLoc BeginLoc = Tok.getLoc();
   // Consume tokens up to code completion token.
@@ -466,6 +467,12 @@ void Parser::consumeDecl(ParserPosition BeginParserPosition, unsigned Flags) {
   State->delayDecl(PersistentParserState::DelayedDeclKind::Decl, Flags,
                    CurDeclContext, { BeginLoc, EndLoc },
                    BeginParserPosition.PreviousLoc);
+
+  if (IsTopLevel) {
+    // Skip the rest of the file to prevent the parser from constructing the
+    // AST for it.  Forward references are not allowed at the top level.
+    skipUntil(tok::eof);
+  }
 }
 
 /// \brief Parse a single syntactic declaration and return a list of decl
@@ -586,7 +593,7 @@ ParserStatus Parser::parseDecl(SmallVectorImpl<Decl*> &Entries,
   if (Status.hasCodeCompletion() && isCodeCompletionFirstPass() &&
       !CurDeclContext->isModuleContext()) {
     // Only consume non-toplevel decls.
-    consumeDecl(BeginParserPosition, Flags);
+    consumeDecl(BeginParserPosition, Flags, /*IsTopLevel=*/false);
 
     // Pretend that there was no error.
     return makeParserSuccess();

Diff for: lib/Parse/ParseStmt.cpp

@@ -133,7 +133,8 @@ static bool isTerminatorForBraceItemListKind(const Token &Tok,
   }
 }
 
-void Parser::consumeTopLevelDecl(ParserPosition BeginParserPosition) {
+void Parser::consumeTopLevelDecl(ParserPosition BeginParserPosition,
+                                 TopLevelCodeDecl *TLCD) {
   backtrackToPosition(BeginParserPosition);
   SourceLoc BeginLoc = Tok.getLoc();
   // Consume tokens up to code completion token.
@@ -146,9 +147,8 @@ void Parser::consumeTopLevelDecl(ParserPosition BeginParserPosition) {
   // this decl that are past the code completion token.
   skipUntilDeclStmtRBrace(tok::l_brace);
   SourceLoc EndLoc = Tok.getLoc();
-  State->delayDecl(PersistentParserState::DelayedDeclKind::TopLevelCodeDecl, 0,
-                   CurDeclContext, { BeginLoc, EndLoc },
-                   BeginParserPosition.PreviousLoc);
+  State->delayTopLevel(TLCD, { BeginLoc, EndLoc },
+                       BeginParserPosition.PreviousLoc);
 
   // Skip the rest of the file to prevent the parser from constructing the AST
   // for it.  Forward references are not allowed at the top level.
@@ -239,7 +239,7 @@ ParserStatus Parser::parseBraceItems(SmallVectorImpl<ExprStmtOrDecl> &Entries,
         NeedParseErrorRecovery = true;
         if (Status.hasCodeCompletion() && IsTopLevel &&
             isCodeCompletionFirstPass()) {
-          consumeTopLevelDecl(BeginParserPosition);
+          consumeDecl(BeginParserPosition, 0U, IsTopLevel);
           return Status;
         }
       }
@@ -258,7 +258,10 @@ ParserStatus Parser::parseBraceItems(SmallVectorImpl<ExprStmtOrDecl> &Entries,
 
       ParserStatus Status = parseExprOrStmt(Result);
       if (Status.hasCodeCompletion() && isCodeCompletionFirstPass()) {
-        consumeTopLevelDecl(BeginParserPosition);
+        consumeTopLevelDecl(BeginParserPosition, TLCD);
+        auto Brace = BraceStmt::create(Context, StartLoc, {}, Tok.getLoc());
+        TLCD->setBody(Brace);
+        Entries.push_back(TLCD);
         return Status;
       }
       if (Status.isError())
@@ -339,14 +342,20 @@ void Parser::parseTopLevelCodeDeclDelayed() {
   // Rewind to the beginning of the top-level code.
   restoreParserPosition(BeginParserPosition);
 
-  // No need to re-enter the scope: parseBraceItems() will create a scope
-  // anyway.
+  // Re-enter the lexical scope.
+  Scope S(this, DelayedState->takeScope());
 
   // Re-enter the top-level decl context.
-  ContextChange CC(*this, DelayedState->ParentContext);
-
-  SmallVector<ExprStmtOrDecl, 4> Entries;
-  parseBraceItems(Entries, true, BraceItemListKind::TopLevelCode);
+  auto *TLCD = cast<TopLevelCodeDecl>(DelayedState->ParentContext);
+  ContextChange CC(*this, TLCD);
+
+  SourceLoc StartLoc = Tok.getLoc();
+  ExprStmtOrDecl Result;
+  parseExprOrStmt(Result);
+  if (!Result.isNull()) {
+    auto Brace = BraceStmt::create(Context, StartLoc, Result, Tok.getLoc());
+    TLCD->setBody(Brace);
+  }
 }
 
 /// Recover from a 'case' or 'default' outside of a 'switch' by consuming up to

Diff for: lib/Parse/Parser.cpp

@@ -115,7 +115,7 @@ void parseDelayedDecl(TranslationUnit *TU, PersistentParserState &ParserState,
     TheParser.parseDeclDelayed();
     break;
   }
-  
+
   if (CodeCompletion)
     CodeCompletion->doneParsing();
 }

Diff for: lib/Parse/PersistentParserState.cpp

@@ -53,3 +53,10 @@ void PersistentParserState::delayDecl(DelayedDeclKind Kind,
       ScopeInfo.saveCurrentScope()));
 }
 
+void PersistentParserState::delayTopLevel(TopLevelCodeDecl *TLCD,
+                                          SourceRange BodyRange,
+                                          SourceLoc PreviousLoc) {
+  delayDecl(DelayedDeclKind::TopLevelCodeDecl, 0U, TLCD, BodyRange,
+            PreviousLoc);
+}
+

Diff for: lib/Sema/TypeCheckDecl.cpp

@@ -2168,7 +2168,8 @@ void TypeChecker::validateDecl(ValueDecl *D, bool resolveTypeParams) {
       // FIXME: This case is hit when code completion occurs in a function
       // parameter list. Previous parameters are definitely in scope, but
       // we don't really know how to type-check them.
-      assert(isa<AbstractFunctionDecl>(D->getDeclContext()));
+      assert(isa<AbstractFunctionDecl>(D->getDeclContext()) ||
+             isa<TopLevelCodeDecl>(D->getDeclContext()));
       D->setType(ErrorType::get(Context));
     }
     break;

Diff for: lib/Sema/TypeCheckREPL.cpp

@@ -715,7 +715,8 @@ void TypeChecker::processREPLTopLevel(SourceFile &SF, unsigned FirstDecl) {
     SF.Decls.push_back(D);
 
     TopLevelCodeDecl *TLCD = dyn_cast<TopLevelCodeDecl>(D);
-    if (TLCD == 0) continue;
+    if (!TLCD || TLCD->getBody()->getElements().empty())
+      continue;
 
     auto Entry = TLCD->getBody()->getElements()[0];
 

Diff for: lib/Sema/TypeChecker.cpp

@@ -715,10 +715,13 @@ bool swift::performTypeLocChecking(ASTContext &Ctx, TypeLoc &T,
   }
 }
 
-bool swift::typeCheckCompletionDecl(ASTContext &Ctx, Decl *D) {
+bool swift::typeCheckCompletionDecl(Decl *D) {
+  auto &Ctx = D->getASTContext();
+
   // Set up a diagnostics engine that swallows diagnostics.
   DiagnosticEngine Diags(Ctx.SourceMgr);
   TypeChecker TC(Ctx, Diags);
+
   TC.typeCheckDecl(D, true);
   return true;
 }
@@ -736,19 +739,31 @@ bool swift::typeCheckCompletionContextExpr(ASTContext &Ctx, DeclContext *DC,
                     && !parsedExpr->getType()->is<ErrorType>();
 }
 
-bool swift::typeCheckAbstractFunctionBodyUntil(ASTContext &Ctx,
-                                               AbstractFunctionDecl *AFD,
+bool swift::typeCheckAbstractFunctionBodyUntil(AbstractFunctionDecl *AFD,
                                                SourceLoc EndTypeCheckLoc) {
   if (AFD->isInvalid())
     return false;
 
+  auto &Ctx = AFD->getASTContext();
+
   // Set up a diagnostics engine that swallows diagnostics.
   DiagnosticEngine Diags(Ctx.SourceMgr);
 
   TypeChecker TC(Ctx, Diags);
   return !TC.typeCheckAbstractFunctionBodyUntil(AFD, EndTypeCheckLoc);
 }
 
+bool swift::typeCheckTopLevelCodeDecl(TopLevelCodeDecl *TLCD) {
+  auto &Ctx = static_cast<Decl *>(TLCD)->getASTContext();
+
+  // Set up a diagnostics engine that swallows diagnostics.
+  DiagnosticEngine Diags(Ctx.SourceMgr);
+
+  TypeChecker TC(Ctx, Diags);
+  TC.typeCheckTopLevelCodeDecl(TLCD);
+  return true;
+}
+
 static void deleteTypeCheckerAndDiags(LazyResolver *resolver) {
   DiagnosticEngine &diags = static_cast<TypeChecker*>(resolver)->Diags;
   delete resolver;