From e106a2de659bf44b76124fe55e7e9140c86b55fa Mon Sep 17 00:00:00 2001
From: Fabian Meumertzheim <fabian@meumertzhe.im>
Date: Sun, 18 Dec 2022 08:42:02 +0100
Subject: [PATCH] Prettify labels in action progress messages with Bzlmod

`UiStateTracker` is provided with the repository mapping of the main
repository after the loading phase has been completed and uses this
mapping to "unmap" canonical labels back to the apparent name used for
them by the main repository.
---
 .../build/lib/actions/AbstractAction.java     | 25 +++++++++++++++--
 .../lib/actions/ActionExecutionMetadata.java  | 14 ++++++++++
 .../pkgcache/LoadingPhaseCompleteEvent.java   | 13 ++++++++-
 .../lib/runtime/SkymeldUiStateTracker.java    |  1 +
 .../build/lib/runtime/UiStateTracker.java     | 12 +++++---
 .../skyframe/TargetPatternPhaseFunction.java  |  3 +-
 .../runtime/SkymeldUiStateTrackerTest.java    |  4 ++-
 .../build/lib/runtime/UiStateTrackerTest.java | 28 +++++++++++++++----
 8 files changed, 85 insertions(+), 15 deletions(-)

diff --git a/src/main/java/com/google/devtools/build/lib/actions/AbstractAction.java b/src/main/java/com/google/devtools/build/lib/actions/AbstractAction.java
index ea4c7889302de2..7d6fd6f920ffad 100644
--- a/src/main/java/com/google/devtools/build/lib/actions/AbstractAction.java
+++ b/src/main/java/com/google/devtools/build/lib/actions/AbstractAction.java
@@ -27,6 +27,7 @@
 import com.google.devtools.build.lib.actions.extra.ExtraActionInfo;
 import com.google.devtools.build.lib.analysis.platform.PlatformInfo;
 import com.google.devtools.build.lib.cmdline.Label;
+import com.google.devtools.build.lib.cmdline.RepositoryMapping;
 import com.google.devtools.build.lib.collect.nestedset.Depset;
 import com.google.devtools.build.lib.collect.nestedset.NestedSet;
 import com.google.devtools.build.lib.collect.nestedset.NestedSetBuilder;
