Opaque type fixes for interesting generic environments

include/swift/SIL/AbstractionPattern.h

@@ -981,7 +981,7 @@ class AbstractionPattern {
         return true;
       }
       if (auto archetype = dyn_cast<ArchetypeType>(type)) {
-        return !isa<OpaqueTypeArchetypeType>(archetype->getRoot());
+        return !isa<OpaqueTypeArchetypeType>(archetype);
       }
       return false;
     }

lib/AST/ASTVerifier.cpp

@@ -2807,6 +2807,19 @@ class Verifier : public ASTWalker {
 
       unsigned currentDepth = DC->getGenericContextDepth();
       if (currentDepth < GTPD->getDepth()) {
+        // If this is actually an opaque type's generic parameter, we're okay.
+        if (auto value = dyn_cast_or_null<ValueDecl>(DC->getAsDecl())) {
+          auto opaqueDecl = dyn_cast<OpaqueTypeDecl>(value);
+          if (!opaqueDecl)
+            opaqueDecl = value->getOpaqueResultTypeDecl();
+          if (opaqueDecl) {
+            if (GTPD->getDepth() ==
+                    opaqueDecl->getOpaqueGenericParams().front()->getDepth()) {
+              return;
+            }
+          }
+        }
+
         Out << "GenericTypeParamDecl has incorrect depth\n";
         abort();
       }

lib/AST/GenericEnvironment.cpp

@@ -362,12 +362,24 @@ GenericEnvironment::getOrCreateArchetypeFromInterfaceType(Type depType) {
       break;
     }
 
-    case Kind::Opaque:
+    case Kind::Opaque: {
+      // If the anchor type isn't rooted in a generic parameter that
+      // represents an opaque declaration, then apply the outer substitutions.
+      // It would be incorrect to build an opaque type archetype here.
+      auto rootGP = requirements.anchor->getRootGenericParam();
+      unsigned opaqueDepth =
+          getOpaqueTypeDecl()->getOpaqueGenericParams().front()->getDepth();
+      if (rootGP->getDepth() < opaqueDepth) {
+        result = maybeApplyOpaqueTypeSubstitutions(requirements.anchor);
+        break;
+      }
+
       result = OpaqueTypeArchetypeType::getNew(this, requirements.anchor,
                                                requirements.protos, superclass,
                                                requirements.layout);
       break;
     }
+    }
   }
 
   if (genericParam)

lib/AST/ProtocolConformance.cpp

