Make adding inherited vars to ActionEnvironment memory efficient

Along the way, this commit optimizes away a few potential allocations
of empty ActionEnvironment instances.

src/main/java/com/google/devtools/build/lib/actions/AbstractAction.java

@@ -699,7 +699,7 @@ public Dict<String, String> getExecutionInfoDict() {
 
   @Override
   public Dict<String, String> getEnv() {
-    return Dict.immutableCopyOf(env.getFixedEnv().toMap());
+    return Dict.immutableCopyOf(env.getFixedEnv());
   }
 
   @Override

src/main/java/com/google/devtools/build/lib/actions/ActionEnvironment.java

@@ -19,6 +19,7 @@
 import com.google.common.collect.ImmutableSet;
 import com.google.devtools.build.lib.util.Fingerprint;
 import java.util.LinkedHashMap;
+import java.util.LinkedHashSet;
 import java.util.Map;
 import java.util.Set;
 import java.util.TreeMap;
@@ -45,23 +46,34 @@
  */
 public final class ActionEnvironment {
 
-  /** A map of environment variables. */
+  /**
+   * A map of environment variables together with a list of variables to inherit from the shell
+   * environment.
+   */
   public interface EnvironmentVariables {
 
     /**
-     * Returns the environment variables as a map.
+     * Returns the fixed environment variables as a map.
+     *
+     * <p>WARNING: this allocates additional objects if the underlying implementation is a {@link
+     * CompoundEnvironmentVariables}; use sparingly.
+     */
+    ImmutableMap<String, String> getFixedEnvironment();
+
+    /**
+     * Returns the inherited environment variables as a set.
      *
-     * <p>WARNING: this allocations additional objects if the underlying implementation is a {@link
+     * <p>WARNING: this allocated additional objects if the underlying implementation is a {@link
      * CompoundEnvironmentVariables}; use sparingly.
      */
-    ImmutableMap<String, String> toMap();
+    ImmutableSet<String> getInheritedEnvironment();
 
     default boolean isEmpty() {
-      return toMap().isEmpty();
+      return getFixedEnvironment().isEmpty() && getInheritedEnvironment().isEmpty();
     }
 
     default int size() {
-      return toMap().size();
+      return getFixedEnvironment().size() + getInheritedEnvironment().size();
     }
   }
 
@@ -70,11 +82,21 @@ default int size() {
    * allocation a new map.
    */
   static class CompoundEnvironmentVariables implements EnvironmentVariables {
+
+    static EnvironmentVariables create(Map<String, String> fixedVars, Set<String> inheritedVars,
+        EnvironmentVariables base) {
+      if (fixedVars.isEmpty() && inheritedVars.isEmpty() && base.isEmpty()) {
+        return EMPTY_ENVIRONMENT_VARIABLES;
+      }
+      return new CompoundEnvironmentVariables(fixedVars, inheritedVars, base);
+    }
+
     private final EnvironmentVariables current;
     private final EnvironmentVariables base;
 
-    CompoundEnvironmentVariables(Map<String, String> vars, EnvironmentVariables base) {
-      this.current = new SimpleEnvironmentVariables(vars);
+    private CompoundEnvironmentVariables(Map<String, String> fixedVars, Set<String> inheritedVars,
+        EnvironmentVariables base) {
+      this.current = SimpleEnvironmentVariables.create(fixedVars, inheritedVars);
       this.base = base;
     }
 
@@ -84,48 +106,66 @@ public boolean isEmpty() {
     }
 
     @Override
-    public ImmutableMap<String, String> toMap() {
+    public ImmutableMap<String, String> getFixedEnvironment() {
       Map<String, String> result = new LinkedHashMap<>();
-      result.putAll(base.toMap());
-      result.putAll(current.toMap());
+      result.putAll(base.getFixedEnvironment());
+      result.putAll(current.getFixedEnvironment());
       return ImmutableMap.copyOf(result);
     }
+
+    @Override
+    public ImmutableSet<String> getInheritedEnvironment() {
+      Set<String> result = new LinkedHashSet<>();
+      result.addAll(base.getInheritedEnvironment());
+      result.addAll(current.getInheritedEnvironment());
+      return ImmutableSet.copyOf(result);
+    }
   }
 
-  /** A simple {@link EnvironmentVariables}. */
+  /**
+   * A simple {@link EnvironmentVariables}.
+   */
   static class SimpleEnvironmentVariables implements EnvironmentVariables {
 
-    static EnvironmentVariables create(Map<String, String> vars) {
-      if (vars.isEmpty()) {
+    static EnvironmentVariables create(Map<String, String> fixedVars, Set<String> inheritedVars) {
+      if (fixedVars.isEmpty() && inheritedVars.isEmpty()) {
         return EMPTY_ENVIRONMENT_VARIABLES;
       }
-      return new SimpleEnvironmentVariables(vars);
+      return new SimpleEnvironmentVariables(fixedVars, inheritedVars);
     }
 
-    private final ImmutableMap<String, String> vars;
+    private final ImmutableMap<String, String> fixedVars;
+    private final ImmutableSet<String> inheritedVars;
+
+    private SimpleEnvironmentVariables(Map<String, String> fixedVars, Set<String> inheritedVars) {
+      this.fixedVars = ImmutableMap.copyOf(fixedVars);
+      this.inheritedVars = ImmutableSet.copyOf(inheritedVars);
+    }
 
-    private SimpleEnvironmentVariables(Map<String, String> vars) {
-      this.vars = ImmutableMap.copyOf(vars);
+    @Override
+    public ImmutableMap<String, String> getFixedEnvironment() {
+      return fixedVars;
     }
 
     @Override
-    public ImmutableMap<String, String> toMap() {
-      return vars;
+    public ImmutableSet<String> getInheritedEnvironment() {
+      return inheritedVars;
     }
   }
 
-  /** An empty {@link EnvironmentVariables}. */
+  /**
+   * An empty {@link EnvironmentVariables}.
+   */
   public static final EnvironmentVariables EMPTY_ENVIRONMENT_VARIABLES =
-      new SimpleEnvironmentVariables(ImmutableMap.of());
+      new SimpleEnvironmentVariables(ImmutableMap.of(), ImmutableSet.of());
 
   /**
    * An empty environment, mainly for testing. Production code should never use this, but instead
    * get the proper environment from the current configuration.
    */
   // TODO(ulfjack): Migrate all production code to use the proper action environment, and then make
   // this @VisibleForTesting or rename it to clarify.
-  public static final ActionEnvironment EMPTY =
-      new ActionEnvironment(EMPTY_ENVIRONMENT_VARIABLES, ImmutableSet.of());
+  public static final ActionEnvironment EMPTY = new ActionEnvironment(EMPTY_ENVIRONMENT_VARIABLES);
 
   /**
    * Splits the given map into a map of variables with a fixed value, and a set of variables that
@@ -145,76 +185,69 @@ public static ActionEnvironment split(Map<String, String> env) {
         inheritedEnv.add(key);
       }
     }
-    return create(new SimpleEnvironmentVariables(fixedEnv), ImmutableSet.copyOf(inheritedEnv));
+    return create(SimpleEnvironmentVariables.create(fixedEnv, inheritedEnv));
   }
 
-  private final EnvironmentVariables fixedEnv;
-  private final ImmutableSet<String> inheritedEnv;
+  private final EnvironmentVariables vars;
 
-  private ActionEnvironment(EnvironmentVariables fixedEnv, ImmutableSet<String> inheritedEnv) {
-    this.fixedEnv = fixedEnv;
-    this.inheritedEnv = inheritedEnv;
+  private ActionEnvironment(EnvironmentVariables vars) {
+    this.vars = vars;
   }
 
   /**
    * Creates a new action environment. The order in which the environments are combined is
    * undefined, so callers need to take care that the key set of the {@code fixedEnv} map and the
    * set of {@code inheritedEnv} elements are disjoint.
    */
-  private static ActionEnvironment create(
-      EnvironmentVariables fixedEnv, ImmutableSet<String> inheritedEnv) {
-    if (fixedEnv.isEmpty() && inheritedEnv.isEmpty()) {
+  private static ActionEnvironment create(EnvironmentVariables vars) {
+    if (vars.isEmpty()) {
       return EMPTY;
     }
-    return new ActionEnvironment(fixedEnv, inheritedEnv);
+    return new ActionEnvironment(vars);
   }
 
   public static ActionEnvironment create(
       Map<String, String> fixedEnv, ImmutableSet<String> inheritedEnv) {
-    return new ActionEnvironment(SimpleEnvironmentVariables.create(fixedEnv), inheritedEnv);
+    return ActionEnvironment.create(SimpleEnvironmentVariables.create(fixedEnv, inheritedEnv));
   }
 
   public static ActionEnvironment create(Map<String, String> fixedEnv) {
-    return new ActionEnvironment(new SimpleEnvironmentVariables(fixedEnv), ImmutableSet.of());
+    return ActionEnvironment.create(SimpleEnvironmentVariables.create(fixedEnv, ImmutableSet.of()));
   }
 
   /**
    * Returns a copy of the environment with the given fixed variables added to it, <em>overwriting
    * any existing occurrences of those variables</em>.
    */
-  public ActionEnvironment addFixedVariables(Map<String, String> vars) {
-    return new ActionEnvironment(new CompoundEnvironmentVariables(vars, fixedEnv), inheritedEnv);
+  public ActionEnvironment addFixedVariables(Map<String, String> fixedVars) {
+    return ActionEnvironment.create(
+        CompoundEnvironmentVariables.create(fixedVars, ImmutableSet.of(),
+            vars));
   }
 
   /**
    * Returns a copy of the environment with the given fixed and inherited variables added to it,
    * <em>overwriting any existing occurrences of those variables</em>.
    */
-  public ActionEnvironment addVariables(Map<String, String> vars, Set<String> inheritedVars) {
-    if (inheritedVars.isEmpty()) {
-      return addFixedVariables(vars);
-    } else {
-      // TODO: inheritedEnv is currently not optimized for allocation avoidance in the same way as
-      //  fixedEnv.
-      ImmutableSet<String> newInheritedEnv = ImmutableSet.<String>builder().addAll(inheritedEnv)
-          .addAll(inheritedVars).build();
-      return new ActionEnvironment(new CompoundEnvironmentVariables(vars, fixedEnv),
-          newInheritedEnv);
-    }
+  public ActionEnvironment addVariables(Map<String, String> fixedVars, Set<String> inheritedVars) {
+    return ActionEnvironment.create(
+        CompoundEnvironmentVariables.create(fixedVars, inheritedVars, vars));
   }
 
-  /** Returns the combined size of the fixed and inherited environments. */
+  /**
+   * Returns the combined size of the fixed and inherited environments.
+   */
   public int size() {
-    return fixedEnv.size() + inheritedEnv.size();
+    return vars.size();
   }
 
   /**
    * Returns the 'fixed' part of the environment, i.e., those environment variables that are set to
    * fixed values and their values. This should only be used for testing and to compute the cache
    * keys of actions. Use {@link #resolve} instead to get the complete environment.
    */
-  public EnvironmentVariables getFixedEnv() {
-    return fixedEnv;
+  public ImmutableMap<String, String> getFixedEnv() {
+    return vars.getFixedEnvironment();
   }
 
   /**
@@ -224,7 +257,7 @@ public EnvironmentVariables getFixedEnv() {
    * get the complete environment.
    */
   public ImmutableSet<String> getInheritedEnv() {
-    return inheritedEnv;
+    return vars.getInheritedEnvironment();
   }
 
   /**
@@ -235,8 +268,8 @@ public ImmutableSet<String> getInheritedEnv() {
    */
   public void resolve(Map<String, String> result, Map<String, String> clientEnv) {
     checkNotNull(clientEnv);
-    result.putAll(fixedEnv.toMap());
-    for (String var : inheritedEnv) {
+    result.putAll(vars.getFixedEnvironment());
+    for (String var : vars.getInheritedEnvironment()) {
       String value = clientEnv.get(var);
       if (value != null) {
         result.put(var, value);
@@ -245,7 +278,7 @@ public void resolve(Map<String, String> result, Map<String, String> clientEnv) {
   }
 
   public void addTo(Fingerprint f) {
-    f.addStringMap(fixedEnv.toMap());
-    f.addStrings(inheritedEnv);
+    f.addStringMap(vars.getFixedEnvironment());
+    f.addStrings(vars.getInheritedEnvironment());
   }
 }

src/main/java/com/google/devtools/build/lib/analysis/actions/SpawnAction.java

@@ -465,7 +465,7 @@ public String describeKey() {
     StringBuilder message = new StringBuilder();
     message.append(getProgressMessage());
     message.append('\n');
-    for (Map.Entry<String, String> entry : env.getFixedEnv().toMap().entrySet()) {
+    for (Map.Entry<String, String> entry : env.getFixedEnv().entrySet()) {
       message.append("  Environment variable: ");
       message.append(ShellEscaper.escapeString(entry.getKey()));
       message.append('=');
@@ -574,7 +574,7 @@ public final ImmutableMap<String, String> getIncompleteEnvironmentForTesting() {
     // ActionEnvironment to avoid developers misunderstanding the purpose of this method. That
     // requires first updating all subclasses and callers to actually handle environments correctly,
     // so it's not a small change.
-    return env.getFixedEnv().toMap();
+    return env.getFixedEnv();
   }
 
   /** Returns the out-of-band execution data for this action. */

src/main/java/com/google/devtools/build/lib/analysis/config/BuildConfigurationValue.java

@@ -485,7 +485,7 @@ public boolean isSiblingRepositoryLayout() {
    */
   @Override
   public ImmutableMap<String, String> getLocalShellEnvironment() {
-    return actionEnv.getFixedEnv().toMap();
+    return actionEnv.getFixedEnv();
   }
 
   /**
@@ -629,7 +629,7 @@ public boolean legacyExternalRunfiles() {
    */
   @Override
   public ImmutableMap<String, String> getTestEnv() {
-    return testEnv.getFixedEnv().toMap();
+    return testEnv.getFixedEnv();
   }
 
   /**

src/main/java/com/google/devtools/build/lib/analysis/test/TestRunnerAction.java

@@ -1051,7 +1051,7 @@ public List<String> getArguments() throws CommandLineExpansionException, Interru
   @Override
   public ImmutableMap<String, String> getIncompleteEnvironmentForTesting()
       throws ActionExecutionException {
-    return getEnvironment().getFixedEnv().toMap();
+    return getEnvironment().getFixedEnv();
   }
 
   @Override

src/main/java/com/google/devtools/build/lib/query2/aquery/ActionGraphTextOutputFormatterCallback.java

@@ -221,7 +221,7 @@ private void writeAction(ActionAnalysisMetadata action, PrintStream printStream)
       // TODO(twerth): This handles the fixed environment. We probably want to output the inherited
       // environment as well.
       Iterable<Map.Entry<String, String>> fixedEnvironment =
-          abstractAction.getEnvironment().getFixedEnv().toMap().entrySet();
+          abstractAction.getEnvironment().getFixedEnv().entrySet();
       if (!Iterables.isEmpty(fixedEnvironment)) {
         stringBuilder
             .append("  Environment: [")

src/main/java/com/google/devtools/build/lib/rules/java/JavaCompileAction.java

@@ -614,7 +614,7 @@ public final ImmutableMap<String, String> getIncompleteEnvironmentForTesting() {
     // ActionEnvironment to avoid developers misunderstanding the purpose of this method. That
     // requires first updating all subclasses and callers to actually handle environments correctly,
     // so it's not a small change.
-    return env.getFixedEnv().toMap();
+    return env.getFixedEnv();
   }
 
   @Override

src/main/java/com/google/devtools/build/lib/skyframe/actiongraph/v2/ActionGraphDump.java

@@ -161,7 +161,7 @@ private void dumpSingleAction(ConfiguredTarget configuredTarget, ActionAnalysisM
       SpawnAction spawnAction = (SpawnAction) action;
       // TODO(twerth): This handles the fixed environment. We probably want to output the inherited
       // environment as well.
-      Map<String, String> fixedEnvironment = spawnAction.getEnvironment().getFixedEnv().toMap();
+      Map<String, String> fixedEnvironment = spawnAction.getEnvironment().getFixedEnv();
       for (Map.Entry<String, String> environmentVariable : fixedEnvironment.entrySet()) {
         actionBuilder.addEnvironmentVariables(
             AnalysisProtosV2.KeyValuePair.newBuilder()

src/test/java/com/google/devtools/build/lib/actions/ActionEnvironmentTest.java

@@ -34,10 +34,10 @@ public void compoundEnvOrdering() {
     // entries added by env2 override the existing entries
     ActionEnvironment env2 = env1.addFixedVariables(ImmutableMap.of("FOO", "foo2"));
 
-    assertThat(env1.getFixedEnv().toMap()).containsExactly("FOO", "foo1", "BAR", "bar");
+    assertThat(env1.getFixedEnv()).containsExactly("FOO", "foo1", "BAR", "bar");
     assertThat(env1.getInheritedEnv()).containsExactly("baz");
 
-    assertThat(env2.getFixedEnv().toMap()).containsExactly("FOO", "foo2", "BAR", "bar");
+    assertThat(env2.getFixedEnv()).containsExactly("FOO", "foo2", "BAR", "bar");
     assertThat(env2.getInheritedEnv()).containsExactly("baz");
   }
 }

src/test/java/com/google/devtools/build/lib/bazel/rules/BazelRuleClassProviderTest.java

@@ -176,8 +176,8 @@ public void strictActionEnv() throws Exception {
             "--action_env=FOO=bar");
 
     ActionEnvironment env = BazelRuleClassProvider.SHELL_ACTION_ENV.apply(options);
-    assertThat(env.getFixedEnv().toMap()).containsEntry("PATH", "/bin:/usr/bin:/usr/local/bin");
-    assertThat(env.getFixedEnv().toMap()).containsEntry("FOO", "bar");
+    assertThat(env.getFixedEnv()).containsEntry("PATH", "/bin:/usr/bin:/usr/local/bin");
+    assertThat(env.getFixedEnv()).containsEntry("FOO", "bar");
   }
 
   @Test