[Mangler] Handle ElementArchetypeType for decls

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

This patch adds some basic mangling logic to handle
mangling the type for a decl, which is sufficient
to fix things like cursor info. More work is needed
to support lazy variables though (now it crashes in
SILGen).

This patch doesn't handle mangling standalone
element archetypes for e.g `printTypeUSR`, we'll
either need to fixup each of the clients there, or
we'll need to rethink how we form local archetypes
to correctly handle that.

rdar://143077965

Author: hamishknight

lib/AST/ASTMangler.cpp

@@ -4038,6 +4038,18 @@ CanType ASTMangler::getDeclTypeForMangling(
     ty = ty->stripConcurrency(/*recurse=*/true, /*dropGlobalActor=*/true);
   }
 
+  // Map any local archetypes out of context.
+  if (ty->hasLocalArchetype()) {
+    CaptureInfo captureInfo;
+    auto *innerDC = decl->getInnermostDeclContext();
+    if (auto fn = AnyFunctionRef::fromDeclContext(innerDC))
+      captureInfo = fn->getCaptureInfo();
+
+    ty = mapLocalArchetypesOutOfContext(
+        ty, innerDC->getGenericSignatureOfContext(),
+        captureInfo.getGenericEnvironments());
+  }
+
   auto canTy = ty->getCanonicalType();
 
   if (auto gft = dyn_cast<GenericFunctionType>(canTy)) {

lib/AST/LocalArchetypeRequirementCollector.cpp

@@ -183,13 +183,29 @@ Type MapLocalArchetypesOutOfContext::getInterfaceType(
     ++depth;
   }
 
+  // If the environment we have is not for a capture, this must be a generic
+  // environment local to the parent function. For element archetypes, the depth
+  // we have is relative to the parent signature, so we can compute the correct
+  // depth by offsetting it by the number of parent captured environments.
+  if (genericEnv->getKind() == GenericEnvironment::Kind::OpenedElement) {
+    ASSERT(rootParam->getDepth() >= baseGenericSig.getNextDepth());
+    return GenericTypeParamType::getType(
+        rootParam->getDepth() + capturedEnvs.size(), rootParam->getIndex(),
+        rootParam->getASTContext());
+  }
+
+  // For opened archetypes, we have no way of computing the correct depth.
   llvm::errs() << "Fell off the end:\n";
   interfaceTy->dump(llvm::errs());
   abort();
 }
 
 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,60 @@
+// 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_qd0__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)
+          }
+        }
+      }
+    }
+  }
+}
+
+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)
+            }
+          }
+        }
+      }
+    }
+  }
+}