AST: ProtocolConformanceRef can store a pack conformance

slavapestov · slavapestov · commit 0f26d02fb140 · 2022-08-23T21:07:30.000-04:00
diff --git a/include/swift/AST/PackConformance.h b/include/swift/AST/PackConformance.h
@@ -77,12 +77,16 @@ class alignas(1 << DeclAlignInBits) PackConformance final
   PackConformance *
   getAssociatedConformance(Type assocType, ProtocolDecl *protocol) const;
 
-  PackConformance *subst(SubstitutionMap subMap,
-                         SubstOptions options=None) const;
-
-  PackConformance *subst(TypeSubstitutionFn subs,
-                         LookupConformanceFn conformances,
-                         SubstOptions options=None) const;
+  /// The ProtocolConformanceRef either stores a pack conformance, or
+  /// it is invalid in the case of substitution failure.
+  ProtocolConformanceRef subst(SubstitutionMap subMap,
+                               SubstOptions options=None) const;
+
+  /// The ProtocolConformanceRef either stores a pack conformance, or
+  /// it is invalid in the case of substitution failure.
+  ProtocolConformanceRef subst(TypeSubstitutionFn subs,
+                               LookupConformanceFn conformances,
+                               SubstOptions options=None) const;
 
   SWIFT_DEBUG_DUMP;
   void dump(llvm::raw_ostream &out, unsigned indent = 0) const;
