From 7a6070e79c6fa9b700eead337920ce829433fb75 Mon Sep 17 00:00:00 2001
From: Shubham Sandeep Rastogi <srastogi22@apple.com>
Date: Thu, 6 Mar 2025 17:22:27 -0800
Subject: [PATCH] [DebugInfo] Use Stmt EndLoc if SILLocation passed in is the
 Stmt EndLoc.

When creating an ExtendedASTNodeLoc from a SILLocation, if the
SILLocation passed in belongs to a swift::Stmt, we only ever use the
Stmt's StartLoc for the SourceLocation. If the SILLocation passed in,
has a SourceLocation that matches the EndLoc of the Stmt, we should
correctly set the primary ASTNodeTy PointerUnion's integer to 1, to
denote that the SourceLocation dervied from the Stmt points to the
EndLoc.
---
 lib/SIL/IR/SILLocation.cpp           |  8 +++++++-
 test/DebugInfo/hop_to_executor.swift | 24 ++++++++++++++++++++++++
 2 files changed, 31 insertions(+), 1 deletion(-)
 create mode 100644 test/DebugInfo/hop_to_executor.swift

diff --git a/lib/SIL/IR/SILLocation.cpp b/lib/SIL/IR/SILLocation.cpp
index 794870ab2910b..ea61a144ad37b 100644
--- a/lib/SIL/IR/SILLocation.cpp
+++ b/lib/SIL/IR/SILLocation.cpp
@@ -290,8 +290,14 @@ RegularLocation::getDebugOnlyExtendedASTNodeLoc(SILLocation L,
     return new (Module) ExtendedASTNodeLoc(Empty, {D, 0});
   if (auto E = L.getAsASTNode<Expr>())
     return new (Module) ExtendedASTNodeLoc(Empty, {E, 0});
-  if (auto S = L.getAsASTNode<Stmt>())
+  if (auto S = L.getAsASTNode<Stmt>()) {
+    // If the source location of the SILLocation passed in matches the EndLoc of
+    // the Stmt, set the primary ASTNodeTy integer to 1, so that
+    // SILLocation::getSourceLoc returns the EndLoc when queried.
+    if (L.getSourceLocForDebugging() == S->getEndLoc())
+      Empty.setInt(1);
     return new (Module) ExtendedASTNodeLoc(Empty, {S, 0});
+  }
   auto P = L.getAsASTNode<Pattern>();
   return new (Module) ExtendedASTNodeLoc(Empty, {P, 0});
 }
diff --git a/test/DebugInfo/hop_to_executor.swift b/test/DebugInfo/hop_to_executor.swift
new file mode 100644
index 0000000000000..840193409a5ac
--- /dev/null
+++ b/test/DebugInfo/hop_to_executor.swift
@@ -0,0 +1,24 @@
+// RUN: %target-swiftc_driver %s -c -g -Onone -o - -Xllvm -sil-print-debuginfo -emit-sil -parse-as-library -module-name m |  %FileCheck %s
+
+// This test ensures that the hop_to_executor source location matches the end of the do block
+
+func getTimestamp(x: Int) async -> Int {
+  return 40 + x
+}
+func work() {}
+func foo() async {
+  do {
+    work()
+    async let timestamp2 = getTimestamp(x:2)
+    print(await timestamp2)
+  // CHECK: %[[REG:[0-9]+]] = function_ref @swift_asyncLet_finish : $@convention(thin) @async (Builtin.RawPointer, Builtin.RawPointer) -> (), loc {{.*}}:[[@LINE+3]]
+  // CHECK-NEXT: %{{[0-9]+}} = apply %[[REG]](%{{[0-9]+}}, %{{[0-9]+}}) : $@convention(thin) @async (Builtin.RawPointer, Builtin.RawPointer) -> (), loc{{.*}}:[[@LINE+2]]
+  // CHECK-NEXT: hop_to_executor %0, loc * {{.*}}:[[@LINE+1]]
+  }
+  work()
+}
+@main enum entry {
+  static func main() async {
+    await foo()
+  }
+}