Merge pull request #38284 from slavapestov/requirement-machine-canonical-types

RequirementMachine: Implement GenericSignature::getCanonicalTypeInContext()

Author: slavapestov

include/swift/AST/RequirementMachine.h

@@ -24,6 +24,7 @@ namespace swift {
 class ASTContext;
 class AssociatedTypeDecl;
 class CanType;
+class GenericTypeParamType;
 class LayoutConstraint;
 class ProtocolDecl;
 class Requirement;
@@ -49,7 +50,7 @@ class RequirementMachine final {
   void addGenericSignature(CanGenericSignature sig);
 
   bool isComplete() const;
-  void computeCompletion(CanGenericSignature sig);
+  void computeCompletion();
 
 public:
   ~RequirementMachine();
@@ -61,6 +62,8 @@ class RequirementMachine final {
   GenericSignature::RequiredProtocols getRequiredProtocols(Type depType) const;
   bool isConcreteType(Type depType) const;
   bool areSameTypeParameterInContext(Type depType1, Type depType2) const;
+  Type getCanonicalTypeInContext(Type type,
+                      TypeArrayView<GenericTypeParamType> genericParams) const;
 
   void dump(llvm::raw_ostream &out) const;
 };

lib/AST/GenericSignature.cpp

@@ -743,8 +743,41 @@ CanType GenericSignatureImpl::getCanonicalTypeInContext(Type type) const {
   if (!type->hasTypeParameter())
     return CanType(type);
 
-  auto &builder = *getGenericSignatureBuilder();
-  return builder.getCanonicalTypeInContext(type, { })->getCanonicalType();
+  auto computeViaGSB = [&]() {
+    auto &builder = *getGenericSignatureBuilder();
+    return builder.getCanonicalTypeInContext(type, { })->getCanonicalType();
+  };
+
+  auto computeViaRQM = [&]() {
+    auto *machine = getRequirementMachine();
+    return machine->getCanonicalTypeInContext(type, { })->getCanonicalType();
+  };
+
+  auto &ctx = getASTContext();
+  if (ctx.LangOpts.EnableRequirementMachine) {
+    auto rqmResult = computeViaRQM();
+
+#ifndef NDEBUG
+    auto gsbResult = computeViaGSB();
+
+    if (gsbResult != rqmResult) {
+      llvm::errs() << "RequirementMachine::getCanonicalTypeInContext() is broken\n";
+      llvm::errs() << "Generic signature: " << GenericSignature(this) << "\n";
+      llvm::errs() << "Dependent type: "; type.dump(llvm::errs());
+      llvm::errs() << "GenericSignatureBuilder says: " << gsbResult << "\n";
+      gsbResult.dump(llvm::errs());
+      llvm::errs() << "RequirementMachine says: " << rqmResult << "\n";
+      rqmResult.dump(llvm::errs());
+      llvm::errs() << "\n";
+      getRequirementMachine()->dump(llvm::errs());
+      abort();
+    }
+#endif
+
+    return rqmResult;
+  } else {
+    return computeViaGSB();
+  }
 }
 
 ArrayRef<CanTypeWrapper<GenericTypeParamType>>

lib/AST/RequirementMachine/EquivalenceClassMap.cpp

@@ -50,6 +50,8 @@
 
 #include "swift/AST/Decl.h"
 #include "swift/AST/LayoutConstraint.h"
+#include "swift/AST/Module.h"
+#include "swift/AST/ProtocolConformance.h"
 #include "swift/AST/TypeMatcher.h"
 #include "swift/AST/Types.h"
 #include "llvm/Support/raw_ostream.h"
@@ -102,6 +104,45 @@ static unsigned getGenericParamIndex(Type type) {
   return paramTy->getIndex();
 }
 
+/// Reverses the transformation performed by
+/// RewriteSystemBuilder::getConcreteSubstitutionSchema().
+static Type getTypeFromSubstitutionSchema(Type schema,
+                                          ArrayRef<Term> substitutions,
+                              TypeArrayView<GenericTypeParamType> genericParams,
+                                          const ProtocolGraph &protos,
+                                          RewriteContext &ctx) {
+  assert(!schema->isTypeParameter() && "Must have a concrete type here");
+
+  if (!schema->hasTypeParameter())
+    return schema;
+
+  return schema.transformRec([&](Type t) -> Optional<Type> {
+    if (t->is<GenericTypeParamType>()) {
+      auto index = getGenericParamIndex(t);
+
+      return ctx.getTypeForTerm(substitutions[index],
+                                genericParams, protos);
+    }
+
+    assert(!t->isTypeParameter());
+    return None;
+  });
+}
+
+/// Get the concrete type of this equivalence class.
+///
+/// Asserts if this equivalence class is not concrete.
+Type EquivalenceClass::getConcreteType(
+    TypeArrayView<GenericTypeParamType> genericParams,
+    const ProtocolGraph &protos,
+    RewriteContext &ctx) const {
+  return getTypeFromSubstitutionSchema(ConcreteType->getConcreteType(),
+                                       ConcreteType->getSubstitutions(),
+                                       genericParams,
+                                       protos,
+                                       ctx);
+}
+
 /// Given a concrete type that is a structural sub-component of a concrete
 /// type produced by RewriteSystemBuilder::getConcreteSubstitutionSchema(),
 /// collect the subset of referenced substitutions and renumber the generic
@@ -472,6 +513,157 @@ void EquivalenceClassMap::addProperty(
                           inducedRules, DebugConcreteUnification);
 }
 
+void EquivalenceClassMap::concretizeNestedTypesFromConcreteParents(
+    SmallVectorImpl<std::pair<MutableTerm, MutableTerm>> &inducedRules) const {
+  for (const auto &equivClass : Map) {
+    if (equivClass->isConcreteType() &&
+        !equivClass->getConformsTo().empty()) {
+      if (DebugConcretizeNestedTypes) {
+        llvm::dbgs() << "^ Concretizing nested types of ";
+        equivClass->dump(llvm::dbgs());
+        llvm::dbgs() << "\n";
+      }
+
+      concretizeNestedTypesFromConcreteParent(
+          equivClass->getKey(),
+          equivClass->ConcreteType->getConcreteType(),
+          equivClass->ConcreteType->getSubstitutions(),
+          equivClass->getConformsTo(),
+          inducedRules);
+    }
+  }
+}
+
+/// If we have an equivalence class T => { conforms_to: [ P ], concrete: Foo },
+/// then for each associated type A of P, we generate a new rule:
+///
+///   T.[P:A].[concrete: Foo.A] => T.[P:A]  (if Foo.A is concrete)
+///   T.[P:A] => T.(Foo.A)                  (if Foo.A is abstract)
+///
+void EquivalenceClassMap::concretizeNestedTypesFromConcreteParent(
+    const MutableTerm &key,
+    CanType concreteType, ArrayRef<Term> substitutions,
+    ArrayRef<const ProtocolDecl *> conformsTo,
+    SmallVectorImpl<std::pair<MutableTerm, MutableTerm>> &inducedRules) const {
+  for (auto *proto : conformsTo) {
+    // FIXME: Either remove the ModuleDecl entirely from conformance lookup,
+    // or pass the correct one down in here.
+    auto *module = proto->getParentModule();
+
+    auto conformance = module->lookupConformance(concreteType,
+                                                 const_cast<ProtocolDecl *>(proto));
+    if (conformance.isInvalid()) {
+      // FIXME: Diagnose conflict
+      if (DebugConcretizeNestedTypes) {
+        llvm::dbgs() << "^^ " << concreteType << " does not conform to "
+                     << proto->getName() << "\n";
+      }
+
+      continue;
+    }
+
+    // FIXME: Maybe this can happen if the concrete type is an
+    // opaque result type?
+    assert(!conformance.isAbstract());
+
+    auto assocTypes = Protos.getProtocolInfo(proto).AssociatedTypes;
+    if (assocTypes.empty())
+      continue;
+
+    auto *concrete = conformance.getConcrete();
+
+    // We might have duplicates in the list due to diamond inheritance.
+    // FIXME: Filter those out further upstream?
+    // FIXME: This should actually be outside of the loop over the conforming protos...
+    llvm::SmallDenseSet<AssociatedTypeDecl *, 4> visited;
+    for (auto *assocType : assocTypes) {
+      if (!visited.insert(assocType).second)
+        continue;
+
+      // Get the actual protocol in case we inherited this associated type.
+      auto *actualProto = assocType->getProtocol();
+      if (actualProto != proto)
+        continue;
+
+      if (DebugConcretizeNestedTypes) {
+        llvm::dbgs() << "^^ " << "Looking up type witness for "
+                     << proto->getName() << ":" << assocType->getName()
+                     << " on " << concreteType << "\n";
+      }
+
+      auto typeWitness = concrete->getTypeWitness(assocType)
+                                 ->getCanonicalType();
+
+      if (DebugConcretizeNestedTypes) {
+        llvm::dbgs() << "^^ " << "Type witness for " << assocType->getName()
+                     << " of " << concreteType << " is " << typeWitness << "\n";
+      }
+
+      auto nestedType = Atom::forAssociatedType(proto, assocType->getName(),
+                                                Context);
+
+      MutableTerm subjectType = key;
+      subjectType.add(nestedType);
+
+      MutableTerm constraintType;
+
+      if (concreteType == typeWitness) {
+        if (DebugConcretizeNestedTypes) {
+          llvm::dbgs() << "^^ Type witness is the same as the concrete type\n";
+        }
+
+        // Add a rule T.[P:A] => T.
+        constraintType = key;
+
+      } else if (typeWitness->isTypeParameter()) {
+        // The type witness is a type parameter of the form τ_0_n.X.Y...Z,
+        // where 'n' is an index into the substitution array.
+        //
+        // Collect zero or more member type names in reverse order.
+        SmallVector<Atom, 3> atoms;
+        while (auto memberType = dyn_cast<DependentMemberType>(typeWitness)) {
+          atoms.push_back(Atom::forName(memberType->getName(), Context));
+          typeWitness = memberType.getBase();
+        }
+
+        // Get the substitution S corresponding to τ_0_n.
+        unsigned index = getGenericParamIndex(typeWitness);
+        constraintType = MutableTerm(substitutions[index]);
+
+        // Add the member type names.
+        std::reverse(atoms.begin(), atoms.end());
+        for (auto atom : atoms)
+          constraintType.add(atom);
+
+        // Add a rule T => S.X.Y...Z.
+
+      } else {
+        // The type witness is a concrete type.
+        constraintType = subjectType;
+
+        // FIXME: Handle dependent member types here
+        SmallVector<Term, 3> result;
+        auto typeWitnessSchema =
+            remapConcreteSubstitutionSchema(typeWitness, substitutions,
+                                            Context.getASTContext(),
+                                            result);
+        constraintType.add(
+            Atom::forConcreteType(
+                typeWitnessSchema, result, Context));
+
+        // Add a rule T.[P:A].[concrete: Foo.A] => T.[P:A].
+
+      }
+
+      inducedRules.emplace_back(subjectType, constraintType);
+      if (DebugConcretizeNestedTypes) {
+        llvm::dbgs() << "^^ Induced rule " << constraintType
+                     << " => " << subjectType << "\n";
+      }
+    }
+  }
+}
+
 void EquivalenceClassMap::dump(llvm::raw_ostream &out) const {
   out << "Equivalence class map: {\n";
   for (const auto &equivClass : Map) {
@@ -485,10 +677,19 @@ void EquivalenceClassMap::dump(llvm::raw_ostream &out) const {
 /// Build the equivalence class map from all rules of the form T.[p] => T, where
 /// [p] is a property atom.
 ///
-/// Returns true if concrete term unification performed while building the map
-/// introduced new rewrite rules; in this case, the completion procedure must
-/// run again.
-bool RewriteSystem::buildEquivalenceClassMap(EquivalenceClassMap &map) {
+/// Returns a pair consisting of a status and number of iterations executed.
+///
+/// The status is CompletionResult::MaxIterations if we exceed \p maxIterations
+/// iterations.
+///
+/// The status is CompletionResult::MaxDepth if we produce a rewrite rule whose
+/// left hand side has a length exceeding \p maxDepth.
+///
+/// Otherwise, the status is CompletionResult::Success.
+std::pair<RewriteSystem::CompletionResult, unsigned>
+RewriteSystem::buildEquivalenceClassMap(EquivalenceClassMap &map,
+                                        unsigned maxIterations,
+                                        unsigned maxDepth) {
   map.clear();
 
   std::vector<std::pair<MutableTerm, Atom>> properties;
@@ -533,18 +734,25 @@ bool RewriteSystem::buildEquivalenceClassMap(EquivalenceClassMap &map) {
     map.addProperty(pair.first, pair.second, inducedRules);
   }
 
+  // We also need to merge concrete type rules with conformance rules, by
+  // concretizing the associated type witnesses of the concrete type.
+  map.concretizeNestedTypesFromConcreteParents(inducedRules);
+
   // Some of the induced rules might be trivial; only count the induced rules
   // where the left hand side is not already equivalent to the right hand side.
   unsigned addedNewRules = 0;
   for (auto pair : inducedRules) {
-    if (addRule(pair.first, pair.second))
+    if (addRule(pair.first, pair.second)) {
       ++addedNewRules;
-  }
 
-  if (auto *stats = Context.getASTContext().Stats) {
-    stats->getFrontendCounters()
-      .NumRequirementMachineUnifiedConcreteTerms += addedNewRules;
+      const auto &newRule = Rules.back();
+      if (newRule.getLHS().size() > maxDepth)
+        return std::make_pair(CompletionResult::MaxDepth, addedNewRules);
+    }
   }
 
-  return addedNewRules > 0;
+  if (Rules.size() > maxIterations)
+    return std::make_pair(CompletionResult::MaxIterations, addedNewRules);
+
+  return std::make_pair(CompletionResult::Success, addedNewRules);
 }

lib/AST/RequirementMachine/EquivalenceClassMap.h

@@ -86,6 +86,11 @@ class EquivalenceClass {
     return ConcreteType.hasValue();
   }
 
+  Type getConcreteType(
+      TypeArrayView<GenericTypeParamType> genericParams,
+      const ProtocolGraph &protos,
+      RewriteContext &ctx) const;
+
   LayoutConstraint getLayoutConstraint() const {
     return Layout;
   }
@@ -103,7 +108,8 @@ class EquivalenceClassMap {
   RewriteContext &Context;
   std::vector<std::unique_ptr<EquivalenceClass>> Map;
   const ProtocolGraph &Protos;
-  bool DebugConcreteUnification = false;
+  unsigned DebugConcreteUnification : 1;
+  unsigned DebugConcretizeNestedTypes : 1;
 
   EquivalenceClass *getEquivalenceClassIfPresent(const MutableTerm &key) const;
   EquivalenceClass *getOrCreateEquivalenceClass(const MutableTerm &key);
@@ -116,14 +122,28 @@ class EquivalenceClassMap {
 public:
   explicit EquivalenceClassMap(RewriteContext &ctx,
                                const ProtocolGraph &protos)
-      : Context(ctx), Protos(protos) {}
+      : Context(ctx), Protos(protos) {
+    DebugConcreteUnification = false;
+    DebugConcretizeNestedTypes = false;
+  }
 
   EquivalenceClass *lookUpEquivalenceClass(const MutableTerm &key) const;
 
+  void dump(llvm::raw_ostream &out) const;
+
   void clear();
   void addProperty(const MutableTerm &key, Atom property,
                    SmallVectorImpl<std::pair<MutableTerm, MutableTerm>> &inducedRules);
-  void dump(llvm::raw_ostream &out) const;
+  void concretizeNestedTypesFromConcreteParents(
+                   SmallVectorImpl<std::pair<MutableTerm, MutableTerm>> &inducedRules) const;
+
+private:
+  void concretizeNestedTypesFromConcreteParent(
+                   const MutableTerm &key,
+                   CanType concreteType,
+                   ArrayRef<Term> substitutions,
+                   ArrayRef<const ProtocolDecl *> conformsTo,
+                   SmallVectorImpl<std::pair<MutableTerm, MutableTerm>> &inducedRules) const;
 };
 
 } // end namespace rewriting