RequirementMachine: Same-type requirements imply same-shape requirements

slavapestov · slavapestov · commit 4a494199fc0a · 2023-07-03T15:18:04.000-04:00
We want `T.A == U.B` to imply `shape(T) == shape(U)` if T (and thus U)
is a parameter pack.

To do this, we introduce some new rewrite rules:

1) For each associated type symbol `[P:A]`, a rule `([P:A].[shape] =&gt; [P:A])`.
2) For each non-pack generic parameter `τ_d_i`, a rule `τ_d_i.[shape] =&gt; [shape]`.

Now consider a rewrite rule `(τ_d_i.[P:A] =&gt; τ_D_I.[Q:B])`. The left-hand
side overlaps with the rule `([P:A].[shape] =&gt; [shape])` on the term
`τ_d_i.[P:A].[shape]`. Resolving the overlap gives us a new rule

    t_d_i.[shape] =&gt; T_D_I.[shape]

If T is a term corresponding to some type parameter, we say that `T.[shape]` is
a shape term. If `T'.[shape]` is a reduced term, we say that T' is the reduced
shape of T.

Recall that shape requirements are represented as rules of the form:

    τ_d_i.[shape] =&gt; τ_D_I.[shape]

Now, the rules of the first kind reduce our shape term `T.[shape]` to
`τ_d_i.[shape]`, where `τ_d_i` is the root generic parameter of T.

If `τ_d_i` is not a pack, a rule of the second kind reduces it to `[shape]`,
so the reduced shape of a non-pack parameter T is the empty term.

Otherwise, if `τ_d_i` is a pack, `τ_d_i.[shape]` might reduce to `τ_D_I.[shape]`
via a shape requirement. In this case, `τ_D_I` is the reduced shape of T.

