Extend TemplateExpander to handle $(func param) expansion

Rewrite the Expander to use the new functionality; also rewrite the Skylark
expand_location function to use it.

PiperOrigin-RevId: 173280839

10 files changed, 215 insertions, 61 deletions

diff --git a/src/main/java/com/google/devtools/build/lib/analysis/ConfigurationMakeVariableContext.java b/src/main/java/com/google/devtools/build/lib/analysis/ConfigurationMakeVariableContext.java
index a5c7cbfb99..f9c170a0e2 100644
--- a/src/main/java/com/google/devtools/build/lib/analysis/ConfigurationMakeVariableContext.java
+++ b/src/main/java/com/google/devtools/build/lib/analysis/ConfigurationMakeVariableContext.java
@@ -105,4 +105,9 @@ public class ConfigurationMakeVariableContext implements TemplateContext {
     }
     return SkylarkDict.<String, String>copyOf(null, map);
   }
+
+  @Override
+  public String lookupFunction(String name, String param) throws ExpansionException {
+    throw new ExpansionException(String.format("$(%s) not defined", name));
+  }
 }
diff --git a/src/main/java/com/google/devtools/build/lib/analysis/Expander.java b/src/main/java/com/google/devtools/build/lib/analysis/Expander.java
index c4eb867427..9dc224a29c 100644
--- a/src/main/java/com/google/devtools/build/lib/analysis/Expander.java
+++ b/src/main/java/com/google/devtools/build/lib/analysis/Expander.java
@@ -19,6 +19,7 @@ import com.google.common.collect.ImmutableMap;
 import com.google.devtools.build.lib.actions.Artifact;
 import com.google.devtools.build.lib.analysis.LocationExpander.Options;
 import com.google.devtools.build.lib.analysis.stringtemplate.ExpansionException;
+import com.google.devtools.build.lib.analysis.stringtemplate.TemplateContext;
 import com.google.devtools.build.lib.analysis.stringtemplate.TemplateExpander;
 import com.google.devtools.build.lib.cmdline.Label;
 import com.google.devtools.build.lib.shell.ShellUtils;
@@ -38,37 +39,26 @@ public final class Expander {
   }
 
   private final RuleContext ruleContext;
-  private final ConfigurationMakeVariableContext makeVariableContext;
-  @Nullable private final LocationExpander locationExpander;
+  private final TemplateContext templateContext;
 
