[SharedCache] Make it faster to look up which image contains an address

`MemoryRegion` now tracks the start address of the image to which it
belongs. `SharedCache::HeaderForAddress` uses the existing address range
map to quickly find the `MemoryRegion` for a given address, and from
there can look up the image via its start address.

Some extra logic is added to `CacheInfo::Load` to populate the image
start address on `MemoryRegion` when loading from metadata that predates
this change.

The result is that `HeaderForAddress` is no longer the most expensive
part of `FindSymbolAtAddrAndApplyToAddr`.

2 files changed, 44 insertions, 22 deletions

diff --git a/view/sharedcache/core/SharedCache.cpp b/view/sharedcache/core/SharedCache.cpp
index afbd3b58..a5dd65d8 100644
--- a/view/sharedcache/core/SharedCache.cpp
+++ b/view/sharedcache/core/SharedCache.cpp
@@ -1015,6 +1015,7 @@ void SharedCache::PerformInitialLoad(std::lock_guard<std::mutex>& lock)
 				sectionRegion.prettyName = imageHeader.value().identifierPrefix + "::" + std::string(segName);
 				sectionRegion.start = segment.vmaddr;
 				sectionRegion.size = segment.vmsize;
+				sectionRegion.imageStart = start.second;
 				uint32_t flags = SegmentFlagsFromMachOProtections(segment.initprot, segment.maxprot);
 
 				// if we're positive we have an entry point for some reason, force the segment
@@ -1529,18 +1530,16 @@ const SharedCacheMachOHeader* SharedCache::HeaderForAddress(uint64_t address)
 	if (auto it = m_cacheInfo->headers.find(address); it != m_cacheInfo->headers.end())
 		return &it->second;
 
-	// We _could_ mark each page with the image start? :grimacing emoji:
-	// But that'd require mapping pages :grimacing emoji: :grimacing emoji:
-	// There's not really any other hacks that could make this faster, that I can think of...
-	for (const auto& [start, header] : m_cacheInfo->headers)
-	{
-		for (const auto& segment : header.segments)
-		{
-			if (segment.vmaddr <= address && segment.vmaddr + segment.vmsize > address)
-				return &header;
-		}
-	}
-	return {};
+	auto it = m_cacheInfo->memoryRegions.find(address);
+	if (it == m_cacheInfo->memoryRegions.end())
+		return nullptr;
+
+	if (auto headerAddress = it->second.imageStart)
+		return HeaderForAddress(headerAddress);
+
+	// Found a region, but its `imageStart` was 0. This should mean it doesn't belong to an image.
+	assert(it->second.type != MemoryRegion::Type::Image);
+	return nullptr;
 }
 
 std::string SharedCache::NameForAddress(uint64_t address)
@@ -1561,16 +1560,9 @@ std::string SharedCache::ImageNameForAddress(uint64_t address)
 
 bool SharedCache::LoadImageContainingAddress(uint64_t address, bool skipObjC)
 {
-	for (const auto& [start, header] : m_cacheInfo->headers)
-	{
-		for (const auto& segment : header.segments)
-		{
-			if (segment.vmaddr <= address && segment.vmaddr + segment.vmsize > address)
-			{
-				std::lock_guard lock(m_mutex);
-				return LoadImageWithInstallName(lock, header.installName, skipObjC);
-			}
-		}
+	if (auto header = HeaderForAddress(address)) {
+		std::lock_guard lock(m_mutex);
+		return LoadImageWithInstallName(lock, header->installName, skipObjC);
 	}
 
 	return false;
@@ -3387,6 +3379,29 @@ std::optional<SharedCache::CacheInfo> SharedCache::CacheInfo::Load(Deserializati
 	cacheInfo.MSL(objcOptimizationDataRange);
 	cacheInfo.MSL(baseFilePath);
 	cacheInfo.MSL_CAST(cacheFormat, uint8_t, SharedCacheFormat);
+
+	// Older metadata may be missing the `imageStart` field on `MemoryRegion`.
+	bool regionsMissingImageStart =
+		std::any_of(cacheInfo.memoryRegions.begin(), cacheInfo.memoryRegions.end(), [](const auto& pair) {
+			const auto& region = pair.second;
+			return region.type == MemoryRegion::Type::Image && region.imageStart == 0;
+		});
+
+	if (regionsMissingImageStart)
+	{
+		for (const auto& [start, header] : cacheInfo.headers)
+		{
+			for (const auto& segment : header.segments)
+			{
+				const auto regionIt = cacheInfo.memoryRegions.find(segment.vmaddr);
+				assert(regionIt != cacheInfo.memoryRegions.end());
+				auto& region = regionIt->second;
+				assert(!region.imageStart || region.imageStart == start);
+				region.imageStart = start;
+			}
+		}
+	}
+
 	return cacheInfo;
 }
 void State::Store(SerializationContext& context, std::optional<DSCViewState> viewState) const
diff --git a/view/sharedcache/core/SharedCache.h b/view/sharedcache/core/SharedCache.h
index 97da0880..d696154c 100644
--- a/view/sharedcache/core/SharedCache.h
+++ b/view/sharedcache/core/SharedCache.h
@@ -41,6 +41,9 @@ namespace SharedCacheCore {
 		std::string prettyName;
 		uint64_t start;
 		uint64_t size;
+		// Start address of the image this region belongs to.
+		// 0 if the region does not belong to any image.
+		uint64_t imageStart = 0;
 		BNSegmentFlag flags;
 		Type type;
 
@@ -55,6 +58,7 @@ namespace SharedCacheCore {
 			MSS(prettyName);
 			MSS(start);
 			MSS(size);
+			MSS(imageStart);
 			MSS_CAST(flags, uint64_t);
 			MSS_CAST(type, uint8_t);
 		}
@@ -67,6 +71,9 @@ namespace SharedCacheCore {
 			region.MSL(size);
 			region.MSL_CAST(flags, uint64_t, BNSegmentFlag);
 			region.MSL_CAST(type, uint8_t, Type);
+			if (context.doc.HasMember("imageStart"))
+				region.MSL(imageStart);
+
 			return region;
 		}
 	};