AutoDiff: Use getLocalProtocols() instead of getInheritedProtocols()

Author: slavapestov

lib/Sema/DerivedConformanceDifferentiable.cpp

@@ -337,51 +337,6 @@ static ValueDecl *deriveDifferentiable_move(DerivedConformance &derived) {
       C.TheEmptyTupleType, {deriveBodyDifferentiable_move, nullptr});
 }
 
-/// Pushes all the protocols inherited, directly or transitively, by `decl` to `protos`.
-///
-/// Precondition: `decl` is a nominal type decl or an extension decl.
-void getInheritedProtocols(Decl *decl, SmallPtrSetImpl<ProtocolDecl *> &protos) {
-  ArrayRef<TypeLoc> inheritedTypeLocs;
-  if (auto *nominalDecl = dyn_cast<NominalTypeDecl>(decl))
-    inheritedTypeLocs = nominalDecl->getInherited();
-  else if (auto *extDecl = dyn_cast<ExtensionDecl>(decl))
-    inheritedTypeLocs = extDecl->getInherited();
-  else
-    llvm_unreachable("conformance is not a nominal or an extension");
-
-  std::function<void(Type)> handleInheritedType;
-
-  auto handleProto = [&](ProtocolType *proto) -> void {
-    proto->getDecl()->walkInheritedProtocols([&](ProtocolDecl *p) -> TypeWalker::Action {
-      protos.insert(p);
-      return TypeWalker::Action::Continue;
-    });
-  };
-
-  auto handleProtoComp = [&](ProtocolCompositionType *comp) -> void {
-    for (auto ty : comp->getMembers())
-      handleInheritedType(ty);
-  };
-
-  handleInheritedType = [&](Type ty) -> void {
-    if (auto *proto = ty->getAs<ProtocolType>())
-      handleProto(proto);
-    else if (auto *comp = ty->getAs<ProtocolCompositionType>())
-      handleProtoComp(comp);
-  };
-
-  for (auto loc : inheritedTypeLocs) {
-    if (loc.getTypeRepr())
-      handleInheritedType(
-          TypeResolution::forStructural(cast<DeclContext>(decl), None,
-                                        /*unboundTyOpener*/ nullptr,
-                                        /*placeholderHandler*/ nullptr)
-              .resolveType(loc.getTypeRepr()));
-    else
-      handleInheritedType(loc.getType());
-  }
-}
-
 /// Return associated `TangentVector` struct for a nominal type, if it exists.
 /// If not, synthesize the struct.
 static StructDecl *
@@ -409,12 +364,14 @@ getOrSynthesizeTangentVectorStruct(DerivedConformance &derived, Identifier id) {
   // Note that, for example, this will always find `AdditiveArithmetic` and `Differentiable` because
   // the `Differentiable` protocol itself requires that its `TangentVector` conforms to
   // `AdditiveArithmetic` and `Differentiable`.
-  llvm::SmallPtrSet<ProtocolDecl *, 4> tvDesiredProtos;
-  llvm::SmallPtrSet<ProtocolDecl *, 4> conformanceInheritedProtos;
-  getInheritedProtocols(derived.ConformanceDecl, conformanceInheritedProtos);
+  llvm::SmallSetVector<ProtocolDecl *, 4> tvDesiredProtos;
+
   auto *diffableProto = C.getProtocol(KnownProtocolKind::Differentiable);
   auto *tvAssocType = diffableProto->getAssociatedType(C.Id_TangentVector);
-  for (auto proto : conformanceInheritedProtos) {
+
+  auto localProtos = cast<IterableDeclContext>(derived.ConformanceDecl)
+      ->getLocalProtocols();
+  for (auto proto : localProtos) {
     for (auto req : proto->getRequirementSignature()) {
       if (req.getKind() != RequirementKind::Conformance)
         continue;