Add KeyPathIterable protocol and synthesis. (#21557)

* Add `KeyPathIterable` protocol and synthesis.

The `KeyPathIterable` protocol represents types whose values provide
custom key paths to properties or elements.

For types that conform to `KeyPathIterable`, the compiler can synthesize
a default implementation of `allKeyPaths` based on the type's stored
properties.

---

`KeyPathIterable` enables generic machine learning optimizers and will
be used to replace `ParameterGroup`.

The `ParameterGroup` protocol requirements are not sufficiently general
for many use cases: it is limited to a single parameter type and does not
support joint iteration over parameters for multiple `ParameterGroup`
instances.

`KeyPathIterable` solves these problems. In addition, it is a generally
useful language feature that can model both static stored properties
and custom dynamic properties.

* Address feedback from @rxwei.
- Remove `@inlinable` from `KeyPathIterable` extension methods.
- Enable `KeyPathIterable` synthesis for empty structs, add test.
- Add comment by @rxwei about supporting synthesis for classes in the future.

* Address feedback from @rxwei.

Mark synthesized `allKeyPaths` declaration with `@inlinable`.
Verified that `@inlinable` is correctly added in SILGen:

```
struct Parameters : KeyPathIterable {
  @sil_stored var w: Float { get set }
  @sil_stored var b: Float { get set }
  init(w: Float, b: Float)
  typealias AllKeyPaths = [PartialKeyPath<Parameters>]
  @inlinable var allKeyPaths: [PartialKeyPath<Parameters>] { get }
}
```

Author: dan-zheng

include/swift/AST/DiagnosticsSema.def

@@ -2405,6 +2405,8 @@ ERROR(broken_encodable_requirement,none,
 ERROR(broken_decodable_requirement,none,
       "Decodable protocol is broken: unexpected requirement", ())
 // SWIFT_ENABLE_TENSORFLOW
+ERROR(broken_key_path_iterable_requirement,none,
+      "KeyPathIterable protocol is broken: unexpected requirement", ())
 ERROR(broken_parameter_group_requirement,none,
       "ParameterGroup protocol is broken: unexpected requirement", ())
 ERROR(broken_parameterized_requirement,none,

include/swift/AST/KnownIdentifiers.def

@@ -47,7 +47,6 @@ IDENTIFIER(decode)
 IDENTIFIER(decodeIfPresent)
 IDENTIFIER(Decoder)
 IDENTIFIER(decoder)
-// SWIFT_ENABLE_TENSORFLOW
 IDENTIFIER(dynamicCallable)
 IDENTIFIER(dynamicMember)
 IDENTIFIER(Element)
@@ -118,6 +117,11 @@ IDENTIFIER_WITH_NAME(value_, "_value")
 IDENTIFIER(with)
 
 // SWIFT_ENABLE_TENSORFLOW
+IDENTIFIER(AllKeyPaths)
+IDENTIFIER(allKeyPaths)
+IDENTIFIER(recursivelyAllKeyPaths)
+IDENTIFIER(allWritableKeyPaths)
+IDENTIFIER(recursivelyAllWritableKeyPaths)
 IDENTIFIER(allParameters)
 IDENTIFIER(Parameter)
 IDENTIFIER(Parameters)

include/swift/AST/KnownProtocols.def

@@ -70,6 +70,7 @@ PROTOCOL(Decodable)
 PROTOCOL(AdditiveArithmetic)
 PROTOCOL(Numeric)
 PROTOCOL(FloatingPoint)
+PROTOCOL(KeyPathIterable)
 PROTOCOL(ParameterGroup)
 PROTOCOL(Parameterized)
 PROTOCOL(TensorArrayProtocol)

lib/IRGen/GenMeta.cpp

@@ -4061,6 +4061,7 @@ SpecialProtocol irgen::getSpecialProtocolID(ProtocolDecl *P) {
   case KnownProtocolKind::FloatingPoint:
   case KnownProtocolKind::AdditiveArithmetic:
   case KnownProtocolKind::Numeric:
+  case KnownProtocolKind::KeyPathIterable:
   case KnownProtocolKind::ParameterGroup:
   case KnownProtocolKind::Parameterized:
   case KnownProtocolKind::TensorArrayProtocol:

lib/Sema/CMakeLists.txt

@@ -27,6 +27,7 @@ add_swift_host_library(swiftSema STATIC
   DerivedConformanceEquatableHashable.cpp
   DerivedConformanceError.cpp
   # SWIFT_ENABLE_TENSORFLOW
+  DerivedConformanceKeyPathIterable.cpp
   DerivedConformanceParameterGroup.cpp
   DerivedConformanceParameterized.cpp
   DerivedConformanceRawRepresentable.cpp

lib/Sema/DerivedConformanceKeyPathIterable.cpp

@@ -0,0 +1,140 @@
+//===--- DerivedConformanceKeyPathIterable.cpp ----------------------------===//
+//
+// This source file is part of the Swift.org open source project
+//
+// Copyright (c) 2014 - 2017 Apple Inc. and the Swift project authors
+// Licensed under Apache License v2.0 with Runtime Library Exception
+//
+// See https://swift.org/LICENSE.txt for license information
+// See https://swift.org/CONTRIBUTORS.txt for the list of Swift project authors
+//
+//===----------------------------------------------------------------------===//
+//
+// This file implements explicit derivation of the KeyPathIterable protocol for
+// a nominal type.
+//
+//===----------------------------------------------------------------------===//
+
+#include "CodeSynthesis.h"
+#include "TypeChecker.h"
+#include "swift/AST/Decl.h"
+#include "swift/AST/Expr.h"
+#include "swift/AST/GenericSignature.h"
+#include "swift/AST/Module.h"
+#include "swift/AST/ParameterList.h"
+#include "swift/AST/Pattern.h"
+#include "swift/AST/ProtocolConformance.h"
+#include "swift/AST/Stmt.h"
+#include "swift/AST/Types.h"
+#include "DerivedConformances.h"
+
+using namespace swift;
+
+bool DerivedConformance::canDeriveKeyPathIterable(NominalTypeDecl *nominal) {
+  // Note: we could extend synthesis to support classes.
+  // Subclasses need to append `allKeyPaths` to `super.allKeyPaths`.
+  return isa<StructDecl>(nominal);
+}
+
+// Compute `PartialKeyPathType<Nominal>`, bound to the given nominal
+// declaration's type.
+static Type computePartialKeyPathType(NominalTypeDecl *nominal) {
+  auto &C = nominal->getASTContext();
+  auto nominalType = nominal->getDeclaredInterfaceType();
+  if (!nominalType || nominalType->hasError())
+    return nullptr;
+  auto *partialKeyPathDecl = cast<ClassDecl>(C.getPartialKeyPathDecl());
+  return BoundGenericClassType::get(partialKeyPathDecl, /*parent*/ Type(),
+                                    {nominal->getDeclaredInterfaceType()});
+}
+
+// Compute `AllKeyPaths` associated type for the given nominal declaration.
+// It should be `[PartialKeyPath<Nominal>]`.
+static ArraySliceType *computeAllKeyPathsType(NominalTypeDecl *nominal) {
+  auto partialKeyPathType = computePartialKeyPathType(nominal);
+  return ArraySliceType::get(partialKeyPathType);
+}
+
+// Synthesize body for the `allKeyPaths` computed property getter.
+static void
+deriveBodyKeyPathIterable_allKeyPaths(AbstractFunctionDecl *funcDecl) {
+  auto *nominal = funcDecl->getDeclContext()->getSelfNominalTypeDecl();
+  auto &C = nominal->getASTContext();
+
+  auto *nominalTypeExpr = TypeExpr::createForDecl(SourceLoc(), nominal,
+                                                  funcDecl, /*Implicit*/ true);
+
+  // Create array of key path expressions to stored properties.
+  llvm::SmallVector<Expr *, 2> keyPathExprs;
+  for (auto member : nominal->getStoredProperties()) {
+    auto *dotExpr = new (C)
+        UnresolvedDotExpr(nominalTypeExpr, SourceLoc(), member->getFullName(),
+                          DeclNameLoc(), /*Implicit*/ true);
+    auto *keyPathExpr =
+        new (C) KeyPathExpr(SourceLoc(), dotExpr, nullptr, /*Implicit*/ true);
+    keyPathExprs.push_back(keyPathExpr);
+  }
+  // Return array of all key path expressions.
+  auto keyPathsArrayExpr =
+      ArrayExpr::create(C, SourceLoc(), keyPathExprs, {}, SourceLoc());
+  auto *returnStmt = new (C) ReturnStmt(SourceLoc(), keyPathsArrayExpr);
+  auto *body = BraceStmt::create(C, SourceLoc(), {returnStmt}, SourceLoc(),
+                                 /*Implicit*/ true);
+  funcDecl->setBody(BraceStmt::create(C, SourceLoc(), {body}, SourceLoc(),
+                                      /*Implicit*/ true));
+}
+
+// Synthesize the `allKeyPaths` computed property declaration.
+static ValueDecl *
+deriveKeyPathIterable_allKeyPaths(DerivedConformance &derived) {
+  auto nominal = derived.Nominal;
+  auto &C = derived.TC.Context;
+
+  auto returnInterfaceTy = computeAllKeyPathsType(nominal);
+  auto returnTy =
+      derived.getConformanceContext()->mapTypeIntoContext(returnInterfaceTy);
+
+  // Create `allKeyPaths` property declaration.
+  VarDecl *allKeyPathsDecl;
+  PatternBindingDecl *pbDecl;
+  std::tie(allKeyPathsDecl, pbDecl) = derived.declareDerivedProperty(
+      C.Id_allKeyPaths, returnInterfaceTy, returnTy, /*isStatic*/ false,
+      /*isFinal*/ true);
+
+  // Add `@inlinable` to the `allKeyPaths` declaration.
+  allKeyPathsDecl->getAttrs().add(new (C) InlinableAttr(/*implicit*/ true));
+
+  // Create `allKeyPaths` getter.
+  auto *getterDecl = derived.declareDerivedPropertyGetter(
+      derived.TC, allKeyPathsDecl, returnTy);
+  getterDecl->setBodySynthesizer(deriveBodyKeyPathIterable_allKeyPaths);
+  allKeyPathsDecl->setAccessors(StorageImplInfo::getImmutableComputed(),
+                                SourceLoc(), {getterDecl}, SourceLoc());
+  derived.addMembersToConformanceContext({getterDecl, allKeyPathsDecl, pbDecl});
+
+  return allKeyPathsDecl;
+}
+
+static Type deriveKeyPathIterable_AllKeyPaths(DerivedConformance &derived) {
+  auto *rawInterfaceType = computeAllKeyPathsType(derived.Nominal);
+  return derived.getConformanceContext()->mapTypeIntoContext(rawInterfaceType);
+}
+
+ValueDecl *DerivedConformance::deriveKeyPathIterable(ValueDecl *requirement) {
+  if (requirement->getBaseName() == TC.Context.Id_allKeyPaths) {
+    return deriveKeyPathIterable_allKeyPaths(*this);
+  }
+  TC.diagnose(requirement->getLoc(),
+              diag::broken_key_path_iterable_requirement);
+  return nullptr;
+}
+
+Type DerivedConformance::deriveKeyPathIterable(
+    AssociatedTypeDecl *requirement) {
+  if (requirement->getBaseName() == TC.Context.Id_AllKeyPaths) {
+    return deriveKeyPathIterable_AllKeyPaths(*this);
+  }
+  TC.diagnose(requirement->getLoc(),
+              diag::broken_key_path_iterable_requirement);
+  return nullptr;
+}

lib/Sema/DerivedConformances.cpp

@@ -62,6 +62,10 @@ bool DerivedConformance::derivesProtocolConformance(DeclContext *DC,
     return canDeriveHashable(Nominal);
   }
 
+  // SWIFT_ENABLE_TENSORFLOW
+  if (*knownProtocol == KnownProtocolKind::KeyPathIterable)
+    return canDeriveKeyPathIterable(Nominal);
+
   // SWIFT_ENABLE_TENSORFLOW
   // The only requirement for deriving Parameterized is that there exist some
   // stored properties marked with @TFParameter. The `Parameters` struct can
@@ -206,6 +210,11 @@ ValueDecl *DerivedConformance::getDerivableRequirement(TypeChecker &tc,
     if (name.isSimpleName(ctx.Id_intValue))
       return getRequirement(KnownProtocolKind::CodingKey);
 
+    // SWIFT_ENABLE_TENSORFLOW
+    // KeyPathIterable.allKeyPaths
+    if (name.isSimpleName(ctx.Id_allKeyPaths))
+      return getRequirement(KnownProtocolKind::KeyPathIterable);
+
     // SWIFT_ENABLE_TENSORFLOW
     // Parameterized.allParameters
     if (name.isSimpleName(ctx.Id_allParameters))
@@ -279,6 +288,11 @@ ValueDecl *DerivedConformance::getDerivableRequirement(TypeChecker &tc,
     if (name.isSimpleName(ctx.Id_AllCases))
       return getRequirement(KnownProtocolKind::CaseIterable);
 
+    // SWIFT_ENABLE_TENSORFLOW
+    // KeyPathIterable.AllKeyPaths
+    if (name.isSimpleName(ctx.Id_AllKeyPaths))
+      return getRequirement(KnownProtocolKind::KeyPathIterable);
+
     // SWIFT_ENABLE_TENSORFLOW
     // Parameterized.Parameters
     if (name.isSimpleName(ctx.Id_Parameters))

lib/Sema/DerivedConformances.h

@@ -164,6 +164,22 @@ class DerivedConformance {
   /// \returns the derived member, which will also be added to the type.
   ValueDecl *deriveDecodable(ValueDecl *requirement);
 
+  // SWIFT_ENABLE_TENSORFLOW
+  /// Determine if a KeyPathIterable requirement can be derived for a type.
+  ///
+  /// \returns True if the requirement can be derived.
+  static bool canDeriveKeyPathIterable(NominalTypeDecl *type);
+
+  /// Derive a KeyPathIterable requirement for a nominal type.
+  ///
+  /// \returns the derived member, which will also be added to the type.
+  ValueDecl *deriveKeyPathIterable(ValueDecl *requirement);
+
+  /// Derive a KeyPathIterable type witness for a nominal type.
+  ///
+  /// \returns the derived member, which will also be added to the type.
+  Type deriveKeyPathIterable(AssociatedTypeDecl *assocType);
+
   // SWIFT_ENABLE_TENSORFLOW
   /// Derive a Parameterized requirement for a nominal type.
   ///

lib/Sema/TypeCheckProtocol.cpp

@@ -5276,6 +5276,10 @@ ValueDecl *TypeChecker::deriveProtocolRequirement(DeclContext *DC,
   case KnownProtocolKind::Decodable:
     return derived.deriveDecodable(Requirement);
 
+  // SWIFT_ENABLE_TENSORFLOW
+  case KnownProtocolKind::KeyPathIterable:
+    return derived.deriveKeyPathIterable(Requirement);
+
   // SWIFT_ENABLE_TENSORFLOW
   case KnownProtocolKind::Parameterized:
     return derived.deriveParameterized(Requirement);
@@ -5308,6 +5312,8 @@ Type TypeChecker::deriveTypeWitness(DeclContext *DC,
   case KnownProtocolKind::CaseIterable:
     return derived.deriveCaseIterable(AssocType);
   // SWIFT_ENABLE_TENSORFLOW
+  case KnownProtocolKind::KeyPathIterable:
+    return derived.deriveKeyPathIterable(AssocType);
   case KnownProtocolKind::Parameterized:
     return derived.deriveParameterized(AssocType);
   case KnownProtocolKind::ParameterGroup:

stdlib/public/core/CMakeLists.txt

@@ -88,6 +88,8 @@ set(SWIFTLIB_ESSENTIAL
   IntegerTypes.swift.gyb
   Join.swift
   KeyPath.swift
+  # SWIFT_ENABLE_TENSORFLOW
+  KeyPathIterable.swift
   KeyValuePairs.swift
   LazyCollection.swift
   LazySequence.swift

stdlib/public/core/GroupInfo.json

@@ -140,7 +140,8 @@
     "MemoryLayout.swift"
   ],
   "KeyPaths": [
-    "KeyPath.swift"
+    "KeyPath.swift",
+    "KeyPathIterable.swift"
   ],
   "Reflection": [
     "Dump.swift",