From 1457fe65804d40180ad7f0d4ae164239435823c9 Mon Sep 17 00:00:00 2001
From: Anton Korobeynikov <anton@korobeynikov.info>
Date: Wed, 27 Mar 2024 12:12:47 -0700
Subject: [PATCH] Ensure LoadableByAddress rewrites pointer_to_address (#72619)

if the target address is address of tuple type that should be rewritten.

Fixes #72363

(cherry picked from commit 2ad4a60619914ef35462e7ce80c3b8fb8d0fd9db)
---
 lib/IRGen/LoadableByAddress.cpp               |  3 +-
 .../loadable_by_address_issue72363.swift      | 31 +++++++++++++++++++
 2 files changed, 33 insertions(+), 1 deletion(-)
 create mode 100644 test/IRGen/loadable_by_address_issue72363.swift

diff --git a/lib/IRGen/LoadableByAddress.cpp b/lib/IRGen/LoadableByAddress.cpp
index f7c24f0ff230a..54a6f42f91cb3 100644
--- a/lib/IRGen/LoadableByAddress.cpp
+++ b/lib/IRGen/LoadableByAddress.cpp
@@ -1003,7 +1003,8 @@ void LargeValueVisitor::visitAllocStackInst(AllocStackInst *instr) {
 
 void LargeValueVisitor::visitPointerToAddressInst(PointerToAddressInst *instr) {
   SILType currSILType = instr->getType().getObjectType();
-  if (getInnerFunctionType(currSILType)) {
+  if (pass.containsDifferentFunctionSignature(pass.F->getLoweredFunctionType(),
+                                              currSILType)) {
     pass.pointerToAddrkInstsToMod.push_back(instr);
   }
 }
diff --git a/test/IRGen/loadable_by_address_issue72363.swift b/test/IRGen/loadable_by_address_issue72363.swift
new file mode 100644
index 0000000000000..18b7a5aa35f9a
--- /dev/null
+++ b/test/IRGen/loadable_by_address_issue72363.swift
@@ -0,0 +1,31 @@
+// REQUIRES: differentiable_programming
+// RUN: %target-swift-frontend -emit-ir -verify %s
+
+// This used to trigger an assertion due to LoadableByAddress incorrectly doing incorrect transforms
+// around linear map tuple users (pointer_to_address).
+
+import _Differentiation
+struct H: Differentiable {}
+protocol J: Differentiable {}
+struct L: Differentiable {
+    var p: [P]
+    @differentiable(reverse) func s() -> H {
+        var m = 0.0
+        for i in 0 ..< withoutDerivative(at: p.count) {
+            m += p[i].a
+            m += p[i].a
+            m += p[i].a
+            m += p[i].a
+        }
+        return P.g(p: P(a: 0.0, b: 0.0, c: 0.0, d: m), z: L( p: self.p)).w
+    }
+}
+struct P: J {
+    var a = 0.0
+    var b = 0.0
+    var c = 0.0
+    var d = 0.0
+    var e = 0.0
+    @differentiable(reverse) static func g(p: P, z: L) -> Y<P> {return Y<P>(w: H())}
+}
+struct Y<U: J>: Differentiable {var w: H = H()}