[AutoDiff upstream] Add flag-gated AdditiveArithmetic derivation. (swiftlang#30628)

dan-zheng · web-flow · commit e5cb871428e3 · 2020-03-25T10:31:50.000-07:00
Add `AdditiveArithmetic` derived conformances for structs, gated by the
`-enable-experimential-additive-arithmetic-derivation` flag.

Structs whose stored properties all conform to `AdditiveArithmetic` can derive
`AdditiveArithmetic`:
- `static var zero: Self`
- `static func +(lhs: Self, rhs: Self) -&gt; Self`
- `static func -(lhs: Self, rhs: Self) -&gt; Self`
- An "effective memberwise initializer":
  - Either a synthesized memberwise initializer or a user-defined initializer
    with the same type.

Effective memberwise initializers are used only by derived conformances for
`Self`-returning protocol requirements like `AdditiveArithmetic.+`, which
require memberwise initialization.

Resolves TF-844.
Unblocks TF-845: upstream `Differentiable` derived conformances.
diff --git a/include/swift/AST/Decl.h b/include/swift/AST/Decl.h
@@ -3516,6 +3516,27 @@ class NominalTypeDecl : public GenericTypeDecl, public IterableDeclContext {
   /// or \c nullptr if it does not have one.
   ConstructorDecl *getMemberwiseInitializer() const;
 
+  /// Retrieves the effective memberwise initializer for this declaration, or
+  /// \c nullptr if it does not have one.
+  ///
+  /// An effective memberwise initializer is either a synthesized memberwise
+  /// initializer or a user-defined initializer with the same type.
+  ///
+  /// The access level of the memberwise initializer is set to the minimum of:
+  /// - Public, by default. This enables public nominal types to have public
+  ///   memberwise initializers.
+  ///   - The `public` default is important for synthesized member types, e.g.
+  ///     `TangentVector` structs synthesized during `Differentiable` derived
+  ///     conformances. Manually extending these types to define a public
+  ///     memberwise initializer causes a redeclaration error.
+  /// - The minimum access level of memberwise-initialized properties in the
+  ///   nominal type declaration.
+  ///
+  /// Effective memberwise initializers are used only by derived conformances
+  /// for `Self`-returning protocol requirements like `AdditiveArithmetic.+`.
+  /// Such derived conformances require memberwise initialization.
+  ConstructorDecl *getEffectiveMemberwiseInitializer();
+
   /// Whether this declaration has a synthesized zero parameter default
   /// initializer.
   bool hasDefaultInitializer() const;
diff --git a/include/swift/AST/DiagnosticsSema.def b/include/swift/AST/DiagnosticsSema.def
@@ -2674,6 +2674,8 @@ ERROR(cannot_synthesize_in_crossfile_extension,none,
       "implementation of %0 cannot be automatically synthesized in an extension "
       "in a different file to the type", (Type))
 
+ERROR(broken_additive_arithmetic_requirement,none,
+      "AdditiveArithmetic protocol is broken: unexpected requirement", ())
 ERROR(broken_case_iterable_requirement,none,
       "CaseIterable protocol is broken: unexpected requirement", ())
 ERROR(broken_raw_representable_requirement,none,
diff --git a/include/swift/AST/KnownIdentifiers.def b/include/swift/AST/KnownIdentifiers.def
@@ -133,8 +133,6 @@ IDENTIFIER(withKeywordArguments)
 IDENTIFIER(wrapped)
 IDENTIFIER(wrappedValue)
 IDENTIFIER(wrapperValue)
-IDENTIFIER(differential)
-IDENTIFIER(pullback)
 
 // Kinds of layout constraints
 IDENTIFIER_WITH_NAME(UnknownLayout, "_UnknownLayout")
@@ -206,7 +204,10 @@ IDENTIFIER_(nsError)
 IDENTIFIER(OSLogMessage)
 
 // Differentiable programming
+IDENTIFIER(differential)
+IDENTIFIER(pullback)
 IDENTIFIER(TangentVector)
+IDENTIFIER(zero)
 
 #undef IDENTIFIER
 #undef IDENTIFIER_
diff --git a/include/swift/AST/KnownProtocols.def b/include/swift/AST/KnownProtocols.def
@@ -78,13 +78,13 @@ PROTOCOL_(SwiftNewtypeWrapper)
 PROTOCOL(CodingKey)
 PROTOCOL(Encodable)
 PROTOCOL(Decodable)
-PROTOCOL(AdditiveArithmetic)
 
 PROTOCOL_(ObjectiveCBridgeable)
 PROTOCOL_(DestructorSafeContainer)
 
 PROTOCOL(StringInterpolationProtocol)
 
+PROTOCOL(AdditiveArithmetic)
 PROTOCOL(Differentiable)
 
 EXPRESSIBLE_BY_LITERAL_PROTOCOL(ExpressibleByArrayLiteral, "Array", false)
diff --git a/include/swift/AST/TypeCheckRequests.h b/include/swift/AST/TypeCheckRequests.h
@@ -1600,6 +1600,31 @@ class SynthesizeMemberwiseInitRequest
   bool isCached() const { return true; }
 };
 
+/// Resolves the effective memberwise initializer for a given type.
+///
+/// An effective memberwise initializer is either a synthesized memberwise
+/// initializer or a user-defined initializer with the same type.
+///
+/// See `NominalTypeDecl::getEffectiveMemberwiseInitializer` for details.
+class ResolveEffectiveMemberwiseInitRequest
+    : public SimpleRequest<ResolveEffectiveMemberwiseInitRequest,
+                           ConstructorDecl *(NominalTypeDecl *),
+                           CacheKind::Cached> {
+public:
+  using SimpleRequest::SimpleRequest;
+
+private:
+  friend SimpleRequest;
+
+  // Evaluation.
+  llvm::Expected<ConstructorDecl *> evaluate(Evaluator &evaluator,
+                                             NominalTypeDecl *decl) const;
+
+public:
+  // Caching.
+  bool isCached() const { return true; }
+};
+
 /// Checks whether this type has a synthesized zero parameter default
 /// initializer.
 class HasDefaultInitRequest
diff --git a/include/swift/AST/TypeCheckerTypeIDZone.def b/include/swift/AST/TypeCheckerTypeIDZone.def
@@ -214,6 +214,8 @@ SWIFT_REQUEST(TypeChecker, SPIGroupsRequest,
               Cached, NoLocationInfo)
 SWIFT_REQUEST(TypeChecker, SynthesizeMemberwiseInitRequest,
               ConstructorDecl *(NominalTypeDecl *), Cached, NoLocationInfo)
+SWIFT_REQUEST(TypeChecker, ResolveEffectiveMemberwiseInitRequest,
+              ConstructorDecl *(NominalTypeDecl *), Cached, NoLocationInfo)
 SWIFT_REQUEST(TypeChecker, HasDefaultInitRequest,
               bool(NominalTypeDecl *), Cached, NoLocationInfo)
 SWIFT_REQUEST(TypeChecker, SynthesizeDefaultInitRequest,
diff --git a/include/swift/Basic/LangOptions.h b/include/swift/Basic/LangOptions.h
@@ -327,6 +327,10 @@ namespace swift {
     /// `@differentiable` declaration attribute, etc.
     bool EnableExperimentalDifferentiableProgramming = false;
 
+    /// Whether to enable experimental `AdditiveArithmetic` derived
+    /// conformances.
+    bool EnableExperimentalAdditiveArithmeticDerivedConformances = false;
+
     /// Enable verification when every SubstitutionMap is constructed.
     bool VerifyAllSubstitutionMaps = false;
 
diff --git a/include/swift/Option/Options.td b/include/swift/Option/Options.td
@@ -495,10 +495,16 @@ def disable_bridging_pch : Flag<["-"], "disable-bridging-pch">,
   HelpText<"Disable automatic generation of bridging PCH files">;
 
 // Experimental feature options
-def enable_experimental_differentiable_programming : Flag<["-"], "enable-experimental-differentiable-programming">,
+def enable_experimental_differentiable_programming :
+  Flag<["-"], "enable-experimental-differentiable-programming">,
   Flags<[FrontendOption]>,
   HelpText<"Enable experimental differentiable programming features">;
 
+def enable_experimental_additive_arithmetic_derivation :
+  Flag<["-"], "enable-experimental-additive-arithmetic-derivation">,
+  Flags<[FrontendOption]>,
+  HelpText<"Enable experimental 'AdditiveArithmetic' derived conformances">;
+
 def enable_experimental_concise_pound_file : Flag<["-"],
     "enable-experimental-concise-pound-file">,
   Flags<[FrontendOption]>,
diff --git a/lib/AST/Decl.cpp b/lib/AST/Decl.cpp
@@ -4059,6 +4059,14 @@ ConstructorDecl *NominalTypeDecl::getMemberwiseInitializer() const {
       ctx.evaluator, SynthesizeMemberwiseInitRequest{mutableThis}, nullptr);
 }
 
+ConstructorDecl *NominalTypeDecl::getEffectiveMemberwiseInitializer() {
+  auto &ctx = getASTContext();
+  auto *mutableThis = const_cast<NominalTypeDecl *>(this);
+  return evaluateOrDefault(ctx.evaluator,
+                           ResolveEffectiveMemberwiseInitRequest{mutableThis},
+                           nullptr);
+}
+
 bool NominalTypeDecl::hasDefaultInitializer() const {
   // Currently only structs and classes can have default initializers.
   if (!isa<StructDecl>(this) && !isa<ClassDecl>(this))
diff --git a/lib/Frontend/CompilerInvocation.cpp b/lib/Frontend/CompilerInvocation.cpp
@@ -442,8 +442,15 @@ static bool ParseLangArgs(LangOptions &Opts, ArgList &Args,
   if (Args.hasArg(OPT_fine_grained_dependency_include_intrafile))
     Opts.FineGrainedDependenciesIncludeIntrafileOnes = true;
 
-  if (Args.hasArg(OPT_enable_experimental_differentiable_programming))
+  if (Args.hasArg(OPT_enable_experimental_additive_arithmetic_derivation))
+    Opts.EnableExperimentalAdditiveArithmeticDerivedConformances = true;
+
+  if (Args.hasArg(OPT_enable_experimental_differentiable_programming)) {
     Opts.EnableExperimentalDifferentiableProgramming = true;
+    // Differentiable programming implies `AdditiveArithmetic` derived
+    // conformances.
+    Opts.EnableExperimentalAdditiveArithmeticDerivedConformances = true;
+  }
 
   Opts.DebuggerSupport |= Args.hasArg(OPT_debugger_support);
   if (Opts.DebuggerSupport)
diff --git a/lib/IRGen/GenMeta.cpp b/lib/IRGen/GenMeta.cpp
@@ -4634,8 +4634,8 @@ SpecialProtocol irgen::getSpecialProtocolID(ProtocolDecl *P) {
   case KnownProtocolKind::Encodable:
   case KnownProtocolKind::Decodable:
   case KnownProtocolKind::StringInterpolationProtocol:
-  case KnownProtocolKind::Differentiable:
   case KnownProtocolKind::AdditiveArithmetic:
+  case KnownProtocolKind::Differentiable:
     return SpecialProtocol::None;
   }
 
diff --git a/lib/Sema/CMakeLists.txt b/lib/Sema/CMakeLists.txt
@@ -1,4 +1,3 @@
-
 add_swift_host_library(swiftSema STATIC
   BuilderTransform.cpp
   CSApply.cpp
@@ -16,6 +15,7 @@ add_swift_host_library(swiftSema STATIC
   ConstraintLocator.cpp
   ConstraintSystem.cpp
   DebuggerTestingTransform.cpp
+  DerivedConformanceAdditiveArithmetic.cpp
   DerivedConformanceCaseIterable.cpp
   DerivedConformanceCodable.cpp
   DerivedConformanceCodingKey.cpp
diff --git a/lib/Sema/CodeSynthesis.cpp b/lib/Sema/CodeSynthesis.cpp
@@ -1201,6 +1201,93 @@ SynthesizeMemberwiseInitRequest::evaluate(Evaluator &evaluator,
   return ctor;
 }
 
+llvm::Expected<ConstructorDecl *>
+ResolveEffectiveMemberwiseInitRequest::evaluate(Evaluator &evaluator,
+                                                NominalTypeDecl *decl) const {
+  // Compute the access level for the memberwise initializer. The minimum of:
+  // - Public, by default. This enables public nominal types to have public
+  //   memberwise initializers.
+  //   - The `public` default is important for synthesized member types, e.g.
+  //     `TangentVector` structs synthesized during `Differentiable` derived
+  //     conformances. Manually extending these types to define a public
+  //     memberwise initializer causes a redeclaration error.
+  // - The minimum access level of memberwise-initialized properties in the
+  //   nominal type declaration.
+  auto accessLevel = AccessLevel::Public;
+  for (auto *member : decl->getMembers()) {
+    auto *var = dyn_cast<VarDecl>(member);
+    if (!var ||
+        !var->isMemberwiseInitialized(/*preferDeclaredProperties*/ true))
+      continue;
+    accessLevel = std::min(accessLevel, var->getFormalAccess());
+  }
+  auto &ctx = decl->getASTContext();
+
+  // If a memberwise initializer exists, set its access level and return it.
+  if (auto *initDecl = decl->getMemberwiseInitializer()) {
+    initDecl->overwriteAccess(accessLevel);
+    return initDecl;
+  }
+
+  auto isEffectiveMemberwiseInitializer = [&](ConstructorDecl *initDecl) {
+    // Check for `nullptr`.
+    if (!initDecl)
+      return false;
+    // Get all stored properties, excluding `let` properties with initial
+    // values.
+    SmallVector<VarDecl *, 8> storedProperties;
+    for (auto *vd : decl->getStoredProperties()) {
+      if (vd->isLet() && vd->hasInitialValue())
+        continue;
+      storedProperties.push_back(vd);
+    }
+    // Return false if initializer does not have interface type set. It is not
+    // possible to determine whether it is a memberwise initializer.
+    if (!initDecl->hasInterfaceType())
+      return false;
+    auto initDeclType =
+        initDecl->getMethodInterfaceType()->getAs<AnyFunctionType>();
+    // Return false if initializer does not have a valid interface type.
+    if (!initDeclType)
+      return false;
+    // Return false if stored property count does not have parameter count.
+    if (storedProperties.size() != initDeclType->getNumParams())
+      return false;
+    // Return true if all stored property types/names match initializer
+    // parameter types/labels.
+    return llvm::all_of(
+        llvm::zip(storedProperties, initDeclType->getParams()),
+        [&](std::tuple<VarDecl *, AnyFunctionType::Param> pair) {
+          auto *storedProp = std::get<0>(pair);
+          auto param = std::get<1>(pair);
+          return storedProp->getInterfaceType()->isEqual(
+                     param.getPlainType()) &&
+                 storedProp->getName() == param.getLabel();
+        });
+  };
+
+  // Otherwise, look for a user-defined effective memberwise initializer.
+  ConstructorDecl *memberwiseInitDecl = nullptr;
+  auto initDecls = decl->lookupDirect(DeclBaseName::createConstructor());
+  for (auto *decl : initDecls) {
+    auto *initDecl = dyn_cast<ConstructorDecl>(decl);
+    if (!isEffectiveMemberwiseInitializer(initDecl))
+      continue;
+    assert(!memberwiseInitDecl && "Memberwise initializer already found");
+    memberwiseInitDecl = initDecl;
+  }
+
+  // Otherwise, create a memberwise initializer, set its access level, and
+  // return it.
+  if (!memberwiseInitDecl) {
+    memberwiseInitDecl = createImplicitConstructor(
+        decl, ImplicitConstructorKind::Memberwise, ctx);
+    memberwiseInitDecl->overwriteAccess(accessLevel);
+    decl->addMember(memberwiseInitDecl);
+  }
+  return memberwiseInitDecl;
+}
+
 llvm::Expected<bool>
 HasDefaultInitRequest::evaluate(Evaluator &evaluator,
                                 NominalTypeDecl *decl) const {
@@ -1263,3 +1350,22 @@ SynthesizeDefaultInitRequest::evaluate(Evaluator &evaluator,
   ctor->setBodySynthesizer(synthesizeSingleReturnFunctionBody);
   return ctor;
 }
+
+ValueDecl *swift::getProtocolRequirement(ProtocolDecl *protocol,
+                                         Identifier name) {
+  auto lookup = protocol->lookupDirect(name);
+  // Erase declarations that are not protocol requirements.
+  // This is important for removing default implementations of the same name.
+  llvm::erase_if(lookup, [](ValueDecl *v) {
+    return !isa<ProtocolDecl>(v->getDeclContext()) ||
+           !v->isProtocolRequirement();
+  });
+  assert(lookup.size() == 1 && "Ambiguous protocol requirement");
+  return lookup.front();
+}
+
+bool swift::hasLetStoredPropertyWithInitialValue(NominalTypeDecl *nominal) {
+  return llvm::any_of(nominal->getStoredProperties(), [&](VarDecl *v) {
+    return v->isLet() && v->hasInitialValue();
+  });
+}
diff --git a/lib/Sema/CodeSynthesis.h b/lib/Sema/CodeSynthesis.h
@@ -60,6 +60,12 @@ Expr *buildSelfReference(VarDecl *selfDecl,
 Expr *buildArgumentForwardingExpr(ArrayRef<ParamDecl*> params,
                                   ASTContext &ctx);
 
+/// Returns the protocol requirement with the specified name.
+ValueDecl *getProtocolRequirement(ProtocolDecl *protocol, Identifier name);
+
+// Returns true if given nominal type has a `let` stored with an initial value.
+bool hasLetStoredPropertyWithInitialValue(NominalTypeDecl *nominal);
+
 } // end namespace swift
 
 #endif
diff --git a/lib/Sema/DerivedConformanceAdditiveArithmetic.cpp b/lib/Sema/DerivedConformanceAdditiveArithmetic.cpp
diff --git a/lib/Sema/DerivedConformances.cpp b/lib/Sema/DerivedConformances.cpp
diff --git a/lib/Sema/DerivedConformances.h b/lib/Sema/DerivedConformances.h
diff --git a/lib/Sema/TypeCheckProtocol.cpp b/lib/Sema/TypeCheckProtocol.cpp
diff --git a/test/AutoDiff/Sema/DerivedConformances/Inputs/struct_additive_arithmetic_other_module.swift b/test/AutoDiff/Sema/DerivedConformances/Inputs/struct_additive_arithmetic_other_module.swift
diff --git a/test/AutoDiff/Sema/DerivedConformances/struct_additive_arithmetic.swift b/test/AutoDiff/Sema/DerivedConformances/struct_additive_arithmetic.swift
