[IRGen] Fix logic for finding witness tables in an opaque result type.

Use the same logic as we do for forming the opaque type result descriptor,
of course.

Author: DougGregor

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
 

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
@@ -61,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)
@@ -70,6 +86,16 @@ 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")