[AutoDiff upstream] Add @differentiable ASTScope support.

`@differentiable` attributes may contain `where` clauses referencing generic
parameters from some generic context, just like `@_specialize` attributes.

Without special ASTScope support for `@differentiable` attributes,
ASTScopeLookup.cpp logic tries to resolve the generic parameter `DeclName`s in
`where` clauses based on source location alone
(`ASTScopeImpl::findChildContaining`) and fails.

The fix is to add a special `DifferentiableAttributeScope`, mimicking
`SpecializeAttributeScope`. Every `@differentiable` attribute has its own scope,
derived from the declaration on which it is declared. Unlike `@_specialize`,
`@differentiable` may also be declared on `AbstractStorageDecl` declarations
(subscripts and variables).

Upstreams #27451.

Progress towards TF-828: upstream `@differentiable` attribute type-checking.

Author: dan-zheng

include/swift/AST/ASTScope.h

@@ -1568,6 +1568,41 @@ class SpecializeAttributeScope final : public ASTScopeImpl {
                              DeclConsumer) const override;
 };
 
+/// A `@differentiable` attribute scope.
+///
+/// This exists because `@differentiable` attribute may have a `where` clause
+/// referring to generic parameters from some generic context.
+class DifferentiableAttributeScope final : public ASTScopeImpl {
+public:
+  DifferentiableAttr *const differentiableAttr;
+  ValueDecl *const attributedDeclaration;
+
+  DifferentiableAttributeScope(DifferentiableAttr *diffAttr, ValueDecl *decl)
+      : differentiableAttr(diffAttr), attributedDeclaration(decl) {}
+  virtual ~DifferentiableAttributeScope() {}
+
+  std::string getClassName() const override;
+  SourceRange
+  getSourceRangeOfThisASTNode(bool omitAssertions = false) const override;
+  NullablePtr<const void> addressForPrinting() const override {
+    return differentiableAttr;
+  }
+
+  NullablePtr<AbstractStorageDecl>
+  getEnclosingAbstractStorageDecl() const override;
+
+  NullablePtr<DeclAttribute> getDeclAttributeIfAny() const override {
+    return differentiableAttr;
+  }
+  NullablePtr<const void> getReferrent() const override;
+
+protected:
+  ASTScopeImpl *expandSpecifically(ScopeCreator &) override;
+  bool lookupLocalsOrMembers(ArrayRef<const ASTScopeImpl *>,
+                             DeclConsumer) const override;
+  bool doesContextMatchStartingContext(const DeclContext *) const override;
+};
+
 class SubscriptDeclScope final : public ASTScopeImpl {
 public:
   SubscriptDecl *const decl;

lib/AST/ASTScope.cpp

@@ -228,6 +228,7 @@ DEFINE_GET_CLASS_NAME(ClosureParametersScope)
 DEFINE_GET_CLASS_NAME(ClosureBodyScope)
 DEFINE_GET_CLASS_NAME(TopLevelCodeScope)
 DEFINE_GET_CLASS_NAME(SpecializeAttributeScope)
+DEFINE_GET_CLASS_NAME(DifferentiableAttributeScope)
 DEFINE_GET_CLASS_NAME(SubscriptDeclScope)
 DEFINE_GET_CLASS_NAME(VarDeclScope)
 DEFINE_GET_CLASS_NAME(EnumElementScope)

lib/AST/ASTScopeCreation.cpp

@@ -64,6 +64,9 @@ static SourceRange getRangeableSourceRange(const Rangeable *const p) {
 static SourceRange getRangeableSourceRange(const SpecializeAttr *a) {
   return a->getRange();
 }
+static SourceRange getRangeableSourceRange(const DifferentiableAttr *a) {
+  return a->getRange();
+}
 static SourceRange getRangeableSourceRange(const ASTNode n) {
   return n.getSourceRange();
 }
@@ -98,6 +101,17 @@ static void dumpRangeable(SpecializeAttr *r, llvm::raw_ostream &f) {
   llvm::errs() << "SpecializeAttr\n";
 }
 
+static void dumpRangeable(const DifferentiableAttr *a,
+                          llvm::raw_ostream &f) LLVM_ATTRIBUTE_USED;
+static void dumpRangeable(const DifferentiableAttr *a, llvm::raw_ostream &f) {
+  llvm::errs() << "DifferentiableAttr\n";
+}
+static void dumpRangeable(DifferentiableAttr *a,
+                          llvm::raw_ostream &f) LLVM_ATTRIBUTE_USED;
+static void dumpRangeable(DifferentiableAttr *a, llvm::raw_ostream &f) {
+  llvm::errs() << "DifferentiableAttr\n";
+}
+
 /// For Debugging
 template <typename T>
 bool doesRangeableRangeMatch(const T *x, const SourceManager &SM,
@@ -439,6 +453,22 @@ class ScopeCreator final {
       fn(specializeAttr);
   }
 
+  void forEachDifferentiableAttrInSourceOrder(
+      Decl *decl, function_ref<void(DifferentiableAttr *)> fn) {
+    std::vector<DifferentiableAttr *> sortedDifferentiableAttrs;
+    for (auto *attr : decl->getAttrs())
+      if (auto *diffAttr = dyn_cast<DifferentiableAttr>(attr))
+        // NOTE(TF-835): Skipping implicit `@differentiable` attributes is
+        // necessary to avoid verification failure in
+        // `ASTScopeImpl::verifyThatChildrenAreContainedWithin`.
+        // Perhaps this check may no longer be necessary after TF-835: robust
+        // `@derivative` attribute lowering.
+        if (!diffAttr->isImplicit())
+          sortedDifferentiableAttrs.push_back(diffAttr);
+    for (auto *diffAttr : sortBySourceRange(sortedDifferentiableAttrs))
+      fn(diffAttr);
+  }
+
   std::vector<ASTNode> expandIfConfigClausesThenCullAndSortElementsOrMembers(
       ArrayRef<ASTNode> input) const {
     auto cleanedupNodes = sortBySourceRange(cull(expandIfConfigClauses(input)));
@@ -1039,6 +1069,13 @@ void ScopeCreator::addChildrenForAllLocalizableAccessorsInSourceOrder(
                   return enclosingAbstractStorageDecl == ad->getStorage();
                 });
 
+  // Create scopes for `@differentiable` attributes.
+  forEachDifferentiableAttrInSourceOrder(
+      asd, [&](DifferentiableAttr *diffAttr) {
+        ifUniqueConstructExpandAndInsert<DifferentiableAttributeScope>(
+            parent, diffAttr, asd);
+      });
+
   // Sort in order to include synthesized ones, which are out of order.
   for (auto *accessor : sortBySourceRange(accessorsToScope))
     addToScopeTree(accessor, parent);
@@ -1183,6 +1220,7 @@ NO_NEW_INSERTION_POINT(WholeClosureScope)
 NO_EXPANSION(GenericParamScope)
 NO_EXPANSION(ClosureParametersScope)
 NO_EXPANSION(SpecializeAttributeScope)
+NO_EXPANSION(DifferentiableAttributeScope)
 NO_EXPANSION(ConditionalClausePatternUseScope)
 NO_EXPANSION(LookupParentDiversionScope)
 
@@ -1353,6 +1391,13 @@ void AbstractFunctionDeclScope::expandAScopeThatDoesNotCreateANewInsertionPoint(
         scopeCreator.ifUniqueConstructExpandAndInsert<SpecializeAttributeScope>(
             this, specializeAttr, decl);
       });
+  // Create scopes for `@differentiable` attributes.
+  scopeCreator.forEachDifferentiableAttrInSourceOrder(
+      decl, [&](DifferentiableAttr *diffAttr) {
+        scopeCreator
+            .ifUniqueConstructExpandAndInsert<DifferentiableAttributeScope>(
+                this, diffAttr, decl);
+      });
   // Create scopes for generic and ordinary parameters.
   // For a subscript declaration, the generic and ordinary parameters are in an
   // ancestor scope, so don't make them here.
@@ -1681,6 +1726,10 @@ SpecializeAttributeScope::getEnclosingAbstractStorageDecl() const {
   return getParent().get()->getEnclosingAbstractStorageDecl();
 }
 NullablePtr<AbstractStorageDecl>
+DifferentiableAttributeScope::getEnclosingAbstractStorageDecl() const {
+  return getParent().get()->getEnclosingAbstractStorageDecl();
+}
+NullablePtr<AbstractStorageDecl>
 AbstractFunctionDeclScope::getEnclosingAbstractStorageDecl() const {
   return getParent().get()->getEnclosingAbstractStorageDecl();
 }
@@ -1807,6 +1856,7 @@ GET_REFERRENT(AbstractStmtScope, getStmt())
 GET_REFERRENT(CaptureListScope, getExpr())
 GET_REFERRENT(WholeClosureScope, getExpr())
 GET_REFERRENT(SpecializeAttributeScope, specializeAttr)
+GET_REFERRENT(DifferentiableAttributeScope, differentiableAttr)
 GET_REFERRENT(GenericTypeOrExtensionScope, portion->getReferrentOfScope(this));
 
 const Decl *

lib/AST/ASTScopeLookup.cpp

@@ -208,6 +208,19 @@ bool GenericParamScope::doesContextMatchStartingContext(
   return false;
 }
 
+bool DifferentiableAttributeScope::doesContextMatchStartingContext(
+    const DeclContext *context) const {
+  // Need special logic to handle case where `attributedDeclaration` is an
+  // `AbstractStorageDecl` (`SubscriptDecl` or `VarDecl`). The initial starting
+  // context in `ASTScopeImpl::findStartingScopeForLookup` will be an accessor
+  // of the `attributedDeclaration`.
+  if (auto *asd = dyn_cast<AbstractStorageDecl>(attributedDeclaration))
+    for (auto accessor : asd->getAllAccessors())
+      if (up_cast<DeclContext>(accessor) == context)
+        return true;
+  return false;
+}
+
 #pragma mark lookup methods that run once per scope
 
 void ASTScopeImpl::lookup(SmallVectorImpl<const ASTScopeImpl *> &history,
@@ -438,6 +451,25 @@ bool SpecializeAttributeScope::lookupLocalsOrMembers(
   return false;
 }
 
+bool DifferentiableAttributeScope::lookupLocalsOrMembers(
+    ArrayRef<const ASTScopeImpl *>, DeclConsumer consumer) const {
+  auto visitAbstractFunctionDecl = [&](AbstractFunctionDecl *afd) {
+    if (auto *params = afd->getGenericParams())
+      for (auto *param : params->getParams())
+        if (consumer.consume({param}, DeclVisibilityKind::GenericParameter))
+          return true;
+    return false;
+  };
+  if (auto *afd = dyn_cast<AbstractFunctionDecl>(attributedDeclaration)) {
+    return visitAbstractFunctionDecl(afd);
+  } else if (auto *asd = dyn_cast<AbstractStorageDecl>(attributedDeclaration)) {
+    for (auto *accessor : asd->getAllAccessors())
+      if (visitAbstractFunctionDecl(accessor))
+        return true;
+  }
+  return false;
+}
+
 bool BraceStmtScope::lookupLocalsOrMembers(ArrayRef<const ASTScopeImpl *>,
                                            DeclConsumer consumer) const {
   // All types and functions are visible anywhere within a brace statement

lib/AST/ASTScopeSourceRange.cpp

@@ -193,6 +193,11 @@ SourceRange SpecializeAttributeScope::getSourceRangeOfThisASTNode(
   return specializeAttr->getRange();
 }
 
+SourceRange DifferentiableAttributeScope::getSourceRangeOfThisASTNode(
+    const bool omitAssertions) const {
+  return differentiableAttr->getRange();
+}
+
 SourceRange AbstractFunctionBodyScope::getSourceRangeOfThisASTNode(
     const bool omitAssertions) const {
   return decl->getBodySourceRange();