[AutoDiff] Add TangentStoredPropertyRequest.

Add request that resolves the "tangent stored property" corresponding to an
original stored property in a `Differentiable`-conforming type.

Enables better non-differentiability differentiation transform diagnostics.

Author: dan-zheng

include/swift/AST/ASTTypeIDZone.def

@@ -27,6 +27,7 @@ SWIFT_TYPEID(PropertyWrapperTypeInfo)
 SWIFT_TYPEID(Requirement)
 SWIFT_TYPEID(ResilienceExpansion)
 SWIFT_TYPEID(FragileFunctionKind)
+SWIFT_TYPEID(TangentPropertyInfo)
 SWIFT_TYPEID(Type)
 SWIFT_TYPEID(TypePair)
 SWIFT_TYPEID(TypeWitnessAndDecl)

include/swift/AST/ASTTypeIDs.h

@@ -19,6 +19,7 @@
 
 #include "swift/Basic/LLVM.h"
 #include "swift/Basic/TypeID.h"
+
 namespace swift {
 
 class AbstractFunctionDecl;
@@ -58,14 +59,14 @@ class Requirement;
 enum class ResilienceExpansion : unsigned;
 struct FragileFunctionKind;
 class SourceFile;
+struct TangentPropertyInfo;
 class Type;
-class ValueDecl;
-class VarDecl;
-class Witness;
 class TypeAliasDecl;
-class Type;
 struct TypePair;
 struct TypeWitnessAndDecl;
+class ValueDecl;
+class VarDecl;
+class Witness;
 enum class AncestryFlags : uint8_t;
 enum class ImplicitMemberAction : uint8_t;
 struct FingerprintAndMembers;

include/swift/AST/AutoDiff.h

@@ -35,6 +35,7 @@ class AnyFunctionType;
 class SourceFile;
 class SILFunctionType;
 class TupleType;
+class VarDecl;
 
 /// A function type differentiability kind.
 enum class DifferentiabilityKind : uint8_t {
@@ -459,6 +460,99 @@ class DerivativeFunctionTypeError
   }
 };
 
+/// Describes the "tangent stored property" corresponding to an original stored
+/// property in a `Differentiable`-conforming type.
+///
+/// The tangent stored property is the stored property in the `TangentVector`
+/// struct of the `Differentiable`-conforming type, with the same name as the
+/// original stored property and with the original stored property's
+/// `TangentVector` type.
+struct TangentPropertyInfo {
+  struct Error {
+    enum class Kind {
+      /// The original property is `@noDerivative`.
+      NoDerivativeOriginalProperty,
+      /// The nominal parent type does not conform to `Differentiable`.
+      NominalParentNotDifferentiable,
+      /// The original property's type does not conform to `Differentiable`.
+      OriginalPropertyNotDifferentiable,
+      /// The parent `TangentVector` type is not a struct.
+      ParentTangentVectorNotStruct,
+      /// The parent `TangentVector` struct does not declare a stored property
+      /// with the same name as the original property.
+      TangentPropertyNotFound,
+      /// The tangent property's type is not equal to the original property's
+      /// `TangentVector` type.
+      TangentPropertyWrongType,
+      /// The tangent property is not a stored property.
+      TangentPropertyNotStored
+    };
+
+    /// The error kind.
+    Kind kind;
+
+  private:
+    union Value {
+      Type type;
+      Value(Type type) : type(type) {}
+      Value() {}
+    } value;
+
+  public:
+    Error(Kind kind) : kind(kind), value() {
+      assert(kind == Kind::NoDerivativeOriginalProperty ||
+             kind == Kind::NominalParentNotDifferentiable ||
+             kind == Kind::OriginalPropertyNotDifferentiable ||
+             kind == Kind::ParentTangentVectorNotStruct ||
+             kind == Kind::TangentPropertyNotFound ||
+             kind == Kind::TangentPropertyNotStored);
+    };
+
+    Error(Kind kind, Type type) : kind(kind), value(type) {
+      assert(kind == Kind::TangentPropertyWrongType);
+    };
+
+    Type getType() const {
+      assert(kind == Kind::TangentPropertyWrongType);
+      return value.type;
+    }
+
+    friend bool operator==(const Error &lhs, const Error &rhs);
+  };
+
+  /// The tangent stored property.
+  VarDecl *tangentProperty = nullptr;
+
+  /// An optional error.
+  Optional<Error> error = None;
+
+private:
+  TangentPropertyInfo(VarDecl *tangentProperty, Optional<Error> error)
+      : tangentProperty(tangentProperty), error(error) {}
+
+public:
+  TangentPropertyInfo(VarDecl *tangentProperty)
+      : TangentPropertyInfo(tangentProperty, None) {}
+
+  TangentPropertyInfo(Error::Kind errorKind)
+      : TangentPropertyInfo(nullptr, Error(errorKind)) {}
+
+  TangentPropertyInfo(Error::Kind errorKind, Type errorType)
+      : TangentPropertyInfo(nullptr, Error(errorKind, errorType)) {}
+
+  /// Returns `true` iff this tangent property info is valid.
+  bool isValid() const { return tangentProperty && !error; }
+
+  explicit operator bool() const { return isValid(); }
+
+  friend bool operator==(const TangentPropertyInfo &lhs,
+                         const TangentPropertyInfo &rhs) {
+    return lhs.tangentProperty == rhs.tangentProperty && lhs.error == rhs.error;
+  }
+};
+
+void simple_display(llvm::raw_ostream &OS, TangentPropertyInfo info);
+
 /// The key type used for uniquing `SILDifferentiabilityWitness` in
 /// `SILModule`: original function name, parameter indices, result indices, and
 /// derivative generic signature.