diff --git a/include/swift/AST/ProtocolConformanceRef.h b/include/swift/AST/ProtocolConformanceRef.h
@@ -33,6 +33,7 @@ namespace swift {
 
 class BuiltinProtocolConformance;
 class ConcreteDeclRef;
+class PackConformance;
 class ProtocolConformance;
 enum class EffectKind : uint8_t;
 
@@ -50,7 +51,9 @@ enum class EffectKind : uint8_t;
 /// ProtocolConformanceRef allows the efficient recovery of the protocol
 /// even when the conformance is abstract.
 class ProtocolConformanceRef {
-  using UnionType = llvm::PointerUnion<ProtocolDecl*, ProtocolConformance*>;
+  using UnionType = llvm::PointerUnion<ProtocolDecl *,
+                                       ProtocolConformance *,
+                                       PackConformance *>;
   UnionType Union;
 
   explicit ProtocolConformanceRef(UnionType value) : Union(value) {}
@@ -68,6 +71,12 @@ class ProtocolConformanceRef {
            "cannot construct ProtocolConformanceRef with null");
   }
 
+  /// Create a pack protocol conformance reference.
+  explicit ProtocolConformanceRef(PackConformance *conf) : Union(conf) {
+    assert(conf != nullptr &&
+           "cannot construct ProtocolConformanceRef with null");
+  }
+
   ProtocolConformanceRef(std::nullptr_t = nullptr)
       : Union((ProtocolDecl *)nullptr) {}
 
@@ -98,6 +107,13 @@ class ProtocolConformanceRef {
     return Union.get<ProtocolConformance*>();
   }
 
+  bool isPack() const {
+    return !isInvalid() && Union.is<PackConformance*>();
+  }
+  PackConformance *getPack() const {
+    return Union.get<PackConformance*>();
+  }
+
   bool isAbstract() const {
     return !isInvalid() && Union.is<ProtocolDecl*>();
   }
diff --git a/lib/AST/ASTDumper.cpp b/lib/AST/ASTDumper.cpp
@@ -3260,6 +3260,11 @@ static void dumpProtocolConformanceRec(
     unsigned indent,
     llvm::SmallPtrSetImpl<const ProtocolConformance *> &visited);
 
+static void dumpPackConformanceRec(
+    const PackConformance *conformance, llvm::raw_ostream &out,
+    unsigned indent,
+    llvm::SmallPtrSetImpl<const ProtocolConformance *> &visited);
+
 static void dumpProtocolConformanceRefRec(
     const ProtocolConformanceRef conformance, llvm::raw_ostream &out,
     unsigned indent,
@@ -3268,7 +3273,11 @@ static void dumpProtocolConformanceRefRec(
     out.indent(indent) << "(invalid_conformance)";
   } else if (conformance.isConcrete()) {
     dumpProtocolConformanceRec(conformance.getConcrete(), out, indent, visited);
+  } else if (conformance.isPack()) {
+    dumpPackConformanceRec(conformance.getPack(), out, indent, visited);
   } else {
+    assert(conformance.isAbstract());
+
     out.indent(indent) << "(abstract_conformance protocol="
                        << conformance.getAbstract()->getName();
     PrintWithColorRAII(out, ParenthesisColor) << ')';
@@ -3509,7 +3518,6 @@ void ProtocolConformanceRef::dump(llvm::raw_ostream &out, unsigned indent,
 void ProtocolConformanceRef::print(llvm::raw_ostream &out) const {
   llvm::SmallPtrSet<const ProtocolConformance *, 8> visited;
   dumpProtocolConformanceRefRec(*this, out, 0, visited);
-
 }
 
 void ProtocolConformance::dump() const {
diff --git a/lib/AST/PackConformance.cpp b/lib/AST/PackConformance.cpp
@@ -160,8 +160,8 @@ PackConformance *PackConformance::getAssociatedConformance(
   return PackConformance::get(conformingType, protocol, packConformances);
 }
 
-PackConformance *PackConformance::subst(SubstitutionMap subMap,
-                                        SubstOptions options) const {
+ProtocolConformanceRef PackConformance::subst(SubstitutionMap subMap,
+                                              SubstOptions options) const {
   return subst(QuerySubstitutionMap{subMap},
                LookUpConformanceInSubstitutionMap(subMap),
                options);
@@ -284,9 +284,9 @@ class PackExpander {
                                        ProtocolDecl *proto) -> ProtocolConformanceRef {
         auto substConformance = conformances(origType, substType, proto);
 
+        // If the substituted conformance is a pack, project the jth element.
         if (isRootedInTypeSequenceParameter(origType)) {
-          // FIXME: get the corresponding pack conformance
-          return ProtocolConformanceRef::forInvalid();
+          return substConformance.getPack()->getPatternConformances()[j];
         }
 
         return substConformance;
@@ -376,18 +376,19 @@ class PackConformanceExpander : public PackExpander<PackConformanceExpander> {
 
 }
 
-PackConformance *PackConformance::subst(TypeSubstitutionFn subs,
-                                        LookupConformanceFn conformances,
-                                        SubstOptions options) const {
+ProtocolConformanceRef PackConformance::subst(TypeSubstitutionFn subs,
+                                              LookupConformanceFn conformances,
+                                              SubstOptions options) const {
   PackConformanceExpander expander(subs, conformances, options,
                                    getPatternConformances());
   expander.expand(ConformingType);
 
   auto &ctx = Protocol->getASTContext();
   auto *substConformingType = PackType::get(ctx, expander.substElements);
 
-  return PackConformance::get(substConformingType, Protocol,
-                              expander.substConformances);
+  auto substConformance = PackConformance::get(substConformingType, Protocol,
+                                               expander.substConformances);
+  return ProtocolConformanceRef(substConformance);
 }
 
 void swift::simple_display(llvm::raw_ostream &out, PackConformance *conformance) {
diff --git a/lib/AST/ProtocolConformanceRef.cpp b/lib/AST/ProtocolConformanceRef.cpp
@@ -20,6 +20,8 @@
 #include "swift/AST/Availability.h"
 #include "swift/AST/Decl.h"
 #include "swift/AST/Module.h"
+#include "swift/AST/PackConformance.h"
+#include "swift/AST/ProtocolConformance.h"
 #include "swift/AST/TypeCheckRequests.h"
 #include "swift/AST/Types.h"
 
@@ -44,6 +46,8 @@ ProtocolDecl *ProtocolConformanceRef::getRequirement() const {
 
   if (isConcrete()) {
     return getConcrete()->getProtocol();
+  } else if (isPack()) {
+    return getPack()->getProtocol();
   } else {
     return getAbstract();
   }
@@ -67,11 +71,14 @@ ProtocolConformanceRef::subst(Type origType,
   if (isInvalid())
     return *this;
 
-  // If we have a concrete conformance, we need to substitute the
-  // conformance to apply to the new type.
   if (isConcrete())
     return ProtocolConformanceRef(getConcrete()->subst(subs, conformances,
                                                        options));
+  if (isPack())
+    return getPack()->subst(subs, conformances, options);
+
+  // Handle abstract conformances below:
+
   // If the type is an opaque archetype, the conformance will remain abstract,
   // unless we're specifically substituting opaque types.
   if (auto origArchetype = origType->getAs<ArchetypeType>()) {
@@ -102,17 +109,26 @@ ProtocolConformanceRef::subst(Type origType,
 }
 
 ProtocolConformanceRef ProtocolConformanceRef::mapConformanceOutOfContext() const {
-  if (!isConcrete())
-    return *this;
+  if (isConcrete()) {
+    auto *concrete = getConcrete()->subst(
+        [](SubstitutableType *type) -> Type {
+          if (auto *archetypeType = type->getAs<ArchetypeType>())
+            return archetypeType->getInterfaceType();
+          return type;
+        },
+        MakeAbstractConformanceForGenericType());
+    return ProtocolConformanceRef(concrete);
+  } else if (isPack()) {
+    return getPack()->subst(
+        [](SubstitutableType *type) -> Type {
+          if (auto *archetypeType = type->getAs<ArchetypeType>())
+            return archetypeType->getInterfaceType();
+          return type;
+        },
+        MakeAbstractConformanceForGenericType());
+  }
 
-  auto *concrete = getConcrete()->subst(
-      [](SubstitutableType *type) -> Type {
-        if (auto *archetypeType = type->getAs<ArchetypeType>())
-          return archetypeType->getInterfaceType();
-        return type;
-      },
-      MakeAbstractConformanceForGenericType());
-  return ProtocolConformanceRef(concrete);
+  return *this;
 }
 
 Type
@@ -171,6 +187,12 @@ ProtocolConformanceRef::getConditionalRequirements() const {
 
 Type ProtocolConformanceRef::getAssociatedType(Type conformingType,
                                                Type assocType) const {
+  if (isPack()) {
+    auto *pack = getPack();
+    assert(conformingType->isEqual(pack->getType()));
+    return pack->getAssociatedType(assocType);
+  }
+
   assert(!isConcrete() || getConcrete()->getType()->isEqual(conformingType));
 
   auto type = assocType->getCanonicalType();
@@ -200,6 +222,14 @@ ProtocolConformanceRef
 ProtocolConformanceRef::getAssociatedConformance(Type conformingType,
                                                  Type assocType,
                                                  ProtocolDecl *protocol) const {
+  // If this is a pack conformance, project the associated conformances.
+  if (isPack()) {
+    auto *pack = getPack();
+    assert(conformingType->isEqual(pack->getType()));
+    return ProtocolConformanceRef(
+        pack->getAssociatedConformance(assocType, protocol));
+  }
+
   // If this is a concrete conformance, look up the associated conformance.
   if (isConcrete()) {
     auto conformance = getConcrete();
@@ -221,23 +251,33 @@ ProtocolConformanceRef::getAssociatedConformance(Type conformingType,
 bool ProtocolConformanceRef::isCanonical() const {
   if (isAbstract() || isInvalid())
     return true;
+  if (isPack())
+    return getPack()->isCanonical();
   return getConcrete()->isCanonical();
+
 }
 
 ProtocolConformanceRef
 ProtocolConformanceRef::getCanonicalConformanceRef() const {
   if (isAbstract() || isInvalid())
     return *this;
+  if (isPack())
+    return ProtocolConformanceRef(getPack()->getCanonicalConformance());
   return ProtocolConformanceRef(getConcrete()->getCanonicalConformance());
 }
 
 bool ProtocolConformanceRef::hasUnavailableConformance() const {
-  if (isInvalid())
+  if (isInvalid() || isAbstract())
     return false;
 
-  // Abstract conformances are never unavailable.
-  if (!isConcrete())
+  if (isPack()) {
+    for (auto conformance : getPack()->getPatternConformances()) {
+      if (conformance.hasUnavailableConformance())
+        return true;
+    }
+
     return false;
+  }
 
   // Check whether this conformance is on an unavailable extension.
   auto concrete = getConcrete();
@@ -266,8 +306,17 @@ bool ProtocolConformanceRef::hasMissingConformance(ModuleDecl *module) const {
 bool ProtocolConformanceRef::forEachMissingConformance(
     ModuleDecl *module,
     llvm::function_ref<bool(BuiltinProtocolConformance *missing)> fn) const {
-  if (!isConcrete())
+  if (isInvalid() || isAbstract())
+    return false;
+
+  if (isPack()) {
+    for (auto conformance : getPack()->getPatternConformances()) {
+      if (conformance.forEachMissingConformance(module, fn))
+        return true;
+    }
+
     return false;
+  }
 
   // Is this a missing conformance?
   ProtocolConformance *concreteConf = getConcrete();
@@ -292,6 +341,8 @@ void swift::simple_display(llvm::raw_ostream &out, ProtocolConformanceRef confor
     simple_display(out, conformanceRef.getAbstract());
   } else if (conformanceRef.isConcrete()) {
     simple_display(out, conformanceRef.getConcrete());
+  } else if (conformanceRef.isPack()) {
+    simple_display(out, conformanceRef.getPack());
   }
 }
 