@@ -363,18 +364,36 @@ public boolean showsOutputUnconditionally() {
   @Nullable
   @Override
   public final String getProgressMessage() {
+    return getProgressMessageChecked(null);
+  }
+
+  @Nullable
+  @Override
+  public final String getProgressMessage(RepositoryMapping mainRepositoryMapping) {
+    Preconditions.checkNotNull(mainRepositoryMapping);
+    return getProgressMessageChecked(mainRepositoryMapping);
+  }
+
+  private String getProgressMessageChecked(@Nullable RepositoryMapping mainRepositoryMapping) {
     String message = getRawProgressMessage();
     if (message == null) {
       return null;
     }
-    message = replaceProgressMessagePlaceholders(message);
+    message = replaceProgressMessagePlaceholders(message, mainRepositoryMapping);
     String additionalInfo = getOwner().getAdditionalProgressInfo();
     return additionalInfo == null ? message : message + " [" + additionalInfo + "]";
   }
 
-  private String replaceProgressMessagePlaceholders(String progressMessage) {
+  private String replaceProgressMessagePlaceholders(String progressMessage,
+      @Nullable RepositoryMapping mainRepositoryMapping) {
     if (progressMessage.contains("%{label}") && getOwner().getLabel() != null) {
-      progressMessage = progressMessage.replace("%{label}", getOwner().getLabel().toString());
+      String labelString;
+      if (mainRepositoryMapping != null) {
+        labelString = getOwner().getLabel().getDisplayForm(mainRepositoryMapping);
+      } else {
+        labelString = getOwner().getLabel().toString();
+      }
+      progressMessage = progressMessage.replace("%{label}", labelString);
     }
     if (progressMessage.contains("%{output}") && getPrimaryOutput() != null) {
       progressMessage =
diff --git a/src/main/java/com/google/devtools/build/lib/actions/ActionExecutionMetadata.java b/src/main/java/com/google/devtools/build/lib/actions/ActionExecutionMetadata.java
index 338fb06993fdbc..0b98b2c7274a47 100644
--- a/src/main/java/com/google/devtools/build/lib/actions/ActionExecutionMetadata.java
+++ b/src/main/java/com/google/devtools/build/lib/actions/ActionExecutionMetadata.java
@@ -13,6 +13,7 @@
 // limitations under the License.
 package com.google.devtools.build.lib.actions;
 
+import com.google.devtools.build.lib.cmdline.RepositoryMapping;
 import com.google.devtools.build.lib.concurrent.ThreadSafety.ThreadSafe;
 import javax.annotation.Nullable;
 
@@ -33,6 +34,19 @@ public interface ActionExecutionMetadata extends ActionAnalysisMetadata {
   @Nullable
   String getProgressMessage();
 
+  /**
+   * A variant of {@link #getProgressMessage} that additionally takes the {@link RepositoryMapping}
+   * of the main repository, which can be used by the implementation to emit labels with apparent
+   * instead of canonical repository names. A return value of {@code null} indicates no
+   * message should be reported.
+   *
+   * <p>The default implementation simply returns the result of {@link #getProgressMessage}.
+   */
+  @Nullable
+  default String getProgressMessage(RepositoryMapping mainRepositoryMapping) {
+    return getProgressMessage();
+  }
+
   /**
    * Returns a human-readable description of the inputs to {@link #getKey(ActionKeyContext)}. Used
    * in the output from '--explain', and in error messages for '--check_up_to_date' and
diff --git a/src/main/java/com/google/devtools/build/lib/pkgcache/LoadingPhaseCompleteEvent.java b/src/main/java/com/google/devtools/build/lib/pkgcache/LoadingPhaseCompleteEvent.java
index aad9f721f1ba05..f40f7a2ceb5aae 100644
--- a/src/main/java/com/google/devtools/build/lib/pkgcache/LoadingPhaseCompleteEvent.java
+++ b/src/main/java/com/google/devtools/build/lib/pkgcache/LoadingPhaseCompleteEvent.java
@@ -16,6 +16,7 @@
 import com.google.common.base.Preconditions;
 import com.google.common.collect.ImmutableSet;
 import com.google.devtools.build.lib.cmdline.Label;
+import com.google.devtools.build.lib.cmdline.RepositoryMapping;
 import com.google.devtools.build.lib.events.ExtendedEventHandler;
 
 /**
@@ -24,6 +25,7 @@
 public final class LoadingPhaseCompleteEvent implements ExtendedEventHandler.Postable {
   private final ImmutableSet<Label> labels;
   private final ImmutableSet<Label> filteredLabels;
+  private final RepositoryMapping mainRepositoryMapping;
 
   /**
    * Construct the event.
@@ -33,9 +35,11 @@ public final class LoadingPhaseCompleteEvent implements ExtendedEventHandler.Pos
    */
   public LoadingPhaseCompleteEvent(
       ImmutableSet<Label> labels,
-      ImmutableSet<Label> filteredLabels) {
+      ImmutableSet<Label> filteredLabels,
+      RepositoryMapping mainRepositoryMapping) {
     this.labels = Preconditions.checkNotNull(labels);
     this.filteredLabels = Preconditions.checkNotNull(filteredLabels);
+    this.mainRepositoryMapping = Preconditions.checkNotNull(mainRepositoryMapping);
   }
 
   /**
@@ -53,6 +57,13 @@ public ImmutableSet<Label> getFilteredLabels() {
     return filteredLabels;
   }
 
+  /**
+   * @return The repository mapping of the main repository.
+   */
+  public RepositoryMapping getMainRepositoryMapping() {
+    return mainRepositoryMapping;
+  }
+
   @Override
   public boolean storeForReplay() {
     return true;
diff --git a/src/main/java/com/google/devtools/build/lib/runtime/SkymeldUiStateTracker.java b/src/main/java/com/google/devtools/build/lib/runtime/SkymeldUiStateTracker.java
index bc74b3f694bcab..cb2f1b687cf77d 100644
--- a/src/main/java/com/google/devtools/build/lib/runtime/SkymeldUiStateTracker.java
+++ b/src/main/java/com/google/devtools/build/lib/runtime/SkymeldUiStateTracker.java
@@ -171,6 +171,7 @@ void loadingComplete(LoadingPhaseCompleteEvent event) {
     } else {
       additionalMessage = labelsCount + " targets";
     }
+    mainRepositoryMapping = event.getMainRepositoryMapping();
   }
 
   @Override
diff --git a/src/main/java/com/google/devtools/build/lib/runtime/UiStateTracker.java b/src/main/java/com/google/devtools/build/lib/runtime/UiStateTracker.java
index c38646b2085cd2..050eac447ebc8a 100644
--- a/src/main/java/com/google/devtools/build/lib/runtime/UiStateTracker.java
+++ b/src/main/java/com/google/devtools/build/lib/runtime/UiStateTracker.java
@@ -42,6 +42,7 @@
 import com.google.devtools.build.lib.buildtool.buildevent.TestFilteringCompleteEvent;
 import com.google.devtools.build.lib.clock.Clock;
 import com.google.devtools.build.lib.cmdline.Label;
+import com.google.devtools.build.lib.cmdline.RepositoryMapping;
 import com.google.devtools.build.lib.events.Event;
 import com.google.devtools.build.lib.events.ExtendedEventHandler.FetchProgress;
 import com.google.devtools.build.lib.pkgcache.LoadingPhaseCompleteEvent;
@@ -90,6 +91,8 @@ class UiStateTracker {
 
   protected String status;
   protected String additionalMessage;
+  // Not null after the loading phase has completed.
+  protected RepositoryMapping mainRepositoryMapping;
 
   protected final Clock clock;
 
@@ -428,6 +431,7 @@ void loadingComplete(LoadingPhaseCompleteEvent event) {
     } else {
       additionalMessage = count + " targets";
     }
+    mainRepositoryMapping = event.getMainRepositoryMapping();
   }
 
   /**
@@ -642,11 +646,11 @@ static String suffix(String s, int len) {
    * If possible come up with a human-readable description of the label that fits within the given
    * width; a non-positive width indicates not no restriction at all.
    */
-  private static String shortenedLabelString(Label label, int width) {
+  private String shortenedLabelString(Label label, int width) {
     if (width <= 0) {
-      return label.toString();
+      return label.getDisplayForm(mainRepositoryMapping);
     }
-    String name = label.toString();
+    String name = label.getDisplayForm(mainRepositoryMapping);
     if (name.length() <= width) {
       return name;
     }
@@ -788,7 +792,7 @@ protected String describeAction(
       postfix += " " + strategy;
     }
 
-    String message = action.getProgressMessage();
+    String message = action.getProgressMessage(mainRepositoryMapping);
     if (message == null) {
       message = action.prettyPrint();
     }
diff --git a/src/main/java/com/google/devtools/build/lib/skyframe/TargetPatternPhaseFunction.java b/src/main/java/com/google/devtools/build/lib/skyframe/TargetPatternPhaseFunction.java
index acebe98f8b8e1a..7a2d4d6408a6bb 100644
--- a/src/main/java/com/google/devtools/build/lib/skyframe/TargetPatternPhaseFunction.java
+++ b/src/main/java/com/google/devtools/build/lib/skyframe/TargetPatternPhaseFunction.java
@@ -236,7 +236,8 @@ public TargetPatternPhaseValue compute(SkyKey key, Environment env) throws Inter
                 mapOriginalPatternsToLabels(expandedPatterns, targets.getTargets()),
                 testSuiteExpansions.buildOrThrow()));
     env.getListener()
-        .post(new LoadingPhaseCompleteEvent(result.getTargetLabels(), removedTargetLabels));
+        .post(new LoadingPhaseCompleteEvent(result.getTargetLabels(), removedTargetLabels,
+            repositoryMappingValue.getRepositoryMapping()));
     return result;
   }
 
diff --git a/src/test/java/com/google/devtools/build/lib/runtime/SkymeldUiStateTrackerTest.java b/src/test/java/com/google/devtools/build/lib/runtime/SkymeldUiStateTrackerTest.java
index 727928919cbc84..b3d7b8267f405f 100644
--- a/src/test/java/com/google/devtools/build/lib/runtime/SkymeldUiStateTrackerTest.java
+++ b/src/test/java/com/google/devtools/build/lib/runtime/SkymeldUiStateTrackerTest.java
@@ -24,6 +24,7 @@
 import com.google.devtools.build.lib.buildtool.ExecutionProgressReceiver;
 import com.google.devtools.build.lib.buildtool.buildevent.BuildCompleteEvent;
 import com.google.devtools.build.lib.buildtool.buildevent.ExecutionProgressReceiverAvailableEvent;
+import com.google.devtools.build.lib.cmdline.RepositoryMapping;
 import com.google.devtools.build.lib.pkgcache.LoadingPhaseCompleteEvent;
 import com.google.devtools.build.lib.runtime.SkymeldUiStateTracker.BuildStatus;
 import com.google.devtools.build.lib.skyframe.ConfigurationPhaseStartedEvent;
@@ -74,7 +75,8 @@ public void loadingComplete_stateChanges() {
     uiStateTracker.buildStatus = BuildStatus.TARGET_PATTERN_PARSING;
 
     uiStateTracker.loadingComplete(
-        new LoadingPhaseCompleteEvent(ImmutableSet.of(), ImmutableSet.of()));
+        new LoadingPhaseCompleteEvent(ImmutableSet.of(), ImmutableSet.of(),
+            RepositoryMapping.ALWAYS_FALLBACK));
 
     assertThat(uiStateTracker.buildStatus).isEqualTo(BuildStatus.LOADING_COMPLETE);
   }
diff --git a/src/test/java/com/google/devtools/build/lib/runtime/UiStateTrackerTest.java b/src/test/java/com/google/devtools/build/lib/runtime/UiStateTrackerTest.java
index 16237d01a28e26..aab4405f4006e2 100644
--- a/src/test/java/com/google/devtools/build/lib/runtime/UiStateTrackerTest.java
+++ b/src/test/java/com/google/devtools/build/lib/runtime/UiStateTrackerTest.java
@@ -18,7 +18,10 @@
 import static org.hamcrest.CoreMatchers.containsString;
 import static org.hamcrest.CoreMatchers.not;
 import static org.hamcrest.MatcherAssert.assertThat;
+import static org.mockito.ArgumentMatchers.eq;
 import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.never;
+import static org.mockito.Mockito.verify;
 import static org.mockito.Mockito.when;
 
 import com.google.common.collect.ImmutableList;
@@ -54,6 +57,8 @@
 import com.google.devtools.build.lib.buildtool.buildevent.TestFilteringCompleteEvent;
 import com.google.devtools.build.lib.cmdline.Label;
 import com.google.devtools.build.lib.cmdline.LabelSyntaxException;
+import com.google.devtools.build.lib.cmdline.RepositoryMapping;
+import com.google.devtools.build.lib.cmdline.RepositoryName;
 import com.google.devtools.build.lib.events.ExtendedEventHandler.FetchProgress;
 import com.google.devtools.build.lib.pkgcache.LoadingPhaseCompleteEvent;
 import com.google.devtools.build.lib.runtime.SkymeldUiStateTracker.BuildStatus;
@@ -83,12 +88,18 @@
 import net.starlark.java.syntax.Location;
 import org.junit.Test;
 import org.junit.runner.RunWith;
+import org.mockito.AdditionalMatchers;
 
-/** Tests {@link UiStateTracker}. */
+/**
+ * Tests {@link UiStateTracker}.
+ */
 @RunWith(TestParameterInjector.class)
 public class UiStateTrackerTest extends FoundationTestCase {
 
-  @TestParameter boolean isSkymeld;
+  @TestParameter
+  boolean isSkymeld;
+  static final RepositoryMapping MOCK_REPO_MAPPING = RepositoryMapping.createAllowingFallback(
+      ImmutableMap.of("main", RepositoryName.MAIN));
 
   private UiStateTracker getUiStateTracker(ManualClock clock) {
     if (isSkymeld) {
@@ -186,8 +197,11 @@ private Action mockAction(String progressMessage, String primaryOutput) {
         ActionsTestUtil.createArtifact(ArtifactRoot.asSourceRoot(Root.fromPath(outputBase)), path);
 
     Action action = mock(Action.class);
-    when(action.getProgressMessage()).thenReturn(progressMessage);
+    when(action.getProgressMessage(eq(MOCK_REPO_MAPPING))).thenReturn(progressMessage);
     when(action.getPrimaryOutput()).thenReturn(artifact);
+
+    verify(action, never()).getProgressMessage(AdditionalMatchers.not(eq(MOCK_REPO_MAPPING)));
+    verify(action, never()).getProgressMessage();
     return action;
   }
 
@@ -208,9 +222,13 @@ private ActionOwner dummyActionOwner() throws LabelSyntaxException {
   }
 
   private void simulateExecutionPhase(UiStateTracker uiStateTracker) {
+    uiStateTracker.loadingComplete(
+        new LoadingPhaseCompleteEvent(ImmutableSet.of(), ImmutableSet.of(), MOCK_REPO_MAPPING));
     if (this.isSkymeld) {
       // SkymeldUiStateTracker needs to be in the configuration phase before the execution phase.
       ((SkymeldUiStateTracker) uiStateTracker).buildStatus = BuildStatus.ANALYSIS_COMPLETE;
+    } else {
+      uiStateTracker.analysisComplete();
     }
     uiStateTracker.progressReceiverAvailable(
         new ExecutionProgressReceiverAvailableEvent(dummyExecutionProgressReceiver()));
@@ -253,8 +271,8 @@ public void testLoadingActivity() throws IOException {
     assertThat(loadingOutput).contains(loadingActivity);
 
     // When it is configuring targets.
-    stateTracker.loadingComplete(
-        new LoadingPhaseCompleteEvent(ImmutableSet.of(), ImmutableSet.of()));
+    stateTracker.loadingComplete(new LoadingPhaseCompleteEvent(ImmutableSet.of(), ImmutableSet.of(),
+        MOCK_REPO_MAPPING));
     String additionalMessage = "5 targets";
     stateTracker.additionalMessage = additionalMessage;
     String configuredTargetProgressString = "5 targets configured";