RequirementMachine: Implement GenericSignature::getRequiredProtocols() query

Author: slavapestov

include/swift/AST/RequirementMachine.h

@@ -13,6 +13,8 @@
 #ifndef SWIFT_REQUIREMENTMACHINE_H
 #define SWIFT_REQUIREMENTMACHINE_H
 
+#include "swift/AST/GenericSignature.h"
+
 namespace llvm {
 class raw_ostream;
 }
@@ -21,9 +23,7 @@ namespace swift {
 
 class ASTContext;
 class AssociatedTypeDecl;
-class CanGenericSignature;
 class CanType;
-class GenericSignature;
 class LayoutConstraint;
 class ProtocolDecl;
 class Requirement;
@@ -58,6 +58,7 @@ class RequirementMachine final {
   bool requiresClass(Type depType) const;
   LayoutConstraint getLayoutConstraint(Type depType) const;
   bool requiresProtocol(Type depType, const ProtocolDecl *proto) const;
+  GenericSignature::RequiredProtocols getRequiredProtocols(Type depType) const;
   bool isConcreteType(Type depType) const;
   bool areSameTypeParameterInContext(Type depType1, Type depType2) const;
 

lib/AST/GenericSignature.cpp

@@ -353,26 +353,62 @@ GenericSignature::RequiredProtocols
 GenericSignatureImpl::getRequiredProtocols(Type type) const {
   assert(type->isTypeParameter() && "Expected a type parameter");
 
-  auto &builder = *getGenericSignatureBuilder();
-  auto equivClass =
-    builder.resolveEquivalenceClass(
-                                  type,
-                                  ArchetypeResolutionKind::CompleteWellFormed);
-  if (!equivClass) return { };
+  auto computeViaGSB = [&]() -> GenericSignature::RequiredProtocols {
+    auto &builder = *getGenericSignatureBuilder();
+    auto equivClass =
+      builder.resolveEquivalenceClass(
+                                    type,
+                                    ArchetypeResolutionKind::CompleteWellFormed);
+    if (!equivClass) return { };
 
-  // If this type parameter was mapped to a concrete type, then there
-  // are no requirements.
-  if (equivClass->concreteType) return { };
+    // If this type parameter was mapped to a concrete type, then there
+    // are no requirements.
+    if (equivClass->concreteType) return { };
 
-  // Retrieve the protocols to which this type conforms.
-  GenericSignature::RequiredProtocols result;
-  for (const auto &conforms : equivClass->conformsTo)
-    result.push_back(conforms.first);
+    // Retrieve the protocols to which this type conforms.
+    GenericSignature::RequiredProtocols result;
+    for (const auto &conforms : equivClass->conformsTo)
+      result.push_back(conforms.first);
 
-  // Canonicalize the resulting set of protocols.
-  ProtocolType::canonicalizeProtocols(result);
+    // Canonicalize the resulting set of protocols.
+    ProtocolType::canonicalizeProtocols(result);
 
-  return result;
+    return result;
+  };
+
+  auto computeViaRQM = [&]() {
+    auto *machine = getRequirementMachine();
+    return machine->getRequiredProtocols(type);
+  };
+
+  auto &ctx = getASTContext();
+  if (ctx.LangOpts.EnableRequirementMachine) {
+    auto rqmResult = computeViaRQM();
+
+#ifndef NDEBUG
+    auto gsbResult = computeViaGSB();
+
+    if (gsbResult != rqmResult) {
+      llvm::errs() << "RequirementMachine::getRequiredProtocols() is broken\n";
+      llvm::errs() << "Generic signature: " << GenericSignature(this) << "\n";
+      llvm::errs() << "Dependent type: "; type.dump(llvm::errs());
+      llvm::errs() << "GenericSignatureBuilder says: ";
+      for (auto *otherProto : gsbResult)
+        llvm::errs() << " " << otherProto->getName();
+      llvm::errs() << "\n";
+      llvm::errs() << "RequirementMachine says: ";
+      for (auto *otherProto : rqmResult)
+        llvm::errs() << " " << otherProto->getName();
+      llvm::errs() << "\n";
+      getRequirementMachine()->dump(llvm::errs());
+      abort();
+    }
+#endif
+
+    return rqmResult;
+  } else {
+    return computeViaGSB();
+  }
 }
 
 bool GenericSignatureImpl::requiresProtocol(Type type,

lib/AST/RequirementMachine/RequirementMachine.cpp

@@ -390,6 +390,29 @@ bool RequirementMachine::requiresProtocol(Type depType,
   return false;
 }
 
+GenericSignature::RequiredProtocols
+RequirementMachine::getRequiredProtocols(Type depType) const {
+  auto term = Impl->Context.getMutableTermForType(depType->getCanonicalType(),
+                                                  /*proto=*/nullptr);
+  Impl->System.simplify(term);
+
+  auto *equivClass = Impl->Map.lookUpEquivalenceClass(term);
+  if (!equivClass)
+    return { };
+
+  if (equivClass->isConcreteType())
+    return { };
+
+  GenericSignature::RequiredProtocols result;
+  for (auto *otherProto : equivClass->getConformsTo()) {
+    result.push_back(const_cast<ProtocolDecl *>(otherProto));
+  }
+
+  ProtocolType::canonicalizeProtocols(result);
+
+  return result;
+}
+
 bool RequirementMachine::isConcreteType(Type depType) const {
   auto term = Impl->Context.getMutableTermForType(depType->getCanonicalType(),
                                                   /*proto=*/nullptr);