Fixes crash while instrumenting generic func body embedded in another func.

The crash was caused by attempting to add logging expressions that capture generic variables while using the outer func decl context that was not the generic decl context of the inner (generic) func.

The fix "pushes" the current func decl context into body instrumenting, then "popping" the context on return. Rather than always using the top-level func decl context for all nested func bodies.

Author: chrismiles

lib/Sema/InstrumenterSupport.cpp

@@ -142,3 +142,18 @@ Expr *InstrumenterBase::buildIDArgumentExpr(std::optional<DeclNameRef> name,
   return new (Context) UnresolvedDeclRefExpr(*name, DeclRefKind::Ordinary,
                                              DeclNameLoc(SR.End));
 }
+
+BraceStmt *InstrumenterBase::transformBraceStmtWithLocalDeclContext(
+    BraceStmt *BS, DeclContext *localDC, const ParameterList *PL,
+    bool TopLevel) {
+  // Override TypeCheckDC for the scope of the brace stmt
+  DeclContext *outerTypeCheckDC = TypeCheckDC;
+  TypeCheckDC = localDC;
+
+  BraceStmt *transformedBS = transformBraceStmt(BS, PL, TopLevel);
+
+  // Restore TypeCheckDC
+  TypeCheckDC = outerTypeCheckDC;
+
+  return transformedBS;
+}

lib/Sema/InstrumenterSupport.h

@@ -52,6 +52,11 @@ class InstrumenterBase {
                                         const ParameterList *PL = nullptr,
                                         bool TopLevel = false) = 0;
 
+  BraceStmt *
+  transformBraceStmtWithLocalDeclContext(BraceStmt *BS, DeclContext *localDC,
+                                         const ParameterList *PL = nullptr,
+                                         bool TopLevel = false);
+
   /// Create an expression which retrieves a valid ModuleIdentifier or
   /// FileIdentifier, if available.
   Expr *buildIDArgumentExpr(std::optional<DeclNameRef> name, SourceRange SR);

lib/Sema/PCMacro.cpp

@@ -339,7 +339,8 @@ class Instrumenter : InstrumenterBase {
     if (auto *FD = dyn_cast<FuncDecl>(D)) {
       if (BraceStmt *B = FD->getTypecheckedBody()) {
         const ParameterList *PL = FD->getParameters();
-        BraceStmt *NB = transformBraceStmt(B, PL);
+        BraceStmt *NB =
+            transformBraceStmtWithLocalDeclContext(B, /*localDC =*/FD, PL);
         // Since it would look strange going straight to the first line in a
         // function body, we throw in a before/after pointing at the function
         // decl at the start of the transformed body

lib/Sema/PlaygroundTransform.cpp

@@ -312,7 +312,8 @@ class Instrumenter : InstrumenterBase {
       if (BraceStmt *B = FD->getTypecheckedBody()) {
         const ParameterList *PL = FD->getParameters();
         TargetKindSetter TKS(BracePairs, BracePair::TargetKinds::Return);
-        BraceStmt *NB = transformBraceStmt(B, PL);
+        BraceStmt *NB =
+            transformBraceStmtWithLocalDeclContext(B, /*localDC =*/FD, PL);
         if (NB != B) {
           FD->setBody(NB, AbstractFunctionDecl::BodyKind::TypeChecked);
           TypeChecker::checkFunctionEffects(FD);

test/PlaygroundTransform/generic_func_in_local_scope.swift

@@ -0,0 +1,33 @@
+// RUN: %empty-directory(%t)
+
+// Build PlaygroundSupport module
+// RUN: %target-build-swift -whole-module-optimization -module-name PlaygroundSupport -emit-module-path %t/PlaygroundSupport.swiftmodule -parse-as-library -c -o %t/PlaygroundSupport.o %S/Inputs/SilentPCMacroRuntime.swift %S/Inputs/PlaygroundsRuntime.swift
+
+// RUN: %target-build-swift -swift-version 5 -emit-module -Xfrontend -playground -I=%t %s
+// RUN: %target-build-swift -swift-version 6 -emit-module -Xfrontend -playground -I=%t %s
+
+// RUN: %target-build-swift -swift-version 5 -emit-module -Xfrontend -playground -Xfrontend -pc-macro -I=%t %s
+// RUN: %target-build-swift -swift-version 6 -emit-module -Xfrontend -playground -Xfrontend -pc-macro -I=%t %s
+
+// REQUIRES: executable_test
+
+import PlaygroundSupport
+
+func test1() {
+    func buildBlock<Content>(content: Content) {
+        content
+    }
+}
+
+func test2() {
+    func buildBlock<Content>(content: Content) -> Content {
+        content
+        return content
+    }
+}
+
+func test3<Content>(_ content: Content) {
+    func buildBlock<Content2>(_ content2: Content2) -> Content2 {
+       return content2
+    }
+}