RequirementMachine: Splitting concrete equivalence classes in protocol requirement signatures

slavapestov · slavapestov · commit 441fa1679a1b · 2022-03-22T15:02:06.000-04:00
diff --git a/lib/AST/RequirementMachine/RequirementMachineRequests.cpp b/lib/AST/RequirementMachine/RequirementMachineRequests.cpp
@@ -42,10 +42,11 @@ using namespace rewriting;
 /// concrete type requirements.
 static bool shouldSplitConcreteEquivalenceClass(
     Requirement req,
+    const ProtocolDecl *proto,
     const RequirementMachine *machine) {
   return (req.getKind() == RequirementKind::SameType &&
           req.getSecondType()->isTypeParameter() &&
-          machine->isConcreteType(req.getSecondType()));
+          machine->isConcreteType(req.getSecondType(), proto));
 }
 
 /// Returns true if this generic signature contains abstract same-type
@@ -54,9 +55,24 @@ static bool shouldSplitConcreteEquivalenceClass(
 /// requirements, and minimize the signature again.
 static bool shouldSplitConcreteEquivalenceClasses(
     ArrayRef<Requirement> requirements,
+    const ProtocolDecl *proto,
     const RequirementMachine *machine) {
   for (auto req : requirements) {
-    if (shouldSplitConcreteEquivalenceClass(req, machine))
+    if (shouldSplitConcreteEquivalenceClass(req, proto, machine))
+      return true;
+  }
+
+  return false;
+}
+
+/// Same as the above, but with the requirements of a protocol connected
+/// component.
+static bool shouldSplitConcreteEquivalenceClasses(
+    const llvm::DenseMap<const ProtocolDecl *, RequirementSignature> &protos,
+    const RequirementMachine *machine) {
+  for (const auto &pair : protos) {
+    if (shouldSplitConcreteEquivalenceClasses(pair.second.getRequirements(),
+                                              pair.first, machine))
       return true;
   }
 
@@ -71,14 +87,14 @@ static bool shouldSplitConcreteEquivalenceClasses(
 static void splitConcreteEquivalenceClasses(
     ASTContext &ctx,
     ArrayRef<Requirement> requirements,
+    const ProtocolDecl *proto,
     const RequirementMachine *machine,
     TypeArrayView<GenericTypeParamType> genericParams,
     SmallVectorImpl<StructuralRequirement> &splitRequirements,
     unsigned &attempt) {
   unsigned maxAttempts =
       ctx.LangOpts.RequirementMachineMaxSplitConcreteEquivClassAttempts;
 
-  ++attempt;
   if (attempt >= maxAttempts) {
     llvm::errs() << "Splitting concrete equivalence classes did not "
                  << "reach fixed point after " << attempt << " attempts.\n";
@@ -94,9 +110,9 @@ static void splitConcreteEquivalenceClasses(
   splitRequirements.clear();
 
   for (auto req : requirements) {
-    if (shouldSplitConcreteEquivalenceClass(req, machine)) {
+    if (shouldSplitConcreteEquivalenceClass(req, proto, machine)) {
       auto concreteType = machine->getConcreteType(
-          req.getSecondType(), genericParams);
+          req.getSecondType(), genericParams, proto);
 
       Requirement firstReq(RequirementKind::SameType,
                            req.getFirstType(), concreteType);
@@ -111,6 +127,25 @@ static void splitConcreteEquivalenceClasses(
   }
 }
 
+/// Same as the above, but with the requirements of a protocol connected
+/// component.
+static void splitConcreteEquivalenceClasses(
+    ASTContext &ctx,
+    const llvm::DenseMap<const ProtocolDecl *, RequirementSignature> &protos,
+    const RequirementMachine *machine,
+    llvm::DenseMap<const ProtocolDecl *,
+                   SmallVector<StructuralRequirement, 4>> &splitProtos,
+    unsigned &attempt) {
+  for (const auto &pair : protos) {
+    const auto *proto = pair.first;
+    auto genericParams = proto->getGenericSignature().getGenericParams();
+    splitConcreteEquivalenceClasses(ctx, pair.second.getRequirements(),
+                                    proto, machine, genericParams,
+                                    splitProtos[proto],
+                                    attempt);
+  }
+}
+
 /// Builds the requirement signatures for each protocol in this strongly
 /// connected component.
 llvm::DenseMap<const ProtocolDecl *, RequirementSignature>
@@ -195,87 +230,99 @@ RequirementSignatureRequestRQM::evaluate(Evaluator &evaluator,
       requirements.push_back({req, SourceLoc(), /*inferred=*/false});
   }
 
-  // Heap-allocate the requirement machine to save stack space.
-  std::unique_ptr<RequirementMachine> machine(new RequirementMachine(
-      ctx.getRewriteContext()));
+  unsigned attempt = 0;
+  for (;;) {
+    // Heap-allocate the requirement machine to save stack space.
+    std::unique_ptr<RequirementMachine> machine(new RequirementMachine(
+        ctx.getRewriteContext()));
 
-  auto status = machine->initWithProtocolWrittenRequirements(component, protos);
-  if (status.first != CompletionResult::Success) {
-    // All we can do at this point is diagnose and give each protocol an empty
-    // requirement signature.
-    for (const auto *otherProto : component) {
-      ctx.Diags.diagnose(otherProto->getLoc(),
-                         diag::requirement_machine_completion_failed,
-                         /*protocol=*/1,
-                         unsigned(status.first));
+    auto status = machine->initWithProtocolWrittenRequirements(component, protos);
+    if (status.first != CompletionResult::Success) {
+      // All we can do at this point is diagnose and give each protocol an empty
+      // requirement signature.
+      for (const auto *otherProto : component) {
+        ctx.Diags.diagnose(otherProto->getLoc(),
+                           diag::requirement_machine_completion_failed,
+                           /*protocol=*/1,
+                           unsigned(status.first));
+
+        auto rule = machine->getRuleAsStringForDiagnostics(status.second);
+        ctx.Diags.diagnose(otherProto->getLoc(),
+                           diag::requirement_machine_completion_rule,
+                           rule);
+
+        if (otherProto != proto) {
+          ctx.evaluator.cacheOutput(
+            RequirementSignatureRequestRQM{const_cast<ProtocolDecl *>(otherProto)},
+            RequirementSignature(GenericSignatureErrorFlags::CompletionFailed));
+        }
+      }
 
-      auto rule = machine->getRuleAsStringForDiagnostics(status.second);
-      ctx.Diags.diagnose(otherProto->getLoc(),
-                         diag::requirement_machine_completion_rule,
-                         rule);
+      return RequirementSignature(GenericSignatureErrorFlags::CompletionFailed);
+    }
 
-      if (otherProto != proto) {
-        ctx.evaluator.cacheOutput(
-          RequirementSignatureRequestRQM{const_cast<ProtocolDecl *>(otherProto)},
-          RequirementSignature(GenericSignatureErrorFlags::CompletionFailed));
+    auto minimalRequirements = machine->computeMinimalProtocolRequirements();
+
+    if (!machine->getErrors()) {
+      if (shouldSplitConcreteEquivalenceClasses(minimalRequirements, machine.get())) {
+        ++attempt;
+        splitConcreteEquivalenceClasses(ctx, minimalRequirements,
+                                        machine.get(), protos, attempt);
+        continue;
       }
     }
 
-    return RequirementSignature(GenericSignatureErrorFlags::CompletionFailed);
-  }
-
-  auto minimalRequirements = machine->computeMinimalProtocolRequirements();
+    bool debug = machine->getDebugOptions().contains(DebugFlags::Minimization);
 
-  bool debug = machine->getDebugOptions().contains(DebugFlags::Minimization);
+    // The requirement signature for the actual protocol that the result
+    // was kicked off with.
+    Optional<RequirementSignature> result;
 
-  // The requirement signature for the actual protocol that the result
-  // was kicked off with.
-  Optional<RequirementSignature> result;
+    if (debug) {
+      llvm::dbgs() << "\nRequirement signatures:\n";
+    }
 
-  if (debug) {
-    llvm::dbgs() << "\nRequirement signatures:\n";
-  }
+    for (const auto &pair : minimalRequirements) {
+      auto *otherProto = pair.first;
+      const auto &reqs = pair.second;
 
-  for (const auto &pair : minimalRequirements) {
-    auto *otherProto = pair.first;
-    const auto &reqs = pair.second;
+      // Dump the result if requested.
+      if (debug) {
+        llvm::dbgs() << "- Protocol " << otherProto->getName() << ": ";
 
-    // Dump the result if requested.
-    if (debug) {
-      llvm::dbgs() << "- Protocol " << otherProto->getName() << ": ";
+        auto sig = GenericSignature::get(
+            otherProto->getGenericSignature().getGenericParams(),
+            reqs.getRequirements());
 
-      auto sig = GenericSignature::get(
-          otherProto->getGenericSignature().getGenericParams(),
-          reqs.getRequirements());
+        PrintOptions opts;
+        opts.ProtocolQualifiedDependentMemberTypes = true;
+        sig.print(llvm::dbgs(), opts);
+        llvm::dbgs() << "\n";
+      }
 
-      PrintOptions opts;
-      opts.ProtocolQualifiedDependentMemberTypes = true;
-      sig.print(llvm::dbgs(), opts);
-      llvm::dbgs() << "\n";
+      // Don't call setRequirementSignature() on the original proto; the
+      // request evaluator will do it for us.
+      if (otherProto == proto)
+        result = reqs;
+      else {
+        auto temp = reqs;
+        ctx.evaluator.cacheOutput(
+          RequirementSignatureRequestRQM{const_cast<ProtocolDecl *>(otherProto)},
+          std::move(temp));
+      }
     }
 
-    // Don't call setRequirementSignature() on the original proto; the
-    // request evaluator will do it for us.
-    if (otherProto == proto)
-      result = reqs;
-    else {
-      auto temp = reqs;
-      ctx.evaluator.cacheOutput(
-        RequirementSignatureRequestRQM{const_cast<ProtocolDecl *>(otherProto)},
-        std::move(temp));
+    if (ctx.LangOpts.RequirementMachineProtocolSignatures ==
+        RequirementMachineMode::Enabled) {
+      SmallVector<RequirementError, 4> errors;
+      machine->System.computeRedundantRequirementDiagnostics(errors);
+      diagnoseRequirementErrors(ctx, errors,
+                                /*allowConcreteGenericParams=*/false);
     }
-  }
 
-  if (ctx.LangOpts.RequirementMachineProtocolSignatures ==
-      RequirementMachineMode::Enabled) {
-    SmallVector<RequirementError, 4> errors;
-    machine->System.computeRedundantRequirementDiagnostics(errors);
-    diagnoseRequirementErrors(ctx, errors,
-                              /*allowConcreteGenericParams=*/false);
+    // Return the result for the specific protocol this request was kicked off on.
+    return *result;
   }
-
-  // Return the result for the specific protocol this request was kicked off on.
-  return *result;
 }
 
 /// Builds the top-level generic signature requirements for this rewrite system.
@@ -479,9 +526,12 @@ AbstractGenericSignatureRequestRQM::evaluate(
 
     if (!errorFlags) {
       if (shouldSplitConcreteEquivalenceClasses(result.getRequirements(),
+                                                /*proto=*/nullptr,
                                                 machine.get())) {
+        ++attempt;
         splitConcreteEquivalenceClasses(ctx, result.getRequirements(),
-                                        machine.get(), result.getGenericParams(),
+                                        /*proto=*/nullptr, machine.get(),
+                                        result.getGenericParams(),
                                         requirements, attempt);
         continue;
       }
@@ -646,9 +696,12 @@ InferredGenericSignatureRequestRQM::evaluate(
     if (!errorFlags) {
       // Check if we need to rebuild the signature.
       if (shouldSplitConcreteEquivalenceClasses(result.getRequirements(),
+                                                /*proto=*/nullptr,
                                                 machine.get())) {
+        ++attempt;
         splitConcreteEquivalenceClasses(ctx, result.getRequirements(),
-                                        machine.get(), result.getGenericParams(),
+                                        /*proto=*/nullptr, machine.get(),
+                                        result.getGenericParams(),
                                         requirements, attempt);
         continue;
       }
diff --git a/test/Generics/split_concrete_equivalence_class_in_protocol.swift b/test/Generics/split_concrete_equivalence_class_in_protocol.swift
@@ -0,0 +1,51 @@
+// RUN: not %target-swift-frontend -typecheck %s -debug-generic-signatures 2>&1 | %FileCheck %s
+
+// CHECK-LABEL: .P1@
+// CHECK-NEXT: Requirement signature: <Self where Self.[P1]T : P1, Self.[P1]U == Int, Self.[P1]V == Int>
+protocol P1 {
+  associatedtype T : P1
+  associatedtype U where U == Int
+  associatedtype V where V == T.U
+}
+
+struct G<X> {}
+
+// CHECK-LABEL: .P2@
+// CHECK-NEXT: Requirement signature: <Self where Self.[P2]T : P2, Self.[P2]U == G<Self.[P2]X>, Self.[P2]V == G<Self.[P2]T.[P2]X>>
+protocol P2 {
+  associatedtype T : P2
+  associatedtype U where U == G<X>
+  associatedtype V where V == T.U
+  associatedtype X
+}
+
+// CHECK-LABEL: .P3@
+// CHECK-NEXT: Requirement signature: <Self where Self.[P3]T : P3, Self.[P3]U == G<Self.[P3]X>, Self.[P3]V == G<Self.[P3]X>, Self.[P3]X == Self.[P3]T.[P3]X>
+protocol P3 {
+  associatedtype T : P3
+  associatedtype U where U == G<X>
+  associatedtype V where V == T.U, V == G<X>
+  associatedtype X
+}
+
+// CHECK-LABEL: .P4@
+// CHECK-NEXT: Requirement signature: <Self where Self.[P4]T : P4, Self.[P4]U == G<Self.[P4]X>, Self.[P4]V == G<Self.[P4]X>, Self.[P4]X == Self.[P4]T.[P4]X>
+protocol P4 {
+  associatedtype T : P4
+  associatedtype U where U == G<X>
+  associatedtype V where V == T.U
+  associatedtype X where X == T.X
+}
+
+// We don't split concrete equivalence classes if the signature had an error,
+// but we also shouldn't crash in verify() because of an unsplit concrete
+// equivalence class.
+
+// CHECK-LABEL: .P4Bad@
+// CHECK-NEXT: Requirement signature: <Self where Self.[P4Bad]T : P4Bad, Self.[P4Bad]U == G<Self.[P4Bad]X>, Self.[P4Bad]V == Self.[P4Bad]T.[P4Bad]U>
+protocol P4Bad {
+  associatedtype T : P4Bad
+  associatedtype U where U == G<X>
+  associatedtype V where V == T.U
+  associatedtype X where X == U.X
+}
diff --git a/test/decl/protocol/existential_member_accesses_self_assoctype.swift b/test/decl/protocol/existential_member_accesses_self_assoctype.swift
@@ -1,6 +1,4 @@
-// RUN: %target-typecheck-verify-swift -disable-availability-checking -requirement-machine-protocol-signatures=off -requirement-machine-abstract-signatures=on
-
-// TODO: Get this to pass with  -requirement-machine-protocol-signatures=on.
+// RUN: %target-typecheck-verify-swift -disable-availability-checking -requirement-machine-abstract-signatures=on
 
 //===----------------------------------------------------------------------===//
 // Use of protocols with Self or associated type requirements
diff --git a/test/decl/protocol/req/associated_type_inference_fixed_type.swift b/test/decl/protocol/req/associated_type_inference_fixed_type.swift
@@ -1,7 +1,5 @@
-// RUN: %target-typecheck-verify-swift -requirement-machine-protocol-signatures=off
-// RUN: not %target-swift-frontend -typecheck -dump-type-witness-systems %s -requirement-machine-protocol-signatures=off 2>&1 | %FileCheck %s
-
-// TODO: Get this to pass with  -requirement-machine-protocol-signatures=on.
+// RUN: %target-typecheck-verify-swift
+// RUN: not %target-swift-frontend -typecheck -dump-type-witness-systems %s 2>&1 | %FileCheck %s
 
 protocol P1 where A == Never {
   associatedtype A
@@ -572,7 +570,7 @@ protocol P32e where A == B {
 }
 protocol Q32: P32e, P32a, P32b, P32c, P32d {}
 // expected-error@-1 {{'Self.B' cannot be equal to both 'Never' and 'Int'}}
-// expected-error@-2 {{'Self.B' cannot be equal to both 'Void' and 'Int'}}
+// expected-error@-2 {{'Self.B' cannot be equal to both '()' and 'Int'}}
 // expected-error@-3 {{'Self.A' cannot be equal to both 'Bool' and 'Int'}}
 // expected-note@-4 3 {{same-type constraint 'Self.A' == 'Int' implied here}}
 do {