Fixes rdar://problem/101813873.
diff --git a/lib/AST/RequirementMachine/RequirementMachine.cpp b/lib/AST/RequirementMachine/RequirementMachine.cpp
@@ -307,7 +307,7 @@ RequirementMachine::initWithProtocolSignatureRequirements(
 ///
 /// Returns failure if completion fails within the configured number of steps.
 std::pair<CompletionResult, unsigned>
-RequirementMachine::initWithGenericSignature(CanGenericSignature sig) {
+RequirementMachine::initWithGenericSignature(GenericSignature sig) {
   Sig = sig;
   Params.append(sig.getGenericParams().begin(),
                 sig.getGenericParams().end());
@@ -323,7 +323,8 @@ RequirementMachine::initWithGenericSignature(CanGenericSignature sig) {
   // Collect the top-level requirements, and all transitively-referenced
   // protocol requirement signatures.
   RuleBuilder builder(Context, System.getReferencedProtocols());
-  builder.initWithGenericSignatureRequirements(sig.getRequirements());
+  builder.initWithGenericSignature(sig.getGenericParams(),
+                                   sig.getRequirements());
 
   // Add the initial set of rewrite rules to the rewrite system.
   System.initialize(/*recordLoops=*/false,
@@ -425,7 +426,7 @@ RequirementMachine::initWithWrittenRequirements(
   // Collect the top-level requirements, and all transitively-referenced
   // protocol requirement signatures.
   RuleBuilder builder(Context, System.getReferencedProtocols());
-  builder.initWithWrittenRequirements(requirements);
+  builder.initWithWrittenRequirements(genericParams, requirements);
 
   // Add the initial set of rewrite rules to the rewrite system.
   System.initialize(/*recordLoops=*/true,
diff --git a/lib/AST/RequirementMachine/RequirementMachine.h b/lib/AST/RequirementMachine/RequirementMachine.h
@@ -54,7 +54,7 @@ class RequirementMachine final {
   friend class swift::AbstractGenericSignatureRequest;
   friend class swift::InferredGenericSignatureRequest;
 
-  CanGenericSignature Sig;
+  GenericSignature Sig;
   SmallVector<GenericTypeParamType *, 2> Params;
 
   RewriteContext &Context;
@@ -95,7 +95,7 @@ class RequirementMachine final {
       ArrayRef<const ProtocolDecl *> protos);
 
   std::pair<CompletionResult, unsigned>
-  initWithGenericSignature(CanGenericSignature sig);
+  initWithGenericSignature(GenericSignature sig);
 
   std::pair<CompletionResult, unsigned>
   initWithProtocolWrittenRequirements(
diff --git a/lib/AST/RequirementMachine/RewriteSystem.cpp b/lib/AST/RequirementMachine/RewriteSystem.cpp
@@ -504,8 +504,16 @@ void RewriteSystem::recordRewriteLoop(MutableTerm basepoint,
     return;
 
   // Ignore the rewrite loop if it is not part of our minimization domain.
-  if (!isInMinimizationDomain(basepoint.getRootProtocol()))
+  //
+  // Completion might record a rewrite loop where the basepoint is just
+  // the term [shape]. In this case though, we know it's in our domain,
+  // since completion only checks local rules for overlap. Other callers
+  // of recordRewriteLoop() always pass in a valid basepoint, so we
+  // check.
+  if (basepoint[0].getKind() != Symbol::Kind::Shape &&
+      !isInMinimizationDomain(basepoint.getRootProtocol())) {
     return;
+  }
 
   Loops.push_back(loop);
 }
@@ -555,11 +563,6 @@ void RewriteSystem::verifyRewriteRules(ValidityPolicy policy) const {
         ASSERT_RULE(symbol.getKind() != Symbol::Kind::Shape);
       }
 
-      // A shape symbol must follow a generic param symbol
-      if (symbol.getKind() == Symbol::Kind::Shape) {
-        ASSERT_RULE(index > 0 && lhs[index - 1].getKind() == Symbol::Kind::GenericParam);
-      }
-
       if (!rule.isLHSSimplified() &&
           index != lhs.size() - 1) {
         ASSERT_RULE(symbol.getKind() != Symbol::Kind::ConcreteConformance);
@@ -602,15 +605,10 @@ void RewriteSystem::verifyRewriteRules(ValidityPolicy policy) const {
       ASSERT_RULE(symbol.getKind() != Symbol::Kind::Superclass);
       ASSERT_RULE(symbol.getKind() != Symbol::Kind::ConcreteType);
 
-      if (index != lhs.size() - 1) {
+      if (index != rhs.size() - 1) {
         ASSERT_RULE(symbol.getKind() != Symbol::Kind::Shape);
       }
 
-      // A shape symbol must follow a generic param symbol
-      if (symbol.getKind() == Symbol::Kind::Shape) {
-        ASSERT_RULE(index > 0 && rhs[index - 1].getKind() == Symbol::Kind::GenericParam);
-      }
-
       // Completion can introduce a rule of the form
       //
       // (T.[P] => T.[concrete: C : P])
@@ -635,10 +633,15 @@ void RewriteSystem::verifyRewriteRules(ValidityPolicy policy) const {
       }
     }
 
-    auto lhsDomain = lhs.getRootProtocol();
-    auto rhsDomain = rhs.getRootProtocol();
-
-    ASSERT_RULE(lhsDomain == rhsDomain);
+    if (rhs.size() == 1 && rhs[0].getKind() == Symbol::Kind::Shape) {
+      // We can have a rule like T.[shape] => [shape].
+      ASSERT_RULE(lhs.back().getKind() == Symbol::Kind::Shape);
+    } else {
+      // Otherwise, LHS and RHS must have the same domain.
+      auto lhsDomain = lhs.getRootProtocol();
+      auto rhsDomain = rhs.getRootProtocol();
+      ASSERT_RULE(lhsDomain == rhsDomain);
+    }
   }
 
 #undef ASSERT_RULE
diff --git a/lib/AST/RequirementMachine/RuleBuilder.cpp b/lib/AST/RequirementMachine/RuleBuilder.cpp
@@ -34,7 +34,8 @@ using namespace rewriting;
 
 /// For building a rewrite system for a generic signature from canonical
 /// requirements.
-void RuleBuilder::initWithGenericSignatureRequirements(
+void RuleBuilder::initWithGenericSignature(
+    ArrayRef<GenericTypeParamType *> genericParams,
     ArrayRef<Requirement> requirements) {
   assert(!Initialized);
   Initialized = 1;
@@ -47,6 +48,7 @@ void RuleBuilder::initWithGenericSignatureRequirements(
   }
 
   collectRulesFromReferencedProtocols();
+  collectPackShapeRules(genericParams);
 
   // Add rewrite rules for all top-level requirements.
   for (const auto &req : requirements)
@@ -56,6 +58,7 @@ void RuleBuilder::initWithGenericSignatureRequirements(
 /// For building a rewrite system for a generic signature from user-written
 /// requirements.
 void RuleBuilder::initWithWrittenRequirements(
+    ArrayRef<GenericTypeParamType *> genericParams,
     ArrayRef<StructuralRequirement> requirements) {
   assert(!Initialized);
   Initialized = 1;
@@ -68,6 +71,7 @@ void RuleBuilder::initWithWrittenRequirements(
   }
 
   collectRulesFromReferencedProtocols();
+  collectPackShapeRules(genericParams);
 
   // Add rewrite rules for all top-level requirements.
   for (const auto &req : requirements)
@@ -488,3 +492,77 @@ void RuleBuilder::collectRulesFromReferencedProtocols() {
                          localRules.end());
   }
 }
+
+void RuleBuilder::collectPackShapeRules(ArrayRef<GenericTypeParamType *> genericParams) {
+  if (Dump) {
+    llvm::dbgs() << "adding shape rules\n";
+  }
+
+  if (!llvm::any_of(genericParams,
+                    [](GenericTypeParamType *t) {
+                      return t->isParameterPack();
+                    })) {
+    return;
+  }
+
+  // Each non-pack generic parameter is part of the "scalar shape class", represented
+  // by the empty term.
+  for (auto *genericParam : genericParams) {
+    if (genericParam->isParameterPack())
+      continue;
+
+    // Add the rule (τ_d_i.[shape] => [shape]).
+    MutableTerm lhs;
+    lhs.add(Symbol::forGenericParam(
+        cast<GenericTypeParamType>(genericParam->getCanonicalType()), Context));
+    lhs.add(Symbol::forShape(Context));
+
+    MutableTerm rhs;
+    rhs.add(Symbol::forShape(Context));
+
+    PermanentRules.emplace_back(lhs, rhs);
+  }
+
+  // A member type T.[P:A] is part of the same shape class as its base type T.
+  llvm::DenseSet<Symbol> visited;
+
+  auto addMemberShapeRule = [&](const ProtocolDecl *proto, AssociatedTypeDecl *assocType) {
+    auto symbol = Symbol::forAssociatedType(proto, assocType->getName(), Context);
+    if (!visited.insert(symbol).second)
+      return;
+
+    // Add the rule ([P:A].[shape] => [shape]).
+    MutableTerm lhs;
+    lhs.add(symbol);
+    lhs.add(Symbol::forShape(Context));
+
+    MutableTerm rhs;
+    rhs.add(Symbol::forShape(Context));
+
+    // Consider it an imported rule, since it is not part of our minimization
+    // domain. It would be more logical if we added these in the protocol component
+    // machine for this protocol, but instead we add them in the "leaf" generic
+    // signature machine. This avoids polluting machines that do not involve
+    // parameter packs with these extra rules, which would otherwise just slow
+    // things down.
+    Rule rule(Term::get(lhs, Context), Term::get(rhs, Context));
+    rule.markPermanent();
+    ImportedRules.push_back(rule);
+  };
+
+  for (auto *proto : ProtocolsToImport) {
+    if (Dump) {
+      llvm::dbgs() << "adding member shape rules for protocol " << proto->getName() << "\n";
+    }
+
+    for (auto *assocType : proto->getAssociatedTypeMembers()) {
+      addMemberShapeRule(proto, assocType);
+    }
+
+    for (auto *inheritedProto : Context.getInheritedProtocols(proto)) {
+      for (auto *assocType : inheritedProto->getAssociatedTypeMembers()) {
+        addMemberShapeRule(proto, assocType);
+      }
+    }
+  }
+}
diff --git a/lib/AST/RequirementMachine/RuleBuilder.h b/lib/AST/RequirementMachine/RuleBuilder.h
@@ -95,8 +95,10 @@ struct RuleBuilder {
     Initialized = 0;
   }
 
-  void initWithGenericSignatureRequirements(ArrayRef<Requirement> requirements);
-  void initWithWrittenRequirements(ArrayRef<StructuralRequirement> requirements);
+  void initWithGenericSignature(ArrayRef<GenericTypeParamType *> genericParams,
+                                ArrayRef<Requirement> requirements);
+  void initWithWrittenRequirements(ArrayRef<GenericTypeParamType *> genericParams,
+                                   ArrayRef<StructuralRequirement> requirements);
   void initWithProtocolSignatureRequirements(ArrayRef<const ProtocolDecl *> proto);
   void initWithProtocolWrittenRequirements(
       ArrayRef<const ProtocolDecl *> component,
@@ -106,6 +108,7 @@ struct RuleBuilder {
                                        ArrayRef<Term> substitutions);
   void addReferencedProtocol(const ProtocolDecl *proto);
   void collectRulesFromReferencedProtocols();
+  void collectPackShapeRules(ArrayRef<GenericTypeParamType *> genericParams);
 
 private:
   void addPermanentProtocolRules(const ProtocolDecl *proto);
diff --git a/test/Generics/pack-shape-requirements.swift b/test/Generics/pack-shape-requirements.swift
@@ -1,7 +1,7 @@
 // RUN: %target-swift-frontend -typecheck %s -debug-generic-signatures -disable-availability-checking 2>&1 | %FileCheck %s
 
 protocol P {
-  associatedtype A
+  associatedtype A: P
 }
 
 // CHECK-LABEL: inferSameShape(ts:us:)
@@ -65,3 +65,21 @@ struct Ts<each T> {
 func expandedParameters<each T, each Result>(_ t: repeat each T, transform: repeat (each T) -> each Result) -> (repeat each Result) {
   fatalError()
 }
+
+
+//////
+///
+/// Same-type requirements should imply same-shape requirements.
+///
+//////
+
+// CHECK-LABEL: sameType1
+// CHECK-NEXT: Generic signature: <each T, each U where repeat each T : P, repeat each U : P, repeat (each T).[P]A == (each U).[P]A>
+func sameType1<each T, each U>(_: repeat (each T, each U)) where repeat each T: P, repeat each U: P, repeat each T.A == each U.A {}
+
+// Make sure inherited associated types are handled
+protocol Q: P where A: Q {}
+
+// CHECK-LABEL: sameType2
+// CHECK-NEXT: Generic signature: <each T, each U where repeat each T : Q, repeat each U : Q, repeat (each T).[P]A.[P]A == (each U).[P]A.[P]A>
+func sameType2<each T, each U>(_: repeat (each T, each U)) where repeat each T: Q, repeat each U: Q, repeat each T.A.A == each U.A.A {}
