[Sema] Differentiable conformance derivation for class types. (#25914)

dan-zheng · web-flow · commit d214a911e3b7 · 2019-07-01T15:54:20.000-07:00
`Differentiable` derived conformances now supports class types. Synthesis works just like for struct types, except `TangentVector = Self` is never synthesized even if `Self` conforms to `AdditiveArithmetic`. Class differentiation support requires further differentiation transform changes. Resolves TF-630.
diff --git a/include/swift/AST/DiagnosticsSema.def b/include/swift/AST/DiagnosticsSema.def
@@ -2869,9 +2869,9 @@ ERROR(compiler_evaluable_ref_non_compiler_evaluable,none,
       "@compilerEvaluable functions may not reference non-@compilerEvaluable functions", ())
 
 // @noDerivative attribute
-ERROR(noderivative_only_on_stored_properties_in_differentiable_structs,none,
-      "'@noDerivative' is only allowed on stored properties in structure types "
-      "that declare a conformance to 'Differentiable'", ())
+ERROR(noderivative_only_on_differentiable_struct_or_class_fields,none,
+      "'@noDerivative' is only allowed on stored properties in structure or "
+      "class types that declare a conformance to 'Differentiable'", ())
 
 //------------------------------------------------------------------------------
 // MARK: Type Check Expressions
diff --git a/lib/Sema/DerivedConformanceDifferentiable.cpp b/lib/Sema/DerivedConformanceDifferentiable.cpp
@@ -13,7 +13,7 @@
 // SWIFT_ENABLE_TENSORFLOW
 //
 // This file implements explicit derivation of the Differentiable protocol for
-// struct types.
+// struct and class types.
 //
 //===----------------------------------------------------------------------===//
 
@@ -108,8 +108,7 @@ static StructDecl *getAssociatedStructDecl(DeclContext *DC, Identifier id) {
   auto *diffableProto = C.getProtocol(KnownProtocolKind::Differentiable);
   assert(diffableProto && "`Differentiable` protocol not found");
   auto conf = TypeChecker::conformsToProtocol(DC->getSelfTypeInContext(),
-                                              diffableProto,
-                                              DC, None);
+                                              diffableProto, DC, None);
   assert(conf && "Nominal must conform to `Differentiable`");
   Type assocType = conf->getTypeWitnessByName(DC->getSelfTypeInContext(), id);
   assert(assocType && "`Differentiable` protocol associated type not found");
@@ -120,9 +119,8 @@ static StructDecl *getAssociatedStructDecl(DeclContext *DC, Identifier id) {
 
 bool DerivedConformance::canDeriveDifferentiable(NominalTypeDecl *nominal,
                                                  DeclContext *DC) {
-  // Nominal type must be a struct. (Zero stored properties is okay.)
-  auto *structDecl = dyn_cast<StructDecl>(nominal);
-  if (!structDecl)
+  // Nominal type must be a struct or class. (No stored properties is okay.)
+  if (!isa<StructDecl>(nominal) && !isa<ClassDecl>(nominal))
     return false;
   auto &C = nominal->getASTContext();
   auto *lazyResolver = C.getLazyResolver();
@@ -153,8 +151,7 @@ bool DerivedConformance::canDeriveDifferentiable(NominalTypeDecl *nominal,
     //   `X == X.TangentVector`.
     if (nominal->isImplicit() && structDecl == nominal->getDeclContext() &&
         TypeChecker::conformsToProtocol(structDecl->getDeclaredInterfaceType(),
-                                        diffableProto, DC,
-                                        None))
+                                        diffableProto, DC, None))
       return structDecl;
     // 3. Equal nominal (and conform to `AdditiveArithmetic` if flag is true).
     if (structDecl == nominal) {
@@ -199,7 +196,7 @@ bool DerivedConformance::canDeriveDifferentiable(NominalTypeDecl *nominal,
   //     initializers that initialize all stored properties, including initial
   //     value information.
   SmallVector<VarDecl *, 16> diffProperties;
-  getStoredPropertiesForDifferentiation(structDecl, DC, diffProperties);
+  getStoredPropertiesForDifferentiation(nominal, DC, diffProperties);
   return llvm::all_of(diffProperties, [&](VarDecl *v) {
     if (!v->hasInterfaceType())
       lazyResolver->resolveDeclSignature(v);
@@ -325,7 +322,8 @@ static ValueDecl *deriveDifferentiable_method(
                                     /*Throws*/ false, SourceLoc(),
                                     /*GenericParams=*/nullptr, params,
                                     TypeLoc::withoutLoc(returnType), parentDC);
-  funcDecl->setSelfAccessKind(SelfAccessKind::Mutating);
+  if (!nominal->getSelfClassDecl())
+    funcDecl->setSelfAccessKind(SelfAccessKind::Mutating);
   funcDecl->setImplicit();
   funcDecl->setBodySynthesizer(bodySynthesizer.Fn, bodySynthesizer.Context);
 
@@ -804,6 +802,8 @@ static void checkAndDiagnoseImplicitNoDerivative(TypeChecker &TC,
                                                  DeclContext* DC) {
   auto *diffableProto =
       TC.Context.getProtocol(KnownProtocolKind::Differentiable);
+  bool nominalCanDeriveAdditiveArithmetic =
+      DerivedConformance::canDeriveAdditiveArithmetic(nominal, DC);
   for (auto *vd : nominal->getStoredProperties()) {
     if (!vd->hasInterfaceType())
       TC.resolveDeclSignature(vd);
@@ -814,8 +814,7 @@ static void checkAndDiagnoseImplicitNoDerivative(TypeChecker &TC,
       continue;
     // Check whether to diagnose stored property.
     bool conformsToDifferentiable =
-        TC.conformsToProtocol(varType, diffableProto, nominal,
-                              None).hasValue();
+        TC.conformsToProtocol(varType, diffableProto, nominal, None).hasValue();
     // If stored property should not be diagnosed, continue.
     if (conformsToDifferentiable && !vd->isLet())
       continue;
@@ -829,8 +828,6 @@ static void checkAndDiagnoseImplicitNoDerivative(TypeChecker &TC,
     // `Differentiable` protocol requirements all have default implementations
     // when `Self` conforms to `AdditiveArithmetic`, so `Differentiable`
     // derived conformances will no longer be necessary.
-    bool nominalCanDeriveAdditiveArithmetic =
-        DerivedConformance::canDeriveAdditiveArithmetic(nominal, DC);
     if (!conformsToDifferentiable) {
       TC.diagnose(loc,
                   diag::differentiable_nondiff_type_implicit_noderivative_fixit,
@@ -844,7 +841,6 @@ static void checkAndDiagnoseImplicitNoDerivative(TypeChecker &TC,
                 vd->getName(), nominal->getName(),
                 nominalCanDeriveAdditiveArithmetic)
         .fixItInsert(loc, "@noDerivative ");
-
   }
 }
 
@@ -954,6 +950,7 @@ deriveDifferentiable_AssociatedStruct(DerivedConformance &derived,
   bool hasNoDerivativeStoredProp = diffProperties.size() != numStoredProperties;
 
   // Check conditions for returning `Self`.
+  // - `Self` is not a class type.
   // - No `@noDerivative` stored properties exist.
   // - All stored properties must have specified associated type equal to
   //   `Self`.
@@ -971,7 +968,7 @@ deriveDifferentiable_AssociatedStruct(DerivedConformance &derived,
                             parentDC, None);
 
   // Return `Self` if conditions are met.
-  if (!hasNoDerivativeStoredProp &&
+  if (!hasNoDerivativeStoredProp && !nominal->getSelfClassDecl() &&
       (id == C.Id_AllDifferentiableVariables ||
        (allMembersAssocTypeEqualsSelf && nominalConformsToAddArith))) {
     auto selfType = parentDC->getSelfTypeInContext();
diff --git a/lib/Sema/DerivedConformanceRingMathProtocols.cpp b/lib/Sema/DerivedConformanceRingMathProtocols.cpp
@@ -178,8 +178,10 @@ static void deriveBodyMathOperator(AbstractFunctionDecl *funcDecl,
     // If conformance reference is concrete, then use concrete witness
     // declaration for the operator.
     if (confRef->isConcrete())
-      memberOpDecl = confRef->getConcrete()->getWitnessDecl(
-          operatorReq, C.getLazyResolver());
+      if (auto *concreteMemberMethodDecl =
+              confRef->getConcrete()->getWitnessDecl(operatorReq,
+                                                     C.getLazyResolver()))
+        memberOpDecl = concreteMemberMethodDecl;
     assert(memberOpDecl && "Member operator declaration must exist");
     auto memberOpDRE =
         new (C) DeclRefExpr(memberOpDecl, DeclNameLoc(), /*Implicit*/ true);
diff --git a/lib/Sema/DerivedConformances.cpp b/lib/Sema/DerivedConformances.cpp
@@ -554,7 +554,9 @@ DerivedConformance::declareDerivedPropertySetter(TypeChecker &tc,
     /*GenericParams*/ nullptr, params, TypeLoc(), parentDC);
   setterDecl->setImplicit();
   setterDecl->setStatic(isStatic);
-  setterDecl->setSelfAccessKind(SelfAccessKind::Mutating);
+  // Set mutating if parent is not a class.
+  if (!parentDC->getSelfClassDecl())
+    setterDecl->setSelfAccessKind(SelfAccessKind::Mutating);
 
   // If this is supposed to be a final method, mark it as such.
   assert(isFinal || !parentDC->getSelfClassDecl());
@@ -584,6 +586,10 @@ DerivedConformance::declareDerivedProperty(Identifier name,
   VarDecl *propDecl = new (C) VarDecl(/*IsStatic*/isStatic, VarDecl::Specifier::Var,
                                       /*IsCaptureList*/false, SourceLoc(), name,
                                       parentDC);
+  // SWIFT_ENABLE_TENSORFLOW
+  // TODO: Upstream this change to master.
+  if (isFinal && parentDC->getSelfClassDecl())
+    propDecl->getAttrs().add(new (C) FinalAttr(/*Implicit*/ true));
   propDecl->setImplicit();
   propDecl->copyFormalAccessFrom(Nominal, /*sourceIsParentContext*/ true);
   propDecl->setInterfaceType(propertyInterfaceType);
diff --git a/lib/Sema/TypeCheckAttr.cpp b/lib/Sema/TypeCheckAttr.cpp
@@ -3759,20 +3759,20 @@ void AttributeChecker::visitNoDerivativeAttr(NoDerivativeAttr *attr) {
     return;
   if (!vd || vd->isStatic()) {
     diagnoseAndRemoveAttr(attr,
-        diag::noderivative_only_on_stored_properties_in_differentiable_structs);
+        diag::noderivative_only_on_differentiable_struct_or_class_fields);
     return;
   }
-  auto *structDecl = dyn_cast<StructDecl>(vd->getDeclContext());
-  if (!structDecl) {
+  auto *nominal = vd->getDeclContext()->getSelfNominalTypeDecl();
+  if (!nominal || (!isa<StructDecl>(nominal) && !isa<ClassDecl>(nominal))) {
     diagnoseAndRemoveAttr(attr,
-        diag::noderivative_only_on_stored_properties_in_differentiable_structs);
+        diag::noderivative_only_on_differentiable_struct_or_class_fields);
     return;
   }
   if (!conformsToDifferentiable(
-          structDecl->getDeclaredInterfaceType(),
-          structDecl->getDeclContext())) {
+          nominal->getDeclaredInterfaceType(),
+          nominal->getDeclContext())) {
     diagnoseAndRemoveAttr(attr,
-        diag::noderivative_only_on_stored_properties_in_differentiable_structs);
+        diag::noderivative_only_on_differentiable_struct_or_class_fields);
     return;
   }
 }
diff --git a/test/AutoDiff/derived_differentiable.swift b/test/AutoDiff/derived_differentiable.swift
@@ -18,7 +18,7 @@ public struct Foo : Differentiable {
 // CHECK-AST:   public typealias TangentVector = Foo.AllDifferentiableVariables
 
 // CHECK-SILGEN-LABEL: // Foo.a.getter
-// CHECK-SILGEN-NEXT: sil [transparent] [serialized] [differentiable source 0 wrt 0] [ossa] @$s33derived_differentiable_properties3FooV1aSfvg : $@convention(method) (Foo) -> Float
+// CHECK-SILGEN-NEXT: sil [transparent] [serialized] [differentiable source 0 wrt 0] [ossa] @$s22derived_differentiable3FooV1aSfvg : $@convention(method) (Foo) -> Float
 
 struct AdditiveTangentIsSelf : AdditiveArithmetic, Differentiable {
   var a: Float
@@ -82,9 +82,6 @@ struct GenericTanMember<T : Differentiable> : Differentiable, AdditiveArithmetic
   var x: T.TangentVector
 }
 
-// TODO(TF-316): Revisit after `Differentiable` derived conformances behavior is standardized.
-// `AllDifferentiableVariables` and `TangentVector` structs need not both be synthesized.
-
 // CHECK-AST-LABEL: internal struct GenericTanMember<T> : Differentiable, AdditiveArithmetic where T : Differentiable
 // CHECK-AST:   internal var x: T.TangentVector
 // CHECK-AST:   internal init(x: T.TangentVector)
@@ -105,3 +102,32 @@ extension ConditionallyDifferentiable : Differentiable where T : Differentiable
 // CHECK-AST:         public var x: T
 // CHECK-AST:         internal init(x: T)
 // CHECK-AST:       }
+
+// Verify that `TangentVector` is not synthesized to be `Self` for
+// `AdditiveArithmetic`-conforming classes.
+final class AdditiveArithmeticClass<T : AdditiveArithmetic & Differentiable> : AdditiveArithmetic, Differentiable {
+  var x, y: T
+  init(x: T, y: T) {
+    self.x = x
+    self.y = y
+  }
+
+  // Dummy `AdditiveArithmetic` requirements.
+  static func == (lhs: AdditiveArithmeticClass, rhs: AdditiveArithmeticClass) -> Bool {
+    fatalError()
+  }
+  static var zero: AdditiveArithmeticClass {
+    fatalError()
+  }
+  static func + (lhs: AdditiveArithmeticClass, rhs: AdditiveArithmeticClass) -> Self {
+    fatalError()
+  }
+  static func - (lhs: AdditiveArithmeticClass, rhs: AdditiveArithmeticClass) -> Self {
+    fatalError()
+  }
+}
+
+// CHECK-AST-LABEL: final internal class AdditiveArithmeticClass<T> : AdditiveArithmetic, Differentiable where T : AdditiveArithmetic, T : Differentiable {
+// CHECK-AST:         final internal var x: T, y: T
+// CHECK-AST:         internal struct TangentVector : Differentiable, AdditiveArithmetic
+// CHECK-AST:       }
diff --git a/test/AutoDiff/noderivative-attr.swift b/test/AutoDiff/noderivative-attr.swift
@@ -1,10 +1,10 @@
 // RUN: %target-swift-frontend -typecheck -verify %s
 
-// expected-error @+1 {{'@noDerivative' is only allowed on stored properties in structure types that declare a conformance to 'Differentiable'}}
+// expected-error @+1 {{'@noDerivative' is only allowed on stored properties in structure or class types that declare a conformance to 'Differentiable'}}
 @noDerivative var flag: Bool
 
 struct Foo {
-  // expected-error @+1 {{'@noDerivative' is only allowed on stored properties in structure types that declare a conformance to 'Differentiable'}}
+  // expected-error @+1 {{'@noDerivative' is only allowed on stored properties in structure or class types that declare a conformance to 'Differentiable'}}
   @noDerivative var flag: Bool
 }
 
diff --git a/test/Sema/Inputs/class_differentiable_other_module.swift b/test/Sema/Inputs/class_differentiable_other_module.swift
@@ -0,0 +1,9 @@
+// SWIFT_ENABLE_TENSORFLOW
+
+// expected-note @+1 {{type declared here}}
+class OtherFileNonconforming {}
+
+// expected-note @+1 {{type declared here}}
+class GenericOtherFileNonconforming<T : Differentiable> {
+  var x: T
+}
diff --git a/test/Sema/class_differentiable.swift b/test/Sema/class_differentiable.swift
diff --git a/test/Sema/struct_differentiable.swift b/test/Sema/struct_differentiable.swift
