From 9e07fac65038925023c43e67901fa2de4c732b3f Mon Sep 17 00:00:00 2001
From: Googler <jcater@google.com>
Date: Wed, 8 Nov 2023 02:07:18 -0800
Subject: [PATCH] Streamline code to cache configurations.

Work towards platform-based flags: #19409.

PiperOrigin-RevId: 580452068
Change-Id: If94e0fbd9a815bfe3decca21ef1fd58003b4c366
---
 .../cquery/ProtoOutputFormatterCallback.java  | 22 +++++++++++++++----
 1 file changed, 18 insertions(+), 4 deletions(-)

diff --git a/src/main/java/com/google/devtools/build/lib/query2/cquery/ProtoOutputFormatterCallback.java b/src/main/java/com/google/devtools/build/lib/query2/cquery/ProtoOutputFormatterCallback.java
index c9b63cd4881147..4d1b1ed927566d 100644
--- a/src/main/java/com/google/devtools/build/lib/query2/cquery/ProtoOutputFormatterCallback.java
+++ b/src/main/java/com/google/devtools/build/lib/query2/cquery/ProtoOutputFormatterCallback.java
@@ -25,6 +25,7 @@
 import com.google.devtools.build.lib.analysis.AnalysisProtosV2.CqueryResult;
 import com.google.devtools.build.lib.analysis.AnalysisProtosV2.CqueryResultOrBuilder;
 import com.google.devtools.build.lib.analysis.ConfiguredTarget;
+import com.google.devtools.build.lib.analysis.config.BuildConfigurationValue;
 import com.google.devtools.build.lib.analysis.config.BuildOptions;
 import com.google.devtools.build.lib.analysis.config.ConfigMatchingProvider;
 import com.google.devtools.build.lib.buildeventstream.BuildEventStreamProtos;
@@ -56,6 +57,7 @@
 import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
+import java.util.function.Function;
 
 /** Proto output formatter for cquery results. */
 class ProtoOutputFormatterCallback extends CqueryThreadsafeCallback {
@@ -80,11 +82,24 @@ public String formatName() {
 
   private static class ConfigurationCache {
     private final Map<BuildConfigurationEvent, Integer> cache = new HashMap<>();
+    private final Function<BuildConfigurationKey, BuildConfigurationValue> configurationGetter;
+
+    private ConfigurationCache(
+        Function<BuildConfigurationKey, BuildConfigurationValue> configurationGetter) {
+      this.configurationGetter = configurationGetter;
+    }
 
     public int getId(BuildConfigurationEvent buildConfigurationEvent) {
       return cache.computeIfAbsent(buildConfigurationEvent, event -> cache.size() + 1);
     }
 
+    public int getId(BuildOptions options) {
+      BuildConfigurationValue configurationValue =
+          configurationGetter.apply(BuildConfigurationKey.create(options));
+      BuildConfigurationEvent buildConfigurationEvent = configurationValue.toBuildEvent();
+      return getId(buildConfigurationEvent);
+    }
+
     public ImmutableList<Configuration> getConfigurations() {
       return cache.entrySet().stream()
           .map(
@@ -110,7 +125,8 @@ public ImmutableList<Configuration> getConfigurations() {
   private final OutputType outputType;
   private final AspectResolver resolver;
   private final SkyframeExecutor skyframeExecutor;
-  private final ConfigurationCache configurationCache = new ConfigurationCache();
+  private final ConfigurationCache configurationCache =
+      new ConfigurationCache(this::getConfiguration);
   private final JsonFormat.Printer jsonPrinter = JsonFormat.printer();
   private final RuleClassProvider ruleClassProvider;
 
@@ -261,9 +277,7 @@ public void processOutput(Iterable<ConfiguredTarget> partialResult)
                           .setLabel(labelPrinter.toString(resolvedTransition.label())));
             } else {
               for (BuildOptions options : resolvedTransition.options()) {
-                BuildConfigurationEvent buildConfigurationEvent =
-                    getConfiguration(BuildConfigurationKey.create(options)).toBuildEvent();
-                int configurationId = configurationCache.getId(buildConfigurationEvent);
+                int configurationId = configurationCache.getId(options);
 
                 targetBuilder
                     .getRuleBuilder()