From 1c1a1cf5f63eb09becbbf9cded6fc7292636f144 Mon Sep 17 00:00:00 2001
From: Alex Hoppen <ahoppen@apple.com>
Date: Thu, 5 Dec 2024 20:24:59 -0800
Subject: [PATCH] Cached transformed results in `Cache`

The transform to get the transformed result might be expensive, so we should cache its result.
---
 .../BuildSystemManager.swift                  | 14 ++++++-----
 Sources/SwiftExtensions/Cache.swift           | 25 ++++++++++++++-----
 2 files changed, 27 insertions(+), 12 deletions(-)

diff --git a/Sources/BuildSystemIntegration/BuildSystemManager.swift b/Sources/BuildSystemIntegration/BuildSystemManager.swift
index 1ac1cd79..c8d3deb3 100644
--- a/Sources/BuildSystemIntegration/BuildSystemManager.swift
+++ b/Sources/BuildSystemIntegration/BuildSystemManager.swift
@@ -1009,19 +1009,21 @@ package actor BuildSystemManager: QueueBasedMessageHandler {
       return []
     }
 
+    let request = BuildTargetSourcesRequest(targets: targets.sorted { $0.uri.stringValue < $1.uri.stringValue })
+
     // If we have a cached request for a superset of the targets, serve the result from that cache entry.
     let fromSuperset = await orLog("Getting source files from superset request") {
-      try await cachedTargetSources.get(isolation: self) { request in
-        targets.isSubset(of: request.targets)
-      } transform: { response in
-        return BuildTargetSourcesResponse(items: response.items.filter { targets.contains($0.target) })
-      }
+      try await cachedTargetSources.getDerived(
+        isolation: self,
+        request,
+        canReuseKey: { targets.isSubset(of: $0.targets) },
+        transform: { BuildTargetSourcesResponse(items: $0.items.filter { targets.contains($0.target) }) }
+      )
     }
     if let fromSuperset {
       return fromSuperset.items
     }
 
-    let request = BuildTargetSourcesRequest(targets: targets.sorted { $0.uri.stringValue < $1.uri.stringValue })
     let response = try await cachedTargetSources.get(request, isolation: self) { request in
       try await buildSystemAdapter.send(request)
     }
diff --git a/Sources/SwiftExtensions/Cache.swift b/Sources/SwiftExtensions/Cache.swift
index 7642b20d..debb7a02 100644
--- a/Sources/SwiftExtensions/Cache.swift
+++ b/Sources/SwiftExtensions/Cache.swift
@@ -33,15 +33,28 @@ package class Cache<Key: Sendable & Hashable, Result: Sendable> {
     return try await task.value
   }
 
-  package func get(
+  /// Get the value cached for `key`. If no value exists for `key`, try deriving the result from an existing cache entry
+  /// that satisfies `canReuseKey` by applying `transform` to that result.
+  package func getDerived(
     isolation: isolated any Actor,
-    whereKey keyPredicate: (Key) -> Bool,
-    transform: @Sendable @escaping (Result) -> Result
+    _ key: Key,
+    canReuseKey: @Sendable @escaping (Key) -> Bool,
+    transform: @Sendable @escaping (_ cachedResult: Result) -> Result
   ) async throws -> Result? {
-    for (key, value) in storage {
-      if keyPredicate(key) {
-        return try await transform(value.value)
+    if let cached = storage[key] {
+      // If we have a value for the requested key, prefer that
+      return try await cached.value
+    }
+
+    // See if don't have an entry for this key, see if we can derive the value from a cached entry.
+    for (cachedKey, cachedValue) in storage {
+      guard canReuseKey(cachedKey) else {
+        continue
       }
+      let transformed = Task { try await transform(cachedValue.value) }
+      // Cache the transformed result.
+      storage[key] = transformed
+      return try await transformed.value
     }
     return nil
   }