Allow implicit self in escaping closures when self usage is unlikely to cause cycle (swiftlang#23934)

* WIP implementation

* Cleanup implementation

* Install backedge rather than storing array reference

* Add diagnostics

* Add missing parameter to ResultFinderForTypeContext constructor

* Fix tests for correct fix-it language

* Change to solution without backedge, change lookup behavior

* Improve diagnostics for weak captures and captures under different names

* Remove ghosts of implementations past

* Address review comments

* Reorder member variable initialization

* Fix typos

* Exclude value types from explicit self requirements

* Add tests

* Add implementation for AST lookup

* Add tests

* Begin addressing review comments

* Re-enable AST scope lookup

* Add fixme

* Pull fix-its into a separate function

* Remove capturedSelfContext tracking from type property initializers

* Add const specifiers to arguments

* Address review comments

* Fix string literals

* Refactor implicit self diagnostics

* Add comment

* Remove trailing whitespace

* Add tests for capture list across multiple lines

* Add additional test

* Fix typo

* Remove use of ?: to fix linux build

* Remove second use of ?:

* Rework logic for finding nested self contexts

Author: Jumhyn

include/swift/AST/ASTScope.h

@@ -479,6 +479,10 @@ class ASTScopeImpl {
   /// asked.
   virtual Optional<NullablePtr<DeclContext>> computeSelfDCForParent() const;
 
+  /// Returns the context that should be used when a nested scope (e.g. a
+  /// closure) captures self explicitly.
+  virtual NullablePtr<DeclContext> capturedSelfDC() const;
+
 protected:
   /// Find either locals or members (no scope has both)
   /// \param history The scopes visited since the start of lookup (including
@@ -692,6 +696,15 @@ class Portion {
     /// to compute the selfDC from the history.
     static NullablePtr<DeclContext>
     computeSelfDC(ArrayRef<const ASTScopeImpl *> history);
+    
+    /// If we find a lookup result that requires the dynamic implict self value,
+    /// we need to check the nested scopes to see if any closures explicitly
+    /// captured \c self. In that case, the appropriate selfDC is that of the
+    /// innermost closure which captures a \c self value from one of this type's
+    /// methods.
+    static NullablePtr<DeclContext>
+    checkNestedScopesForSelfCapture(ArrayRef<const ASTScopeImpl *> history,
+                                    size_t start);
 };
 
 /// Behavior specific to representing the trailing where clause of a
@@ -1449,6 +1462,13 @@ class ClosureParametersScope final : public AbstractClosureScope {
   std::string getClassName() const override;
   SourceRange
   getSourceRangeOfThisASTNode(bool omitAssertions = false) const override;
+  
+  /// Since explicit captures of \c self by closures enable the use of implicit
+  /// \c self, we need to make sure that the appropriate \c self is used as the
+  /// base decl for these uses (otherwise, the capture would be marked as
+  /// unused. \c ClosureParametersScope::capturedSelfDC() checks if we have such
+  ///  a capture of self.
+  NullablePtr<DeclContext> capturedSelfDC() const override;
 
 protected:
   ASTScopeImpl *expandSpecifically(ScopeCreator &) override;

include/swift/AST/Decl.h

@@ -334,12 +334,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
@@ -5029,6 +5032,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

@@ -3244,11 +3244,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

@@ -65,6 +65,7 @@ namespace swift {
   class EnumElementDecl;
   class CallExpr;
   class KeyPathExpr;
+  class CaptureListExpr;
 
 enum class ExprKind : uint8_t {
 #define EXPR(Id, Parent) Id,
@@ -3582,6 +3583,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;
 
@@ -3598,16 +3611,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;
   }
@@ -3639,6 +3652,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.
@@ -3704,6 +3719,14 @@ 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; }
+  
+  /// Whether this closure captures the \c self param in its body in such a
+  /// way that implicit \c self is enabled within its body (i.e. \c self is
+  /// captured non-weakly).
+  bool capturesSelfEnablingImplictSelf() const;
+
   static bool classof(const Expr *E) {
     return E->getKind() == ExprKind::Closure;
   }
@@ -3778,6 +3801,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

@@ -70,13 +70,15 @@ struct LookupResultEntry {
   /// 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.
+  /// \code
   ///  class Outer {
   ///    static func s()
   ///    func i()

include/swift/Parse/Lexer.h

@@ -293,7 +293,17 @@ class Lexer {
   /// resides.
   ///
   /// \param Loc The source location of the beginning of a token.
-  static Token getTokenAtLocation(const SourceManager &SM, SourceLoc Loc);
+  ///
+  /// \param CRM How comments should be treated by the lexer. Default is to
+  /// return the comments as tokens. This is needed in situations where
+  /// detecting the next semantically meaningful token is required, such as
+  /// the 'implicit self' diagnostic determining whether a capture list is
+  /// empty (i.e., the opening bracket is immediately followed by a closing
+  /// bracket, possibly with comments in between) in order to insert the
+  /// appropriate fix-it.
+  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

@@ -1514,6 +1514,8 @@ class Parser {
   ///     identifier (',' identifier)* func-signature-result? 'in'
   /// \endverbatim
   ///
+  /// \param bracketRange The range of the brackets enclosing a capture list, if
+  /// present. Needed to offer fix-its for inserting 'self' into a capture list.
   /// \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.
@@ -1522,12 +1524,14 @@ 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/ASTScopeLookup.cpp

@@ -517,12 +517,59 @@ GenericTypeOrExtensionWhereOrBodyPortion::computeSelfDC(
   while (i != 0) {
     Optional<NullablePtr<DeclContext>> maybeSelfDC =
         history[--i]->computeSelfDCForParent();
-    if (maybeSelfDC)
-      return *maybeSelfDC;
+    if (maybeSelfDC) {
+      // If we've found a selfDC, we'll definitely be returning something.
+      // However, we may have captured 'self' somewhere down the tree, so we
+      // can't return outright without checking the nested scopes.
+      NullablePtr<DeclContext> nestedCapturedSelfDC =
+          checkNestedScopesForSelfCapture(history, i);
+      return nestedCapturedSelfDC ? nestedCapturedSelfDC : *maybeSelfDC;
+    }
   }
   return nullptr;
 }
 
+#pragma mark checkNestedScopesForSelfCapture
+
+NullablePtr<DeclContext>
+GenericTypeOrExtensionWhereOrBodyPortion::checkNestedScopesForSelfCapture(
+    ArrayRef<const ASTScopeImpl *> history, size_t start) {
+  NullablePtr<DeclContext> innerCapturedSelfDC;
+  // Start with the next scope down the tree.
+  size_t j = start;
+
+  // Note: even though having this loop nested inside the while loop from
+  // GenericTypeOrExtensionWhereOrBodyPortion::computeSelfDC may appear to
+  // result in quadratic blowup, complexity actually remains linear with respect
+  // to the size of history. This relies on the fact that
+  // GenericTypeOrExtensionScope::computeSelfDCForParent returns a null pointer,
+  // which will cause this method to bail out of the search early. Thus, this
+  // method is called once per type body in the lookup history, and will not
+  // end up re-checking the bodies of nested types that have already been
+  // covered by earlier calls, so the total impact of this method across all
+  // calls in a single lookup is O(n).
+  while (j != 0) {
+      auto *entry = history[--j];
+    Optional<NullablePtr<DeclContext>> selfDCForParent =
+      entry->computeSelfDCForParent();
+
+    // If we encounter a scope that should cause us to forget the self
+    // context (such as a nested type), bail out and use whatever the
+    // the last inner captured context was.
+    if (selfDCForParent && (*selfDCForParent).isNull())
+      break;
+
+    // Otherwise, if we have a captured self context for this scope, then
+    // remember it since it is now the innermost scope we have encountered.
+    NullablePtr<DeclContext> capturedSelfDC = entry->capturedSelfDC();
+    if (!capturedSelfDC.isNull())
+      innerCapturedSelfDC = entry->capturedSelfDC();
+
+    // Continue searching in the next scope down.
+  }
+  return innerCapturedSelfDC;
+}
+
 #pragma mark compute isCascadingUse
 
 Optional<bool> ASTScopeImpl::computeIsCascadingUse(
@@ -631,6 +678,23 @@ MethodBodyScope::computeSelfDCForParent() const {
   return NullablePtr<DeclContext>(decl);
 }
 
+#pragma mark capturedSelfDC
+
+// Closures may explicitly capture the self param, in which case the lookup
+// should use the closure as the context for implicit self lookups.
+
+// By default, there is no such context to return.
+NullablePtr<DeclContext> ASTScopeImpl::capturedSelfDC() const {
+  return NullablePtr<DeclContext>();
+}
+
+// Closures track this information explicitly.
+NullablePtr<DeclContext> ClosureParametersScope::capturedSelfDC() const {
+  if (closureExpr->capturesSelfEnablingImplictSelf())
+    return NullablePtr<DeclContext>(closureExpr);
+  return NullablePtr<DeclContext>();
+}
+
 #pragma mark ifUnknownIsCascadingUseAccordingTo
 
 static bool isCascadingUseAccordingTo(const DeclContext *const dc) {

lib/AST/Decl.cpp

@@ -5249,6 +5249,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

@@ -1179,6 +1179,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) {
@@ -1813,6 +1824,12 @@ bool ClosureExpr::hasEmptyBody() const {
   return getBody()->empty();
 }
 
+bool ClosureExpr::capturesSelfEnablingImplictSelf() const {
+  if (auto *VD = getCapturedSelfDecl())
+    return VD->isSelfParamCapture() && !VD->getType()->is<WeakStorageType>();
+  return false;
+}
+
 FORWARD_SOURCE_LOCS_TO(AutoClosureExpr, Body)
 
 void AutoClosureExpr::setBody(Expr *E) {