Allow implicit self in escaping closures when self usage is unlikely to cause cycle

include/swift/AST/Decl.h

@@ -356,12 +356,15 @@ class alignas(1 << DeclAlignInBits) Decl {
     IsStatic : 1
   );
 
-  SWIFT_INLINE_BITFIELD(VarDecl, AbstractStorageDecl, 1+1+1+1+1+1,
+  SWIFT_INLINE_BITFIELD(VarDecl, AbstractStorageDecl, 1+1+1+1+1+1+1,
     /// Encodes whether this is a 'let' binding.
     Introducer : 1,
 
     /// Whether this declaration was an element of a capture list.
     IsCaptureList : 1,
+
+    /// Whether this declaration captures the 'self' param under the same name.
+    IsSelfParamCapture : 1,
 
     /// Whether this vardecl has an initial value bound to it in a way
     /// that isn't represented in the AST with an initializer in the pattern
@@ -4980,6 +4983,12 @@ class VarDecl : public AbstractStorageDecl {
 
   /// Is this an element in a capture list?
   bool isCaptureList() const { return Bits.VarDecl.IsCaptureList; }
+
+  /// Is this a capture of the self param?
+  bool isSelfParamCapture() const { return Bits.VarDecl.IsSelfParamCapture; }
+  void setIsSelfParamCapture(bool IsSelfParamCapture = true) {
+      Bits.VarDecl.IsSelfParamCapture = IsSelfParamCapture;
+  }
 
   /// Return true if this vardecl has an initial value bound to it in a way
   /// that isn't represented in the AST with an initializer in the pattern

include/swift/AST/DiagnosticsSema.def

@@ -3024,11 +3024,20 @@ ERROR(self_assignment_var,none,
 ERROR(self_assignment_prop,none,
       "assigning a property to itself", ())
 ERROR(property_use_in_closure_without_explicit_self,none,
-      "reference to property %0 in closure requires explicit 'self.' to make"
-      " capture semantics explicit", (Identifier))
+      "reference to property %0 in closure requires explicit use of 'self' to"
+      " make capture semantics explicit", (Identifier))
 ERROR(method_call_in_closure_without_explicit_self,none,
-      "call to method %0 in closure requires explicit 'self.' to make"
+      "call to method %0 in closure requires explicit use of 'self' to make"
       " capture semantics explicit", (Identifier))
+NOTE(note_capture_self_explicitly,none,
+      "capture 'self' explicitly to enable implicit 'self' in this closure", ())
+NOTE(note_reference_self_explicitly,none,
+      "reference 'self.' explicitly", ())
+NOTE(note_other_self_capture,none,
+      "variable other than 'self' captured here under the name 'self' does not "
+      "enable implicit 'self'", ())
+NOTE(note_self_captured_weakly,none,
+      "weak capture of 'self' here does not enable implicit 'self'", ())
 ERROR(implicit_use_of_self_in_closure,none,
       "implicit use of 'self' in closure; use 'self.' to make"
       " capture semantics explicit", ())

include/swift/AST/Expr.h

@@ -61,6 +61,7 @@ namespace swift {
   class EnumElementDecl;
   class CallExpr;
   class KeyPathExpr;
+  class CaptureListExpr;
 
 enum class ExprKind : uint8_t {
 #define EXPR(Id, Parent) Id,
@@ -3557,6 +3558,18 @@ class SerializedAbstractClosureExpr : public SerializedLocalDeclContext {
 /// \endcode
 class ClosureExpr : public AbstractClosureExpr {
 
+  /// The range of the brackets of the capture list, if present.
+  SourceRange BracketRange;
+
+  /// The (possibly null) VarDecl captured by this closure with the literal name
+  /// "self". In order to recover this information at the time of name lookup,
+  /// we must be able to access it from the associated DeclContext.
+  /// Because the DeclContext inside a closure is the closure itself (and not
+  /// the CaptureListExpr which would normally maintain this sort of
+  /// information about captured variables), we need to have some way to access
+  /// this information directly on the ClosureExpr.
+  VarDecl *CapturedSelfDecl;
+
   /// The location of the "throws", if present.
   SourceLoc ThrowsLoc;
 
@@ -3573,16 +3586,16 @@ class ClosureExpr : public AbstractClosureExpr {
   /// The body of the closure, along with a bit indicating whether it
   /// was originally just a single expression.
   llvm::PointerIntPair<BraceStmt *, 1, bool> Body;
-
 public:
-  ClosureExpr(ParameterList *params, SourceLoc throwsLoc, SourceLoc arrowLoc,
+  ClosureExpr(SourceRange bracketRange, VarDecl *capturedSelfDecl,
+              ParameterList *params, SourceLoc throwsLoc, SourceLoc arrowLoc,
               SourceLoc inLoc, TypeLoc explicitResultType,
               unsigned discriminator, DeclContext *parent)
     : AbstractClosureExpr(ExprKind::Closure, Type(), /*Implicit=*/false,
                           discriminator, parent),
+      BracketRange(bracketRange), CapturedSelfDecl(capturedSelfDecl),
       ThrowsLoc(throwsLoc), ArrowLoc(arrowLoc), InLoc(inLoc),
-      ExplicitResultType(explicitResultType),
-      Body(nullptr) {
+      ExplicitResultType(explicitResultType), Body(nullptr) {
     setParameterList(params);
     Bits.ClosureExpr.HasAnonymousClosureVars = false;
   }
@@ -3614,6 +3627,8 @@ class ClosureExpr : public AbstractClosureExpr {
   /// explicitly-specified result type.
   bool hasExplicitResultType() const { return ArrowLoc.isValid(); }
 
+  /// Retrieve the range of the \c '[' and \c ']' that enclose the capture list.
+  SourceRange getBracketRange() const { return BracketRange; }
 
   /// Retrieve the location of the \c '->' for closures with an
   /// explicit result type.
@@ -3679,6 +3694,9 @@ class ClosureExpr : public AbstractClosureExpr {
   /// Is this a completely empty closure?
   bool hasEmptyBody() const;
 
+  /// VarDecl captured by this closure under the literal name \c self , if any.
+  VarDecl *getCapturedSelfDecl() const { return CapturedSelfDecl; }
+
   static bool classof(const Expr *E) {
     return E->getKind() == ExprKind::Closure;
   }
@@ -3746,6 +3764,8 @@ struct CaptureListEntry {
   CaptureListEntry(VarDecl *Var, PatternBindingDecl *Init)
   : Var(Var), Init(Init) {
   }
+
+  bool isSimpleSelfCapture() const;
 };
 
 /// CaptureListExpr - This expression represents the capture list on an explicit

include/swift/AST/NameLookup.h

@@ -73,13 +73,15 @@ namespace swift {
     /// When finding \c bar() from the function body of \c foo(), \c BaseDC is
     /// the method \c foo().
     ///
-    /// \c BaseDC will be the method if \c self is needed for the lookup,
-    /// and will be the type if not.
-    /// In other words: If \c baseDC is a method, it means you found an instance
-    /// member and you should add an implicit 'self.' (Each method has its own
-    /// implicit self decl.) There's one other kind of non-method context that
-    /// has a 'self.' -- a lazy property initializer, which unlike a non-lazy
-    /// property can reference \c self) Hence: \code
+    /// \c BaseDC will be the type if \c self is not needed for the lookup. If
+    /// \c self is needed, \c baseDC will be either the method or a closure
+    /// which explicitly captured \c self.
+    /// In other words: If \c baseDC is a method or a closure, it means you
+    /// found an instance member and you should add an implicit 'self.' (Each
+    /// method has its own implicit self decl.) There's one other kind of
+    /// non-method, non-closure context that has a 'self.' -- a lazy property
+    /// initializer, which unlike a non-lazy property can reference \c self)
+    /// Hence: \code
     ///  class Outer {
     ///    static func s()
     ///    func i()

include/swift/Parse/Lexer.h

@@ -293,7 +293,12 @@ class Lexer {
   /// resides.
   ///
   /// \param Loc The source location of the beginning of a token.
-  static Token getTokenAtLocation(const SourceManager &SM, SourceLoc Loc);
+  ///
+  /// \param CRM The how comments should be treated by the lexer. Default is to
+  /// return the comments as tokens.
+  static Token getTokenAtLocation(
+      const SourceManager &SM, SourceLoc Loc,
+      CommentRetentionMode CRM = CommentRetentionMode::ReturnAsTokens);
 
 
   /// Retrieve the source location that points just past the

include/swift/Parse/Parser.h

@@ -1407,6 +1407,8 @@ class Parser {
   ///     identifier (',' identifier)* func-signature-result? 'in'
   /// \endverbatim
   ///
+  /// \param bracketRange The range of the brackets enclosing a capture list, if
+  /// present
   /// \param captureList The entries in the capture list.
   /// \param params The parsed parameter list, or null if none was provided.
   /// \param arrowLoc The location of the arrow, if present.
@@ -1415,12 +1417,10 @@ class Parser {
   ///
   /// \returns true if an error occurred, false otherwise.
   bool parseClosureSignatureIfPresent(
-                                SmallVectorImpl<CaptureListEntry> &captureList,
-                                      ParameterList *&params,
-                                      SourceLoc &throwsLoc,
-                                      SourceLoc &arrowLoc,
-                                      TypeRepr *&explicitResultType,
-                                      SourceLoc &inLoc);
+      SourceRange &bracketRange, SmallVectorImpl<CaptureListEntry> &captureList,
+      VarDecl *&capturedSelfParamDecl, ParameterList *&params,
+      SourceLoc &throwsLoc, SourceLoc &arrowLoc, TypeRepr *&explicitResultType,
+      SourceLoc &inLoc);
 
   Expr *parseExprAnonClosureArg();
   ParserResult<Expr> parseExprList(tok LeftTok, tok RightTok,

lib/AST/Decl.cpp

@@ -1526,7 +1526,7 @@ VarDecl *PatternBindingDecl::getSingleVar() const {
     return getPatternList()[0].getPattern()->getSingleVar();
   return nullptr;
 }
-
+    
 bool VarDecl::isInitExposedToClients() const {
   auto parent = dyn_cast<NominalTypeDecl>(getDeclContext());
   if (!parent) return false;
@@ -4963,6 +4963,7 @@ VarDecl::VarDecl(DeclKind kind, bool isStatic, VarDecl::Introducer introducer,
 {
   Bits.VarDecl.Introducer = unsigned(introducer);
   Bits.VarDecl.IsCaptureList = isCaptureList;
+  Bits.VarDecl.IsSelfParamCapture = false;
   Bits.VarDecl.IsDebuggerVar = false;
   Bits.VarDecl.IsLazyStorageProperty = false;
   Bits.VarDecl.HasNonPatternBindingInit = false;

lib/AST/Expr.cpp

@@ -1228,6 +1228,17 @@ UnresolvedSpecializeExpr *UnresolvedSpecializeExpr::create(ASTContext &ctx,
                                              UnresolvedParams, RAngleLoc);
 }
 
+bool CaptureListEntry::isSimpleSelfCapture() const {
+  if (Init->getPatternList().size() != 1)
+    return false;
+  if (auto *DRE = dyn_cast<DeclRefExpr>(Init->getInit(0)))
+    if (auto *VD = dyn_cast<VarDecl>(DRE->getDecl())) {
+      return (VD->isSelfParameter() || VD->isSelfParamCapture())
+             && VD->getName() == Var->getName();
+    }
+  return false;
+}
+
 CaptureListExpr *CaptureListExpr::create(ASTContext &ctx,
                                          ArrayRef<CaptureListEntry> captureList,
                                          ClosureExpr *closureBody) {

lib/AST/NameLookup.cpp

@@ -56,6 +56,13 @@ ValueDecl *LookupResultEntry::getBaseDecl() const {
     return selfDecl;
   }
 
+  if (auto *CE = dyn_cast<ClosureExpr>(BaseDC)) {
+    auto *selfDecl = CE->getCapturedSelfDecl();
+    assert(selfDecl);
+    assert(selfDecl->isSelfParamCapture());
+    return selfDecl;
+  }
+
   auto *nominalDecl = BaseDC->getSelfNominalTypeDecl();
   assert(nominalDecl);
   return nominalDecl;

lib/AST/UnqualifiedLookup.cpp

@@ -198,6 +198,13 @@ namespace {
     bool recordedIsCascadingUse = false;
     unsigned resultsSizeBeforeLocalsPass = ~0;
 
+  private:
+    // When performing a lookup, we may come across a capture of 'self'. We will
+    // need to remember the DeclContext of the innermost captured self so that
+    // it can be used as the base DeclContext if we find a lookup result in the
+    // enclosing type.
+    DeclContext *capturedSelfContext;
+
   public:
     // clang-format off
     UnqualifiedLookupFactory(DeclName Name,
@@ -463,6 +470,7 @@ UnqualifiedLookupFactory::UnqualifiedLookupFactory(
   options(options),
   isOriginallyTypeLookup(options.contains(Flags::TypeLookup)),
   baseNLOptions(computeBaseNLOptions(options, isOriginallyTypeLookup)),
+  capturedSelfContext(nullptr),
   #ifdef NDEBUG
   Consumer(Name, Results, isOriginallyTypeLookup),
   #else
@@ -693,9 +701,12 @@ void UnqualifiedLookupFactory::lookupNamesIntroducedByMemberFunction(
         // If we're not in the body of the function (for example, we
         // might be type checking a default argument expression and
         // performing name lookup from there), the base declaration
-        // is the nominal type, not 'self'.
+        // is the nominal type, not 'self'. If we've captured self
+        // somewhere down the tree, we should use that as the context
+        // for lookup.
         DeclContext *const BaseDC =
-            isOutsideBodyOfFunction(AFD) ? fnDeclContext : AFD;
+            isOutsideBodyOfFunction(AFD) ? fnDeclContext
+            : capturedSelfContext ?:  AFD;
         // If we are inside of a method, check to see if there are any ivars in
         // scope, and if so, whether this is a reference to one of them.
         // FIXME: We should persist this information between lookups.
@@ -725,8 +736,16 @@ void UnqualifiedLookupFactory::lookupNamesIntroducedByPureFunction(
 
 void UnqualifiedLookupFactory::lookupNamesIntroducedByClosure(
     AbstractClosureExpr *ACE, Optional<bool> isCascadingUse) {
-  if (auto *CE = dyn_cast<ClosureExpr>(ACE))
+  if (auto *CE = dyn_cast<ClosureExpr>(ACE)) {
     lookForLocalVariablesIn(CE);
+    // If we don't already have a captured self context, and this closure
+    // captures the self param (not weakly, so that implicit self is available),
+    // remember that.
+    if (capturedSelfContext == nullptr)
+      if (auto *VD = CE->getCapturedSelfDecl())
+        if (VD->isSelfParamCapture() && !VD->getType()->is<WeakStorageType>())
+          capturedSelfContext = CE;
+  }
   ifNotDoneYet([&] {
     // clang-format off
     finishLookingInContext(

lib/Parse/Lexer.cpp

@@ -2489,7 +2489,8 @@ void Lexer::lexImpl() {
   }
 }
 
-Token Lexer::getTokenAtLocation(const SourceManager &SM, SourceLoc Loc) {
+Token Lexer::getTokenAtLocation(const SourceManager &SM, SourceLoc Loc,
+                                CommentRetentionMode CRM) {
   // Don't try to do anything with an invalid location.
   if (!Loc.isValid())
     return Token();
@@ -2508,7 +2509,7 @@ Token Lexer::getTokenAtLocation(const SourceManager &SM, SourceLoc Loc) {
   // (making this option irrelevant), or the caller lexed comments and
   // we need to lex just the comment token.
   Lexer L(FakeLangOpts, SM, BufferID, nullptr, LexerMode::Swift,
-          HashbangMode::Allowed, CommentRetentionMode::ReturnAsTokens);
+          HashbangMode::Allowed, CRM);
   L.restoreState(State(Loc));
   return L.peekNextToken();
 }