include/swift/AST/TypeCheckRequests.h

@@ -2191,6 +2191,27 @@ class DerivativeAttrOriginalDeclRequest
   bool isCached() const { return true; }
 };
 
+/// Resolves the "tangent stored property" corresponding to an original stored
+/// property in a `Differentiable`-conforming type.
+class TangentStoredPropertyRequest
+    : public SimpleRequest<TangentStoredPropertyRequest,
+                           TangentPropertyInfo(VarDecl *),
+                           RequestFlags::Cached> {
+public:
+  using SimpleRequest::SimpleRequest;
+
+private:
+  friend SimpleRequest;
+
+  // Evaluation.
+  TangentPropertyInfo evaluate(Evaluator &evaluator,
+                               VarDecl *originalField) const;
+
+public:
+  // Caching.
+  bool isCached() const { return true; }
+};
+
 /// Checks whether a type eraser has a viable initializer.
 class TypeEraserHasViableInitRequest
     : public SimpleRequest<TypeEraserHasViableInitRequest,

include/swift/AST/TypeCheckerTypeIDZone.def

@@ -203,6 +203,8 @@ SWIFT_REQUEST(TypeChecker, SuperclassTypeRequest,
 SWIFT_REQUEST(TypeChecker, SynthesizeAccessorRequest,
               AccessorDecl *(AbstractStorageDecl *, AccessorKind),
               SeparatelyCached, NoLocationInfo)
+SWIFT_REQUEST(TypeChecker, TangentStoredPropertyRequest,
+              llvm::Expected<VarDecl *>(VarDecl *), Cached, NoLocationInfo)
 SWIFT_REQUEST(TypeChecker, TypeCheckFunctionBodyRequest,
               bool(AbstractFunctionDecl *), Cached, NoLocationInfo)
 SWIFT_REQUEST(TypeChecker, TypeCheckFunctionBodyAtLocRequest,

lib/AST/AutoDiff.cpp

@@ -421,3 +421,128 @@ void DerivativeFunctionTypeError::log(raw_ostream &OS) const {
   }
   }
 }
