[Mangler] Handle local archetypes in getDeclTypeForMangling

Element archetypes can occur here when mangling the USR for local
variables for e.g SourceKit cursor info, as well as for regular
compilation for things like lazy variables.

Update `getDeclTypeForMangling` to map local archetypes out of
context, using both the captured generic environments and the
archetypes present in the type. More work is needed to support lazy
variable though (now it crashes in SILGen).

This patch doesn't handle mangling standalone element archetypes for
e.g `printTypeUSR`, ideally we'd fix the clients there to not pass
local archetypes.

rdar://143077965

Author: hamishknight

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)
+            }
+          }
+        }
+      }
+    }
+  }
+}