-  private Expander(
-      RuleContext ruleContext,
-      ConfigurationMakeVariableContext makeVariableContext,
-      @Nullable LocationExpander locationExpander) {
+  Expander(RuleContext ruleContext, TemplateContext templateContext) {
     this.ruleContext = ruleContext;
-    this.makeVariableContext = makeVariableContext;
-    this.locationExpander = locationExpander;
-  }
-
-  Expander(
-      RuleContext ruleContext,
-      ConfigurationMakeVariableContext makeVariableContext) {
-    this(ruleContext, makeVariableContext, null);
+    this.templateContext = templateContext;
   }
 
   /**
    * Returns a new instance that also expands locations using the default configuration of
-   * {@link LocationExpander}.
+   * {@link LocationTemplateContext}.
    */
   public Expander withLocations(Options... options) {
-    LocationExpander newLocationExpander =
-        new LocationExpander(ruleContext, options);
-    return new Expander(ruleContext, makeVariableContext, newLocationExpander);
+    TemplateContext newTemplateContext =
+        new LocationTemplateContext(templateContext, ruleContext, options);
+    return new Expander(ruleContext, newTemplateContext);
   }
 
   /**
    * Returns a new instance that also expands locations, passing {@link Options#ALLOW_DATA} to the
-   * underlying {@link LocationExpander}.
+   * underlying {@link LocationTemplateContext}.
    */
   public Expander withDataLocations() {
     return withLocations(Options.ALLOW_DATA);
@@ -76,7 +66,7 @@ public final class Expander {
 
   /**
    * Returns a new instance that also expands locations, passing {@link Options#ALLOW_DATA} and
-   * {@link Options#EXEC_PATHS} to the underlying {@link LocationExpander}.
+   * {@link Options#EXEC_PATHS} to the underlying {@link LocationTemplateContext}.
    */
   public Expander withDataExecLocations() {
     return withLocations(Options.ALLOW_DATA, Options.EXEC_PATHS);
@@ -84,12 +74,12 @@ public final class Expander {
 
   /**
    * Returns a new instance that also expands locations, passing the given location map, as well as
-   * {@link Options#EXEC_PATHS} to the underlying {@link LocationExpander}.
+   * {@link Options#EXEC_PATHS} to the underlying {@link LocationTemplateContext}.
    */
   public Expander withExecLocations(ImmutableMap<Label, ImmutableCollection<Artifact>> locations) {
-    LocationExpander newLocationExpander =
-        new LocationExpander(ruleContext, locations, Options.EXEC_PATHS);
-    return new Expander(ruleContext, makeVariableContext, newLocationExpander);
+    TemplateContext newTemplateContext =
+        new LocationTemplateContext(templateContext, ruleContext, locations, Options.EXEC_PATHS);
+    return new Expander(ruleContext, newTemplateContext);
   }
 
   /**
@@ -145,14 +135,15 @@ public final class Expander {
    * @param expression the string to expand.
    * @return the expansion of "expression".
    */
-  public String expand(String attributeName, String expression) {
-    if (locationExpander != null) {
-      expression = locationExpander.expandAttribute(attributeName, expression);
-    }
+  public String expand(@Nullable String attributeName, String expression) {
     try {
-      return TemplateExpander.expand(expression, makeVariableContext);
+      return TemplateExpander.expand(expression, templateContext);
     } catch (ExpansionException e) {
-      ruleContext.attributeError(attributeName, e.getMessage());
+      if (attributeName == null) {
+        ruleContext.ruleError(e.getMessage());
+      } else {
+        ruleContext.attributeError(attributeName, e.getMessage());
+      }
       return expression;
     }
   }
@@ -214,7 +205,7 @@ public final class Expander {
   @Nullable
   public String expandSingleMakeVariable(String attrName, String expression) {
     try {
-      return TemplateExpander.expandSingleVariable(expression, makeVariableContext);
+      return TemplateExpander.expandSingleVariable(expression, templateContext);
     } catch (ExpansionException e) {
       ruleContext.attributeError(attrName, e.getMessage());
       return expression;
diff --git a/src/main/java/com/google/devtools/build/lib/analysis/LocationExpander.java b/src/main/java/com/google/devtools/build/lib/analysis/LocationExpander.java
index e43aa71bf4..7bba08b03f 100644
--- a/src/main/java/com/google/devtools/build/lib/analysis/LocationExpander.java
+++ b/src/main/java/com/google/devtools/build/lib/analysis/LocationExpander.java
@@ -60,7 +60,7 @@ import javax.annotation.Nullable;
  *
  * <p>DO NOT USE DIRECTLY! Use RuleContext.getExpander() instead.
  */
-public class LocationExpander {
+final class LocationExpander {
 
   /**
    * List of options to tweak the LocationExpander.
@@ -329,7 +329,7 @@ public class LocationExpander {
    * @param labelMap map of labels to build artifacts
    * @return map of all possible target locations
    */
-  private static Map<Label, Collection<Artifact>> buildLocationMap(
+  static Map<Label, Collection<Artifact>> buildLocationMap(
       RuleContext ruleContext,
       Map<Label, ? extends Collection<Artifact>> labelMap,
       boolean allowDataAttributeEntriesInLabel) {
diff --git a/src/main/java/com/google/devtools/build/lib/analysis/LocationTemplateContext.java b/src/main/java/com/google/devtools/build/lib/analysis/LocationTemplateContext.java
new file mode 100644
index 0000000000..d2888200e0
--- /dev/null
+++ b/src/main/java/com/google/devtools/build/lib/analysis/LocationTemplateContext.java
@@ -0,0 +1,111 @@
+// Copyright 2017 The Bazel Authors. All rights reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//    http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package com.google.devtools.build.lib.analysis;
+
+import com.google.common.base.Supplier;
+import com.google.common.base.Suppliers;
+import com.google.common.collect.ImmutableCollection;
+import com.google.common.collect.ImmutableMap;
+import com.google.common.collect.ImmutableSet;
+import com.google.devtools.build.lib.actions.Artifact;
+import com.google.devtools.build.lib.analysis.LocationExpander.LocationFunction;
+import com.google.devtools.build.lib.analysis.LocationExpander.Options;
+import com.google.devtools.build.lib.analysis.stringtemplate.ExpansionException;
+import com.google.devtools.build.lib.analysis.stringtemplate.TemplateContext;
+import com.google.devtools.build.lib.cmdline.Label;
+import java.util.Collection;
+import java.util.Map;
+import java.util.function.Function;
+import javax.annotation.Nullable;
+
+/**
+ * Expands $(location) and $(locations) tags inside target attributes. You can specify something
+ * like this in the BUILD file:
+ *
+ * <pre>
+ * somerule(name='some name',
+ *          someopt = [ '$(location //mypackage:myhelper)' ],
+ *          ...)
+ * </pre>
+ *
+ * and location will be substituted with //mypackage:myhelper executable output.
+ *
+ * <p>Note that this expander will always expand labels in srcs, deps, and tools attributes, with
+ * data being optional.
+ *
+ * <p>DO NOT USE DIRECTLY! Use RuleContext.getExpander() instead.
+ */
+final class LocationTemplateContext implements TemplateContext {
+  private final TemplateContext delegate;
+  private final Function<String, String> locationFunction;
+  private final Function<String, String> locationsFunction;
+
+  private LocationTemplateContext(
+      TemplateContext delegate,
+      Label root,
+      Supplier<Map<Label, Collection<Artifact>>> locationMap,
+      boolean execPaths) {
+    this.delegate = delegate;
+    this.locationFunction = new LocationFunction(root, locationMap, execPaths, false);
+    this.locationsFunction = new LocationFunction(root, locationMap, execPaths, true);
+  }
+
+  private LocationTemplateContext(
+      TemplateContext delegate,
+      RuleContext ruleContext,
+      @Nullable ImmutableMap<Label, ImmutableCollection<Artifact>> labelMap,
+      ImmutableSet<Options> options) {
+    this(
+        delegate,
+        ruleContext.getLabel(),
+        // Use a memoizing supplier to avoid eagerly building the location map.
+        Suppliers.memoize(
+            () -> LocationExpander.buildLocationMap(
+                ruleContext, labelMap, options.contains(Options.ALLOW_DATA))),
+        options.contains(Options.EXEC_PATHS));
+  }
+
+  public LocationTemplateContext(
+      TemplateContext delegate,
+      RuleContext ruleContext,
+      @Nullable ImmutableMap<Label, ImmutableCollection<Artifact>> labelMap,
+      LocationExpander.Options... options) {
+    this(delegate, ruleContext, labelMap, ImmutableSet.copyOf(options));
+  }
+
+  public LocationTemplateContext(
+      TemplateContext delegate, RuleContext ruleContext, LocationExpander.Options... options) {
+    this(delegate, ruleContext, null, ImmutableSet.copyOf(options));
+  }
+
+  @Override
+  public String lookupVariable(String name) throws ExpansionException {
+    return delegate.lookupVariable(name);
+  }
+
+  @Override
+  public String lookupFunction(String name, String param) throws ExpansionException {
+    try {
+      if ("location".equals(name)) {
+        return locationFunction.apply(param);
+      } else if ("locations".equals(name)) {
+        return locationsFunction.apply(param);
+      }
+    } catch (IllegalStateException e) {
+      throw new ExpansionException(e.getMessage(), e);
+    }
+    return delegate.lookupFunction(name, param);
+  }
+}
diff --git a/src/main/java/com/google/devtools/build/lib/analysis/RuleContext.java b/src/main/java/com/google/devtools/build/lib/analysis/RuleContext.java
index 94423b6482..15337837fc 100644
--- a/src/main/java/com/google/devtools/build/lib/analysis/RuleContext.java
+++ b/src/main/java/com/google/devtools/build/lib/analysis/RuleContext.java
@@ -48,6 +48,7 @@ import com.google.devtools.build.lib.analysis.config.FragmentCollection;
 import com.google.devtools.build.lib.analysis.config.PatchTransition;
 import com.google.devtools.build.lib.analysis.configuredtargets.RuleConfiguredTarget.Mode;
 import com.google.devtools.build.lib.analysis.fileset.FilesetProvider;
+import com.google.devtools.build.lib.analysis.stringtemplate.TemplateContext;
 import com.google.devtools.build.lib.cmdline.Label;
 import com.google.devtools.build.lib.cmdline.RepositoryName;
 import com.google.devtools.build.lib.collect.ImmutableSortedKeyListMultimap;
@@ -954,8 +955,8 @@ public final class RuleContext extends TargetContext
     initConfigurationMakeVariableContext(ImmutableList.copyOf(makeVariableSuppliers));
   }
 
-  public Expander getExpander(ConfigurationMakeVariableContext makeVariableContext) {
-    return new Expander(this, makeVariableContext);
+  public Expander getExpander(TemplateContext templateContext) {
+    return new Expander(this, templateContext);
   }
 
   public Expander getExpander() {
diff --git a/src/main/java/com/google/devtools/build/lib/analysis/skylark/SkylarkRuleImplementationFunctions.java b/src/main/java/com/google/devtools/build/lib/analysis/skylark/SkylarkRuleImplementationFunctions.java
index f71730b10a..53aaf47e79 100644
--- a/src/main/java/com/google/devtools/build/lib/analysis/skylark/SkylarkRuleImplementationFunctions.java
+++ b/src/main/java/com/google/devtools/build/lib/analysis/skylark/SkylarkRuleImplementationFunctions.java
@@ -20,11 +20,12 @@ import com.google.devtools.build.lib.actions.Artifact;
 import com.google.devtools.build.lib.analysis.AliasProvider;
 import com.google.devtools.build.lib.analysis.CommandHelper;
 import com.google.devtools.build.lib.analysis.FileProvider;
-import com.google.devtools.build.lib.analysis.LocationExpander;
 import com.google.devtools.build.lib.analysis.Runfiles;
 import com.google.devtools.build.lib.analysis.RunfilesProvider;
 import com.google.devtools.build.lib.analysis.TransitiveInfoCollection;
 import com.google.devtools.build.lib.analysis.configuredtargets.AbstractConfiguredTarget;
+import com.google.devtools.build.lib.analysis.stringtemplate.ExpansionException;
+import com.google.devtools.build.lib.analysis.stringtemplate.TemplateContext;
 import com.google.devtools.build.lib.cmdline.Label;
 import com.google.devtools.build.lib.events.Location;
 import com.google.devtools.build.lib.skylarkinterface.Param;
@@ -318,11 +319,12 @@ public class SkylarkRuleImplementationFunctions {
             throws EvalException {
           ctx.checkMutable("expand_location");
           try {
-            return new LocationExpander(
-                    ctx.getRuleContext(),
-                    makeLabelMap(targets.getContents(TransitiveInfoCollection.class, "targets")),
-                    LocationExpander.Options.EXEC_PATHS)
-                .expand(input);
+            return ctx
+                .getRuleContext()
+                .getExpander(new FakeTemplateContext())
+                .withExecLocations(
+                    makeLabelMap(targets.getContents(TransitiveInfoCollection.class, "targets")))
+                .expand(/*attributeName=*/null, input);
           } catch (IllegalStateException ise) {
             throw new EvalException(loc, ise);
           }
@@ -330,6 +332,23 @@ public class SkylarkRuleImplementationFunctions {
       };
 
   /**
+   * This class only exists for backwards compatibility. Previously, we were using LocationExpander
+   * directly, which passes through all unrecognized $() expressions.
+   */
+  private static final class FakeTemplateContext implements TemplateContext {
+    @Override
+    public String lookupVariable(String name) throws ExpansionException {
+      return String.format("$$(%s)", name);
+    }
+
+    @Override
+    public String lookupFunction(String name, String param) throws ExpansionException {
+      // Variables get recursively expanded, functions don't.
+      return String.format("$(%s %s)", name, param);
+    }
+  }
+
+  /**
    * Builds a map: Label -> List of files from the given labels
    *
    * @param knownLabels List of known labels
diff --git a/src/main/java/com/google/devtools/build/lib/analysis/stringtemplate/ExpansionException.java b/src/main/java/com/google/devtools/build/lib/analysis/stringtemplate/ExpansionException.java
index 66f5437667..10542e770c 100644
--- a/src/main/java/com/google/devtools/build/lib/analysis/stringtemplate/ExpansionException.java
+++ b/src/main/java/com/google/devtools/build/lib/analysis/stringtemplate/ExpansionException.java
@@ -21,4 +21,8 @@ public class ExpansionException extends Exception {
   public ExpansionException(String message) {
     super(message);
   }
+
+  public ExpansionException(String message, Throwable cause) {
+    super(message, cause);
+  }
 }
\ No newline at end of file
diff --git a/src/main/java/com/google/devtools/build/lib/analysis/stringtemplate/TemplateContext.java b/src/main/java/com/google/devtools/build/lib/analysis/stringtemplate/TemplateContext.java
index d4aa40dcf0..dc6d27598d 100644
--- a/src/main/java/com/google/devtools/build/lib/analysis/stringtemplate/TemplateContext.java
+++ b/src/main/java/com/google/devtools/build/lib/analysis/stringtemplate/TemplateContext.java
@@ -14,18 +14,29 @@
 package com.google.devtools.build.lib.analysis.stringtemplate;
 
 /**
- * Interface to be implemented by callers of MakeVariableExpander which defines the expansion of
- * each "Make" variable.
+ * Interface to be implemented by callers of {@link TemplateExpander} which defines the expansion of
+ * each template variable and function.
  */
 public interface TemplateContext {
 
   /**
-   * Returns the expansion of the specified "Make" variable.
+   * Returns the expansion of the specified template variable.
    *
-   * @param name the variable to expand.
-   * @return the expansion of the variable.
-   * @throws ExpansionException if the variable "var" was not defined or
-   *     there was any other error while expanding "var".
+   * @param name the variable to expand
+   * @return the expansion of the variable
+   * @throws ExpansionException if the given variable was not defined or there was any other error
+   *     during expansion
    */
   String lookupVariable(String name) throws ExpansionException;
+
+  /**
+   * Returns the expansion of the specified template function with the given parameter.
+   *
+   * @param name the function name
+   * @param param the function parameter
+   * @return the expansion of the function applied to the parameter
+   * @throws ExpansionException if the function was not defined, or if the function application
+   *     failed for some reason
+   */
+  String lookupFunction(String name, String param) throws ExpansionException;
 }
\ No newline at end of file
diff --git a/src/main/java/com/google/devtools/build/lib/analysis/stringtemplate/TemplateExpander.java b/src/main/java/com/google/devtools/build/lib/analysis/stringtemplate/TemplateExpander.java
index 9807a7c3cc..71ca2ca451 100644
--- a/src/main/java/com/google/devtools/build/lib/analysis/stringtemplate/TemplateExpander.java
+++ b/src/main/java/com/google/devtools/build/lib/analysis/stringtemplate/TemplateExpander.java
@@ -15,12 +15,10 @@ package com.google.devtools.build.lib.analysis.stringtemplate;
 
 /**
  * Simple string template expansion. String templates consist of text interspersed with
- * <code>$(variable)</code> references, which are replaced by strings.
- *
- * <p>Note that neither <code>$(location x)</code> nor Make-isms are treated specially in any way by
- * this class.
+ * <code>$(variable)</code> or <code>$(function value)</code> references, which are replaced by
+ * strings.
  */
-public class TemplateExpander {
+public final class TemplateExpander {
   private final char[] buffer;
   private final int length;
   private int offset;
@@ -81,13 +79,21 @@ public class TemplateExpander {
           result.append('$');
         } else {
           String var = scanVariable();
-          String value = context.lookupVariable(var);
-          // To prevent infinite recursion for the ignored shell variables
-          if (!value.equals(var)) {
-            // recursively expand using Make's ":=" semantics:
-            value = expand(value, context, depth + 1);
+          int spaceIndex = var.indexOf(' ');
+          if (spaceIndex < 0) {
+            String value = context.lookupVariable(var);
+            // To prevent infinite recursion for the ignored shell variables
+            if (!value.equals(var)) {
+              // recursively expand using Make's ":=" semantics:
+              value = expand(value, context, depth + 1);
+            }
+            result.append(value);
+          } else {
+            String name = var.substring(0, spaceIndex);
+            String param = var.substring(spaceIndex + 1);
+            String value = context.lookupFunction(name, param);
+            result.append(value);
           }
-          result.append(value);
         }
       } else {
         result.append(c);
@@ -109,7 +115,7 @@ public class TemplateExpander {
   private String scanVariable() throws ExpansionException {
     char c = buffer[offset];
     switch (c) {
-      case '(': { // $(SRCS)
+      case '(': { // looks like $(SRCS)
         offset++;
         int start = offset;
         while (offset < length && buffer[offset] != ')') {
@@ -120,7 +126,8 @@ public class TemplateExpander {
         }
         return new String(buffer, start, offset - start);
       }
-      case '{': { // ${SRCS}
+      // We only parse ${variable} syntax to provide a better error message.
+      case '{': { // looks like ${SRCS}
         offset++;
         int start = offset;
         while (offset < length && buffer[offset] != '}') {
diff --git a/src/main/java/com/google/devtools/build/lib/rules/proto/ProtoCompileActionBuilder.java b/src/main/java/com/google/devtools/build/lib/rules/proto/ProtoCompileActionBuilder.java
index e0410dbb4d..b91e57f14b 100644
--- a/src/main/java/com/google/devtools/build/lib/rules/proto/ProtoCompileActionBuilder.java
+++ b/src/main/java/com/google/devtools/build/lib/rules/proto/ProtoCompileActionBuilder.java
@@ -205,6 +205,11 @@ public class ProtoCompileActionBuilder {
                 }
                 return value.toString();
               }
+
+              @Override
+              public String lookupFunction(String name, String param) throws ExpansionException {
+                throw new ExpansionException(String.format("$(%s) not defined", name));
+              }
             });
       } catch (ExpansionException e) {
         // Squeelch. We don't throw this exception in the lookupMakeVariable implementation above,