+
+bool swift::operator==(const TangentPropertyInfo::Error &lhs,
+                       const TangentPropertyInfo::Error &rhs) {
+  if (lhs.kind != rhs.kind)
+    return false;
+  switch (lhs.kind) {
+  case TangentPropertyInfo::Error::Kind::NoDerivativeOriginalProperty:
+  case TangentPropertyInfo::Error::Kind::NominalParentNotDifferentiable:
+  case TangentPropertyInfo::Error::Kind::OriginalPropertyNotDifferentiable:
+  case TangentPropertyInfo::Error::Kind::ParentTangentVectorNotStruct:
+  case TangentPropertyInfo::Error::Kind::TangentPropertyNotFound:
+  case TangentPropertyInfo::Error::Kind::TangentPropertyNotStored:
+    return true;
+  case TangentPropertyInfo::Error::Kind::TangentPropertyWrongType:
+    return lhs.getType()->isEqual(rhs.getType());
+  }
+}
+
+void swift::simple_display(llvm::raw_ostream &os, TangentPropertyInfo info) {
+  os << "{ ";
+  os << "tangent property: "
+     << (info.tangentProperty ? info.tangentProperty->printRef() : "null");
+  if (info.error) {
+    os << ", error: ";
+    switch (info.error->kind) {
+    case TangentPropertyInfo::Error::Kind::NoDerivativeOriginalProperty:
+      os << "'@noDerivative' original property has no tangent property";
+      break;
+    case TangentPropertyInfo::Error::Kind::NominalParentNotDifferentiable:
+      os << "nominal parent does not conform to 'Differentiable'";
+      break;
+    case TangentPropertyInfo::Error::Kind::OriginalPropertyNotDifferentiable:
+      os << "original property type does not conform to 'Differentiable'";
+      break;
+    case TangentPropertyInfo::Error::Kind::ParentTangentVectorNotStruct:
+      os << "'TangentVector' type is not a struct";
+      break;
+    case TangentPropertyInfo::Error::Kind::TangentPropertyNotFound:
+      os << "'TangentVector' struct does not have stored property with the "
+            "same name as the original property";
+      break;
+    case TangentPropertyInfo::Error::Kind::TangentPropertyWrongType:
+      os << "tangent property's type is not equal to the original property's "
+            "'TangentVector' type";
+      break;
+    case TangentPropertyInfo::Error::Kind::TangentPropertyNotStored:
+      os << "'TangentVector' property '" << info.tangentProperty->getName()
+         << "' is not a stored property";
+      break;
+    }
+  }
+  os << " }";
+}
+
+TangentPropertyInfo
+TangentStoredPropertyRequest::evaluate(Evaluator &evaluator,
+                                       VarDecl *originalField) const {
+  assert(originalField->hasStorage() && originalField->isInstanceMember() &&
+         "Expected stored property");
+  auto *parentDC = originalField->getDeclContext();
+  assert(parentDC->isTypeContext());
+  auto parentType = parentDC->getDeclaredTypeInContext();
+  auto *moduleDecl = originalField->getModuleContext();
+  auto parentTan = parentType->getAutoDiffTangentSpace(
+      LookUpConformanceInModule(moduleDecl));
+  // Error if parent nominal type does not conform to `Differentiable`.
+  if (!parentTan) {
+    return TangentPropertyInfo(
+        TangentPropertyInfo::Error::Kind::NominalParentNotDifferentiable);
+  }
+  // Error if original stored property is `@noDerivative`.
+  if (originalField->getAttrs().hasAttribute<NoDerivativeAttr>()) {
+    return TangentPropertyInfo(
+        TangentPropertyInfo::Error::Kind::NoDerivativeOriginalProperty);
+  }
+  // Error if original property's type does not conform to `Differentiable`.
+  auto originalFieldTan = originalField->getType()->getAutoDiffTangentSpace(
+      LookUpConformanceInModule(moduleDecl));
+  if (!originalFieldTan) {
+    return TangentPropertyInfo(
+        TangentPropertyInfo::Error::Kind::OriginalPropertyNotDifferentiable);
+  }
+  auto parentTanType = parentTan->getType();
+  auto *parentTanStruct = parentTanType->getStructOrBoundGenericStruct();
+  // Error if parent `TangentVector` is not a struct.
+  if (!parentTanStruct) {
+    return TangentPropertyInfo(
+        TangentPropertyInfo::Error::Kind::ParentTangentVectorNotStruct);
+  }
+  // Find the corresponding field in the tangent space.
+  VarDecl *tanField = nullptr;
+  // If `TangentVector` is the original struct, then the tangent property is the
+  // original property.
+  if (parentTanStruct == parentDC->getSelfStructDecl()) {
+    tanField = originalField;
+  }
+  // Otherwise, look up the field by name.
+  else {
+    auto tanFieldLookup =
+        parentTanStruct->lookupDirect(originalField->getName());
+    llvm::erase_if(tanFieldLookup,
+                   [](ValueDecl *v) { return !isa<VarDecl>(v); });
+    // Error if tangent property could not be found.
+    if (tanFieldLookup.empty()) {
+      return TangentPropertyInfo(
+          TangentPropertyInfo::Error::Kind::TangentPropertyNotFound);
+    }
+    tanField = cast<VarDecl>(tanFieldLookup.front());
+  }
+  // Error if tangent property's type is not equal to the original property's
+  // `TangentVector` type.
+  auto originalFieldTanType = originalFieldTan->getType();
+  if (!originalFieldTanType->isEqual(tanField->getType())) {
+    return TangentPropertyInfo(
+        TangentPropertyInfo::Error::Kind::TangentPropertyWrongType,
+        originalFieldTanType);
+  }
+  // Error if tangent property is not a stored property.
+  if (!tanField->hasStorage()) {
+    return TangentPropertyInfo(
+        TangentPropertyInfo::Error::Kind::TangentPropertyNotStored);
+  }
+  // Otherwise, tangent property is valid.
+  return TangentPropertyInfo(tanField);
+}

