From c97dfd6b05441ad11c91632f5de890af8b4f5b07 Mon Sep 17 00:00:00 2001
From: Augusto Noronha <anoronha@apple.com>
Date: Fri, 11 Jul 2025 09:23:45 -0700
Subject: [PATCH] [RemoteAddress] Handle comparison of addresses in different
 spaces

Sometimes it makes sense to compares addresses from different address
spaces.

rdar://148361743
---
 include/swift/Remote/RemoteAddress.h          | 26 +++++++++++++++++++
 .../RemoteInspection/ReflectionContext.h      |  2 +-
 .../swift/RemoteInspection/TypeRefBuilder.h   |  4 +++
 .../RemoteInspection/TypeRefBuilder.cpp       | 10 ++++---
 4 files changed, 37 insertions(+), 5 deletions(-)

diff --git a/include/swift/Remote/RemoteAddress.h b/include/swift/Remote/RemoteAddress.h
index 5e75cfe0ba9d7..91313cbd31b92 100644
--- a/include/swift/Remote/RemoteAddress.h
+++ b/include/swift/Remote/RemoteAddress.h
@@ -55,18 +55,44 @@ class RemoteAddress {
     return !operator==(other);
   }
 
+  bool inRange(const RemoteAddress &begin, const RemoteAddress &end) const {
+    assert(begin.AddressSpace != end.AddressSpace &&
+           "Unexpected address spaces");
+    if (AddressSpace != begin.AddressSpace)
+      return false;
+    return begin <= *this && *this < end;
+  }
+
   bool operator<(const RemoteAddress rhs) const {
     assert(AddressSpace == rhs.AddressSpace &&
            "Comparing remote addresses of different address spaces");
     return Data < rhs.Data;
   }
 
+  /// Less than operator to be used for ordering purposes. The default less than
+  /// operator asserts if the address spaces are different, this one takes it
+  /// into account to determine the order of the addresses.
+  bool orderedLessThan(const RemoteAddress rhs) const {
+    if (AddressSpace == rhs.AddressSpace)
+      return Data < rhs.Data;
+    return AddressSpace < rhs.AddressSpace;
+  }
+
   bool operator<=(const RemoteAddress rhs) const {
     assert(AddressSpace == rhs.AddressSpace &&
            "Comparing remote addresses of different address spaces");
     return Data <= rhs.Data;
   }
 
+  /// Less than or equal operator to be used for ordering purposes. The default
+  /// less than or equal operator asserts if the address spaces are different,
+  /// this one takes it into account to determine the order of the addresses.
+  bool orderedLessThanOrEqual(const RemoteAddress rhs) const {
+    if (AddressSpace == rhs.AddressSpace)
+      return Data <= rhs.Data;
+    return AddressSpace <= rhs.AddressSpace;
+  }
+
   bool operator>(const RemoteAddress &rhs) const {
     assert(AddressSpace == rhs.AddressSpace &&
            "Comparing remote addresses of different address spaces");
diff --git a/include/swift/RemoteInspection/ReflectionContext.h b/include/swift/RemoteInspection/ReflectionContext.h
index dd9ceb3de8b53..45d1ecee7c0c6 100644
--- a/include/swift/RemoteInspection/ReflectionContext.h
+++ b/include/swift/RemoteInspection/ReflectionContext.h
@@ -907,7 +907,7 @@ class ReflectionContext
     for (auto Range : ranges) {
       auto Start = std::get<0>(Range);
       auto End = std::get<1>(Range);
-      if (Start <= Address && Address < End)
+      if (Address.inRange(Start, End))
         return true;
     }
 
diff --git a/include/swift/RemoteInspection/TypeRefBuilder.h b/include/swift/RemoteInspection/TypeRefBuilder.h
index 8d46283855588..f2174133cdcd9 100644
--- a/include/swift/RemoteInspection/TypeRefBuilder.h
+++ b/include/swift/RemoteInspection/TypeRefBuilder.h
@@ -65,6 +65,10 @@ class ReflectionSection {
 
   bool containsRemoteAddress(remote::RemoteAddress remoteAddr,
                              uint64_t size) const {
+    if (Start.getRemoteAddress().getAddressSpace() !=
+        remoteAddr.getAddressSpace())
+      return false;
+
     return Start.getRemoteAddress() <= remoteAddr &&
            remoteAddr + size <= Start.getRemoteAddress() + Size;
   }
diff --git a/stdlib/public/RemoteInspection/TypeRefBuilder.cpp b/stdlib/public/RemoteInspection/TypeRefBuilder.cpp
index 781599b454bca..fb214d05f5025 100644
--- a/stdlib/public/RemoteInspection/TypeRefBuilder.cpp
+++ b/stdlib/public/RemoteInspection/TypeRefBuilder.cpp
@@ -64,7 +64,7 @@ TypeRefBuilder::ReflectionTypeDescriptorFinder::
                   .TypeReference.startAddress()
                   .getRemoteAddress();
 
-          return typeReferenceAStart < typeReferenceBStart;
+          return typeReferenceAStart.orderedLessThan(typeReferenceBStart);
         });
   }
 
@@ -75,9 +75,11 @@ TypeRefBuilder::ReflectionTypeDescriptorFinder::
       ReflectionInfoIndexesSortedByTypeReferenceRange.begin(),
       ReflectionInfoIndexesSortedByTypeReferenceRange.end(), remoteAddr,
       [&](uint32_t ReflectionInfoIndex, remote::RemoteAddress remoteAddr) {
-        return ReflectionInfos[ReflectionInfoIndex]
-                   .TypeReference.endAddress()
-                   .getRemoteAddress() <= remoteAddr;
+        auto reflectionInfoAddress = ReflectionInfos[ReflectionInfoIndex]
+                                         .TypeReference.endAddress()
+                                         .getRemoteAddress();
+
+        return reflectionInfoAddress.orderedLessThanOrEqual(remoteAddr);
       });
 
   if (possiblyMatchingReflectionInfoIndex ==