@@ -115,7 +115,7 @@ ProtocolConformanceRef::subst(Type origType,
   // unless we're specifically substituting opaque types.
   if (auto origArchetype = origType->getAs<ArchetypeType>()) {
     if (!options.contains(SubstFlags::SubstituteOpaqueArchetypes)
-        && isa<OpaqueTypeArchetypeType>(origArchetype->getRoot())) {
+        && isa<OpaqueTypeArchetypeType>(origArchetype)) {
       return *this;
     }
   }

lib/AST/SubstitutionMap.cpp

@@ -326,7 +326,7 @@ SubstitutionMap::lookupConformance(CanType type, ProtocolDecl *proto) const {
   // If we have an archetype, map out of the context so we can compute a
   // conformance access path.
   if (auto archetype = dyn_cast<ArchetypeType>(type)) {
-    if (!isa<OpaqueTypeArchetypeType>(archetype->getRoot())) {
+    if (!isa<OpaqueTypeArchetypeType>(archetype)) {
       type = archetype->getInterfaceType()->getCanonicalType();
     }
   }

lib/AST/TypeWalker.cpp

@@ -207,8 +207,8 @@ class Traversal : public TypeVisitor<Traversal, bool>
   bool visitArchetypeType(ArchetypeType *ty) {
     // If the root is an opaque archetype, visit its substitution replacement
     // types.
-    if (auto opaqueRoot = dyn_cast<OpaqueTypeArchetypeType>(ty->getRoot())) {
-      for (auto arg : opaqueRoot->getSubstitutions().getReplacementTypes()) {
+    if (auto opaque = dyn_cast<OpaqueTypeArchetypeType>(ty)) {
+      for (auto arg : opaque->getSubstitutions().getReplacementTypes()) {
         if (doIt(arg)) {
           return true;
         }

lib/IRGen/DebugTypeInfo.h

@@ -93,7 +93,7 @@ class DebugTypeInfo {
   bool isContextArchetype() const {
     if (auto archetype =
             Type->getWithoutSpecifierType()->getAs<ArchetypeType>()) {
-      return !isa<OpaqueTypeArchetypeType>(archetype->getRoot());
+      return !isa<OpaqueTypeArchetypeType>(archetype);
     }
     return false;
   }

lib/IRGen/GenArchetype.cpp

@@ -537,16 +537,32 @@ llvm::Value *irgen::emitOpaqueTypeWitnessTableRef(IRGenFunction &IGF,
                                           ProtocolDecl *protocol) {
   auto accessorFn = IGF.IGM.getGetOpaqueTypeConformanceFn();
   auto opaqueDecl = archetype->getDecl();
-
+  assert(archetype->isRoot() && "Can only follow from the root");
 
   llvm::Value *descriptor = getAddressOfOpaqueTypeDescriptor(IGF, opaqueDecl);
 
-  auto foundProtocol = std::find(archetype->getConformsTo().begin(),
-                                 archetype->getConformsTo().end(),
-                                 protocol);
-  assert(foundProtocol != archetype->getConformsTo().end());
-
-  unsigned index = foundProtocol - archetype->getConformsTo().begin() + 1;
+  // Compute the index at which this witness table resides.
+  unsigned index = opaqueDecl->getOpaqueGenericParams().size();
+  auto opaqueReqs =
+      opaqueDecl->getOpaqueInterfaceGenericSignature().getRequirements();
+  bool found = false;
+  for (const auto &req : opaqueReqs) {
+    auto reqProto = opaqueTypeRequiresWitnessTable(opaqueDecl, req);
+    if (!reqProto)
+      continue;
+
+    // Is this requirement the one we're looking for?
+    if (reqProto == protocol &&
+        req.getFirstType()->isEqual(archetype->getInterfaceType())) {
+      found = true;
+      break;
+    }
+
+    ++index;
+  }
+
+  (void)found;
+  assert(found && "Opaque type does not conform to protocol");
   auto indexValue = llvm::ConstantInt::get(IGF.IGM.SizeTy, index);
 
   llvm::CallInst *result = nullptr;

lib/IRGen/GenMeta.cpp

@@ -2007,29 +2007,7 @@ namespace {
     ///
     /// When it does, returns the protocol.
     ProtocolDecl *requiresWitnessTable(const Requirement &req) const {
-      // We only care about conformance requirements.
-      if (req.getKind() != RequirementKind::Conformance)
-        return nullptr;
-
-      // The protocol must require a witness table.
-      auto proto = req.getProtocolDecl();
-      if (!Lowering::TypeConverter::protocolRequiresWitnessTable(proto))
-        return nullptr;
-
-      // The type itself must be anchored on one of the generic parameters of
-      // the opaque type (not an outer context).
-      Type subject = req.getFirstType();
-      while (auto depMember = subject->getAs<DependentMemberType>()) {
-        subject = depMember->getBase();
-      }
-
-      if (auto genericParam = subject->getAs<GenericTypeParamType>()) {
-        unsigned opaqueDepth = O->getOpaqueGenericParams().front()->getDepth();
-        if (genericParam->getDepth() == opaqueDepth)
-          return proto;
-      }
-
-      return nullptr;
+      return opaqueTypeRequiresWitnessTable(O, req);
     }
 
   public:
@@ -2155,6 +2133,33 @@ namespace {
   };
 } // end anonymous namespace
 
+ProtocolDecl *irgen::opaqueTypeRequiresWitnessTable(
+    OpaqueTypeDecl *opaque, const Requirement &req) {
+  // We only care about conformance requirements.
+  if (req.getKind() != RequirementKind::Conformance)
+    return nullptr;
+
+  // The protocol must require a witness table.
+  auto proto = req.getProtocolDecl();
+  if (!Lowering::TypeConverter::protocolRequiresWitnessTable(proto))
+    return nullptr;
+
+  // The type itself must be anchored on one of the generic parameters of
+  // the opaque type (not an outer context).
+  Type subject = req.getFirstType();
+  while (auto depMember = subject->getAs<DependentMemberType>()) {
+    subject = depMember->getBase();
+  }
+
+  if (auto genericParam = subject->getAs<GenericTypeParamType>()) {
+    unsigned opaqueDepth = opaque->getOpaqueGenericParams().front()->getDepth();
+    if (genericParam->getDepth() == opaqueDepth)
+      return proto;
+  }
+
+  return nullptr;
+}
+
 static void eraseExistingTypeContextDescriptor(IRGenModule &IGM,
                                                NominalTypeDecl *type) {
   // We may have emitted a partial type context descriptor with some empty

lib/IRGen/GenMeta.h

@@ -187,6 +187,16 @@ namespace irgen {
                                               llvm::Function *function,
                                               LinkEntity entity,
                                               Size size);
+
+  /// Determine whether the given opaque type requires a witness table for the
+  /// given requirement.
+  ///
+  /// \returns the protocol when a witness table is required, or \c nullptr
+  /// if the requirement isn't a conformance requirement or doesn't require a
+  /// witness table.
+  ProtocolDecl *opaqueTypeRequiresWitnessTable(
+      OpaqueTypeDecl *opaque, const Requirement &req);
+
 } // end namespace irgen
 } // end namespace swift
 

lib/IRGen/GenReflection.cpp

@@ -200,7 +200,7 @@ getRuntimeVersionThatSupportsDemanglingType(CanType type) {
   if (type->hasOpaqueArchetype()) {
     auto hasOpaqueAssocType = type.findIf([](CanType t) -> bool {
       if (auto a = dyn_cast<ArchetypeType>(t)) {
-        return isa<OpaqueTypeArchetypeType>(a->getRoot()) &&
+        return isa<OpaqueTypeArchetypeType>(a) &&
           a->getInterfaceType()->is<DependentMemberType>();
       }
       return false;

lib/IRGen/Outlining.cpp

@@ -117,11 +117,10 @@ irgen::getTypeAndGenericSignatureForManglingOutlineFunction(SILType type) {
     GenericEnvironment *env = nullptr;
     loweredType.findIf([&env](Type t) -> bool {
         if (auto arch = t->getAs<ArchetypeType>()) {
-          auto root = arch->getRoot();
-          if (!isa<PrimaryArchetypeType>(root) &&
-              !isa<VariadicSequenceType>(root))
+          if (!isa<PrimaryArchetypeType>(arch) &&
+              !isa<VariadicSequenceType>(arch))
             return false;
-          env = root->getGenericEnvironment();
+          env = arch->getGenericEnvironment();
           return true;
         }
         return false;

lib/SIL/Verifier/SILVerifier.cpp

@@ -86,18 +86,17 @@ static llvm::cl::opt<bool> AllowCriticalEdges("allow-critical-edges",
 /// Returns true if A is an opened existential type or is equal to an
 /// archetype from F's generic context.
 static bool isArchetypeValidInFunction(ArchetypeType *A, const SILFunction *F) {
-  auto root = A->getRoot();
-  if (!isa<PrimaryArchetypeType>(root) && !isa<SequenceArchetypeType>(root))
+  if (!isa<PrimaryArchetypeType>(A) && !isa<SequenceArchetypeType>(A))
     return true;
-  if (isa<OpenedArchetypeType>(root))
+  if (isa<OpenedArchetypeType>(A))
     return true;
-  if (isa<OpaqueTypeArchetypeType>(root))
+  if (isa<OpaqueTypeArchetypeType>(A))
     return true;
 
   // Ok, we have a primary archetype, make sure it is in the nested generic
   // environment of our caller.
   if (auto *genericEnv = F->getGenericEnvironment())
-    if (root->getGenericEnvironment() == genericEnv)
+    if (A->getGenericEnvironment() == genericEnv)
       return true;
 
   return false;

lib/SILOptimizer/Utils/Generics.cpp

@@ -1092,11 +1092,11 @@ shouldBePartiallySpecialized(Type Replacement,
   llvm::SmallSetVector<ArchetypeType *, 2> UsedArchetypes;
   Replacement.visit([&](Type Ty) {
     if (auto Archetype = Ty->getAs<ArchetypeType>()) {
-      if (auto Primary = dyn_cast<PrimaryArchetypeType>(Archetype->getRoot())) {
+      if (auto Primary = dyn_cast<PrimaryArchetypeType>(Archetype)) {
         UsedArchetypes.insert(Primary);
       }
 
-      if (auto Seq = dyn_cast<SequenceArchetypeType>(Archetype->getRoot())) {
+      if (auto Seq = dyn_cast<SequenceArchetypeType>(Archetype)) {
         UsedArchetypes.insert(Seq);
       }
     }
@@ -1322,11 +1322,11 @@ void FunctionSignaturePartialSpecializer::collectUsedCallerArchetypes(
     // Add used generic parameters/archetypes.
     Replacement.visit([&](Type Ty) {
       if (auto Archetype = Ty->getAs<ArchetypeType>()) {
-        if (auto Primary = dyn_cast<PrimaryArchetypeType>(Archetype->getRoot())) {
+        if (auto Primary = dyn_cast<PrimaryArchetypeType>(Archetype)) {
           UsedCallerArchetypes.insert(Primary);
         }
 
-        if (auto Seq = dyn_cast<SequenceArchetypeType>(Archetype->getRoot())) {
+        if (auto Seq = dyn_cast<SequenceArchetypeType>(Archetype)) {
           UsedCallerArchetypes.insert(Seq);
         }
       }

lib/Sema/CSSimplify.cpp

@@ -2849,8 +2849,8 @@ ConstraintSystem::matchDeepEqualityTypes(Type type1, Type type2,
   // Handle opaque archetypes.
   if (auto arch1 = type1->getAs<ArchetypeType>()) {
     auto arch2 = type2->castTo<ArchetypeType>();
-    auto opaque1 = cast<OpaqueTypeArchetypeType>(arch1->getRoot());
-    auto opaque2 = cast<OpaqueTypeArchetypeType>(arch2->getRoot());
+    auto opaque1 = cast<OpaqueTypeArchetypeType>(arch1);
+    auto opaque2 = cast<OpaqueTypeArchetypeType>(arch2);
     assert(opaque1->getDecl() == opaque2->getDecl());
     assert(opaque1->getCanonicalInterfaceType(arch1->getInterfaceType())->isEqual(
                opaque2->getCanonicalInterfaceType(arch2->getInterfaceType())));

lib/Sema/TypeCheckConstraints.cpp

@@ -692,14 +692,15 @@ static Type replaceArchetypesWithTypeVariables(ConstraintSystem &cs,
         return found->second;
 
       if (auto archetypeType = dyn_cast<ArchetypeType>(origType)) {
-        auto root = archetypeType->getRoot();
         // We leave opaque types and their nested associated types alone here.
         // They're globally available.
-        if (isa<OpaqueTypeArchetypeType>(root))
+        if (isa<OpaqueTypeArchetypeType>(archetypeType))
           return origType;
+
+        auto root = archetypeType->getRoot();
         // For other nested types, fail here so the default logic in subst()
         // for nested types applies.
-        else if (root != archetypeType)
+        if (root != archetypeType)
           return Type();
 
         auto locator = cs.getConstraintLocator({});

test/Interpreter/named_opaque_result_types.swift

@@ -1,5 +1,5 @@
 // RUN: %empty-directory(%t)
-// RUN: %target-build-swift -swift-version 5 -Xfrontend -enable-experimental-named-opaque-types %s -o %t/a.out
+// RUN: %target-build-swift -g -swift-version 5 -Xfrontend -enable-experimental-named-opaque-types %s -o %t/a.out
 // RUN: %target-run %t/a.out | %FileCheck %s
 
 // REQUIRES: executable_test
@@ -34,8 +34,7 @@ struct X<T, U: Hashable>: P {
   func f() -> <
       R1: Collection, R2: Collection where R1.Element == T, R2.Element == U
   > (R1, R2) {
-    // FIXME: Use unzipCollection here
-    return (Array(data.map { $0.0 }), Set(data.map { $0.1 }))
+    return unzipCollection(data)
   }
 }
 
@@ -62,6 +61,22 @@ if #available(iOS 13, macOS 10.15, tvOS 13, watchOS 6, *) {
   // CHECK: {{["Hello", "World", "!"]|too old}}
   print(x3.1)
 
+  // CHECK: Integer loop
+  // CHECK-NEXT: 1
+  // CHECK-NEXT: 2
+  // CHECK-NEXT: 3
+  print("Integer loop")
+  for i in x3.0 {
+    print(i)
+  }
+
+  // CHECK: Hello
+  // CHECK-NEXT: World
+  // CHECK-NEXT: !
+  for index in x3.1.indices {
+    print(x3.1[index])
+  }
+
   // CHECK: {{(Array<Int>, Set<String>)|too old}}
   let paType = getP_A(X<Int, String>.self)
   print(paType)
@@ -71,6 +86,17 @@ if #available(iOS 13, macOS 10.15, tvOS 13, watchOS 6, *) {
   print("too old")
   print("too old")
   print("too old")
+  print("too old")
+
+  print("Integer loop")
+  print("1")
+  print("2")
+  print("3")
+
+  print("Hello")
+  print("World")
+  print("!")
+
   print("too old")
   print("too old")
 }