lib/Sema/DerivedConformanceDifferentiable.cpp

@@ -26,6 +26,7 @@
 #include "swift/AST/PropertyWrappers.h"
 #include "swift/AST/ProtocolConformance.h"
 #include "swift/AST/Stmt.h"
+#include "swift/AST/TypeCheckRequests.h"
 #include "swift/AST/Types.h"
 #include "DerivedConformances.h"
 
@@ -646,39 +647,43 @@ getOrSynthesizeTangentVectorStruct(DerivedConformance &derived, Identifier id) {
   structDecl->setImplicit();
   structDecl->copyFormalAccessFrom(nominal, /*sourceIsParentContext*/ true);
 
-  // Add members to `TangentVector` struct.
+  // Add stored properties to the `TangentVector` struct.
   for (auto *member : diffProperties) {
-    // Add this member's corresponding `TangentVector` type to the parent's
-    // `TangentVector` struct.
-    // Note: `newMember` is not marked as implicit here, because that
-    // incorrectly affects memberwise initializer synthesis.
-    auto *newMember = new (C) VarDecl(
+    // Add a tangent stored property to the `TangentVector` struct, with the
+    // name and `TangentVector` type of the original property.
+    auto *tangentProperty = new (C) VarDecl(
         member->isStatic(), member->getIntroducer(), member->isCaptureList(),
         /*NameLoc*/ SourceLoc(), member->getName(), structDecl);
-
+    // Note: `tangentProperty` is not marked as implicit here, because that
+    // incorrectly affects memberwise initializer synthesis.
     auto memberContextualType =
         parentDC->mapTypeIntoContext(member->getValueInterfaceType());
     auto memberTanType =
         getTangentVectorInterfaceType(memberContextualType, parentDC);
-    newMember->setInterfaceType(memberTanType);
-    Pattern *memberPattern = NamedPattern::createImplicit(C, newMember);
+    tangentProperty->setInterfaceType(memberTanType);
+    Pattern *memberPattern = NamedPattern::createImplicit(C, tangentProperty);
     memberPattern->setType(memberTanType);
     memberPattern =
         TypedPattern::createImplicit(C, memberPattern, memberTanType);
     memberPattern->setType(memberTanType);
     auto *memberBinding = PatternBindingDecl::createImplicit(
         C, StaticSpellingKind::None, memberPattern, /*initExpr*/ nullptr,
         structDecl);
-    structDecl->addMember(newMember);
+    structDecl->addMember(tangentProperty);
     structDecl->addMember(memberBinding);
-    newMember->copyFormalAccessFrom(member, /*sourceIsParentContext*/ true);
-    newMember->setSetterAccess(member->getFormalAccess());
-
-    // Now that this member is in the `TangentVector` type, it should be marked
-    // `@differentiable` so that the differentiation transform will synthesize
-    // derivative functions for it. We only add this to public stored
-    // properties, because their access outside the module will go through a
-    // call to the getter.
+    tangentProperty->copyFormalAccessFrom(member,
+                                          /*sourceIsParentContext*/ true);
+    tangentProperty->setSetterAccess(member->getFormalAccess());
+
+    // Cache the tangent property.
+    C.evaluator.cacheOutput(TangentStoredPropertyRequest{member},
+                            TangentPropertyInfo(tangentProperty));
+
+    // Now that the original property has a corresponding tangent property, it
+    // should be marked `@differentiable` so that the differentiation transform
+    // will synthesize derivative functions for its accessors. We only add this
+    // to public stored properties, because their access outside the module will
+    // go through accessor declarations.
     if (member->getEffectiveAccess() > AccessLevel::Internal &&
         !member->getAttrs().hasAttribute<DifferentiableAttr>()) {
       auto *getter = member->getSynthesizedAccessor(AccessorKind::Get);