Merge pull request #78855 from hamishknight/elemental

[Mangler] Handle local archetypes in `getDeclTypeForMangling`

lib/AST/ASTMangler.cpp

@@ -4016,6 +4016,43 @@ static bool isMethodDecl(const Decl *decl) {
     && decl->getDeclContext()->isTypeContext();
 }
 
+/// Map any local archetypes in a decl's interface type out of context, such
+/// that the resulting type is suitable for mangling.
+///
+/// Note this does not guarantee that different archetypes produce different
+/// interface types across decls, but it is guaranteed within a single decl
+/// type. This is okay though since local decls are assigned discriminators.
+static Type mapLocalArchetypesOutOfContextForDecl(const ValueDecl *decl,
+                                                  Type ty) {
+  if (!ty->hasLocalArchetype())
+    return ty;
+
+  ASSERT(decl->getDeclContext()->isLocalContext());
+
+  CaptureInfo captureInfo;
+  auto *innerDC = decl->getInnermostDeclContext();
+  auto genericSig = innerDC->getGenericSignatureOfContext();
+  if (auto fn = AnyFunctionRef::fromDeclContext(innerDC))
+    captureInfo = fn->getCaptureInfo();
+
+  // Record any captured generic environments we have.
+  llvm::SmallSetVector<GenericEnvironment *, 4> capturedEnvs;
+  for (auto *genericEnv : captureInfo.getGenericEnvironments())
+    capturedEnvs.insert(genericEnv);
+
+  // We may still have archetypes local to the current context, e.g for
+  // decls in local for loops over pack expansions. In this case, collect
+  // any remaining generic environments from the type.
+  ty.visit([&](Type t) {
+    if (auto *archetypeTy = t->getAs<LocalArchetypeType>()) {
+      capturedEnvs.insert(archetypeTy->getGenericEnvironment());
+    }
+  });
+
+  return swift::mapLocalArchetypesOutOfContext(ty, genericSig,
+                                               capturedEnvs.getArrayRef());
+}
+
 CanType ASTMangler::getDeclTypeForMangling(
                                        const ValueDecl *decl,
                                        GenericSignature &genericSig,
@@ -4048,6 +4085,9 @@ CanType ASTMangler::getDeclTypeForMangling(
     ty = ty->stripConcurrency(/*recurse=*/true, /*dropGlobalActor=*/true);
   }
 
+  // Map any local archetypes out of context.
+  ty = mapLocalArchetypesOutOfContextForDecl(decl, ty);
+
   auto canTy = ty->getCanonicalType();
 
   if (auto gft = dyn_cast<GenericFunctionType>(canTy)) {

lib/AST/LocalArchetypeRequirementCollector.cpp

@@ -189,7 +189,11 @@ Type MapLocalArchetypesOutOfContext::getInterfaceType(
 }
 
 Type MapLocalArchetypesOutOfContext::operator()(SubstitutableType *type) const {
-  auto *archetypeTy = cast<ArchetypeType>(type);
+  // Local archetypes can appear in interface types alongside generic param
+  // types, ignore them here.
+  auto *archetypeTy = dyn_cast<ArchetypeType>(type);
+  if (!archetypeTy)
+    return type;
 
   // Primary archetypes just map out of context.
   if (isa<PrimaryArchetypeType>(archetypeTy) ||

test/SourceKit/CursorInfo/issue-78690.swift

@@ -0,0 +1,70 @@
+// https://github.com/swiftlang/swift/issues/78690
+// rdar://143077965
+
+protocol P {
+  associatedtype R
+}
+
+func foo<each T>(_ xs: repeat each T) {
+  for x in repeat each xs {
+    // RUN: %sourcekitd-test -req=cursor -pos=%(line-1):7 %s -- %s | %FileCheck %s --check-prefix=BASIC
+    // BASIC: s:4main3fooyyxxQpRvzlF1xL_qd__vp
+    for x2 in repeat each xs {
+      // RUN: %sourcekitd-test -req=cursor -pos=%(line-1):9 %s -- %s | %FileCheck %s --check-prefix=NESTED
+      // NESTED: s:4main3fooyyxxQpRvzlF2x2L_qd__vp
+    }
+  }
+}
+
+func bar<each T, each U>(xs: repeat each T, ys: repeat each U) {
+  for x1 in repeat each xs {
+    for y1 in repeat each ys {
+      func localFn() {
+        for x2 in repeat each xs {
+          for y2 in repeat each ys {
+            let k = (x1, y1, x2, y2)
+            // RUN: %sourcekitd-test -req=cursor -pos=%(line-1):17 %s -- %s | %FileCheck %s --check-prefix=DOUBLENESTED
+            // DOUBLENESTED: s:4main3bar2xs2ysyxxQp_q_q_QptRvzRv_r0_lF7localFnL_yyRvzRv_r0_lF1kL_qd___qd0__qd1__qd2__tvp
+            // -> (A1, A2, A3, A4)
+          }
+        }
+      }
+      _ = {
+        for x2 in repeat each xs {
+          for y2 in repeat each ys {
+            let k = (x1, y1, x2, y2)
+            // RUN: %sourcekitd-test -req=cursor -pos=%(line-1):17 %s -- %s | %FileCheck %s --check-prefix=DOUBLENESTED-CLOSURE
+            // DOUBLENESTED-CLOSURE: s:4main3bar2xs2ysyxxQp_q_q_QptRvzRv_r0_lFyycfU_1kL_qd___qd0__qd1__qd2__tvp
+            // -> (A1, A2, A3, A4)
+          }
+        }
+      }
+    }
+  }
+}
+
+func baz<each T, each U, each V>(
+  ts: repeat each T, us: repeat each U, urs: repeat (each U).R, vs: repeat each V
+) where repeat each U: P, (repeat (each U, each V)): Any {
+  for t in repeat each ts {
+    func localFn() {
+      for y in repeat each us {
+        func genericLocalFn<A>(_ a: A) {
+          for (u, ur, v) in repeat (each us, each urs, each vs) {
+            let k = (a, t, y, u, ur, v)
+            // RUN: %sourcekitd-test -req=cursor -pos=%(line-1):17 %s -- %s | %FileCheck %s --check-prefix=TUPLE1
+            // TUPLE1: s:4main3baz2ts2us3urs2vsyxxQp_q_q_Qp1RQy_q_Qpq0_q_QptRvzRv_Rv0_AA1PR_q0_Rh_r1_lF7localFnL_yyRvzRv_Rv0_AaIR_q0_Rh_r1_lF012genericLocalH0L_yyqd__RvzRv_Rv0_AaIR_q0_Rh_r1__lF1kL_qd___qd0__qd1__qd2__AgaIPQyd2__qd2_0_tvp
+            // -> (A1, A2, A3, A4, A4.main.P.R, B4)
+
+            _ = {
+              let k = (a, t, y, u, ur, v)
+              // RUN: %sourcekitd-test -req=cursor -pos=%(line-1):19 %s -- %s | %FileCheck %s --check-prefix=TUPLE2
+              // TUPLE2: s:4main3baz2ts2us3urs2vsyxxQp_q_q_Qp1RQy_q_Qpq0_q_QptRvzRv_Rv0_AA1PR_q0_Rh_r1_lF7localFnL_yyRvzRv_Rv0_AaIR_q0_Rh_r1_lF012genericLocalH0L_yyqd__RvzRv_Rv0_AaIR_q0_Rh_r1__lFyycfU_1kL_qd___qd0__qd1__qd2__AgaIPQyd2__qd2_0_tvp
+              // -> (A1, A2, A3, A4, A4.main.P.R, B4)
+            }
+          }
+        }
+      }
+    }
+  }
+}