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

13 files changed, 356 insertions, 148 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,
diff --git a/src/test/java/com/google/devtools/build/lib/analysis/stringtemplate/TemplateExpanderTest.java b/src/test/java/com/google/devtools/build/lib/analysis/stringtemplate/TemplateExpanderTest.java
index fa61f091b3..33a2f5bcc2 100644
--- a/src/test/java/com/google/devtools/build/lib/analysis/stringtemplate/TemplateExpanderTest.java
+++ b/src/test/java/com/google/devtools/build/lib/analysis/stringtemplate/TemplateExpanderTest.java
@@ -18,132 +18,186 @@ import static org.junit.Assert.fail;
 
 import java.util.HashMap;
 import java.util.Map;
+import java.util.function.Function;
 import org.junit.Before;
 import org.junit.Test;
 import org.junit.runner.RunWith;
 import org.junit.runners.JUnit4;
 
 /**
- * Unit tests for the {@link
- * com.google.devtools.build.lib.analysis.stringtemplate.TemplateExpander}, which expands variable
- * references of the form <code>"$x"</code> and <code>"$(foo)"</code> into their corresponding
- * values.
+ * Unit tests for the {@link TemplateExpander}.
  */
 @RunWith(JUnit4.class)
 public class TemplateExpanderTest {
+  private static final class TemplateContextImpl implements TemplateContext {
+    private final Map<String, String> vars = new HashMap<>();
+    private final Map<String, Function<String, String>> functions = new HashMap<>();
+
+    @Override
+    public String lookupVariable(String name)
+        throws ExpansionException {
+      // Not a Make variable. Let the shell handle the expansion.
+      if (name.startsWith("$")) {
+        return name;
+      }
+      if (!vars.containsKey(name)) {
+        throw new ExpansionException(String.format("$(%s) not defined", name));
+      }
+      return vars.get(name);
+    }
 
-  private TemplateContext context;
+    @Override
+    public String lookupFunction(String name, String param) throws ExpansionException {
+      if (!functions.containsKey(name)) {
+        throw new ExpansionException(String.format("$(%s) not defined", name));
+      }
+      return functions.get(name).apply(param);
+    }
+  }
 
-  private Map<String, String> vars = new HashMap<>();
+  private TemplateContextImpl context;
 
   @Before
   public final void createContext() throws Exception  {
-    context = new TemplateContext() {
-        @Override
-        public String lookupVariable(String name)
-            throws ExpansionException {
-          // Not a Make variable. Let the shell handle the expansion.
-          if (name.startsWith("$")) {
-            return name;
-          }
-          if (!vars.containsKey(name)) {
-            throw new ExpansionException(String.format("$(%s) not defined", name));
-          }
-          return vars.get(name);
-        }
-      };
-
-    vars.put("SRCS", "src1 src2");
+    context = new TemplateContextImpl();
   }
 
-  private void assertExpansionEquals(String expected, String cmd)
-      throws ExpansionException {
-    assertThat(TemplateExpander.expand(cmd, context)).isEqualTo(expected);
+  private String expand(String value) throws ExpansionException {
+    return TemplateExpander.expand(value, context);
   }
 
-  private void assertExpansionFails(String expectedErrorSuffix, String cmd) {
+  private ExpansionException expansionFailure(String cmd) {
     try {
-      TemplateExpander.expand(cmd, context);
+      expand(cmd);
       fail("Expansion of " + cmd + " didn't fail as expected");
-    } catch (Exception e) {
-      assertThat(e).hasMessageThat().isEqualTo(expectedErrorSuffix);
+      throw new AssertionError();
+    } catch (ExpansionException e) {
+      return e;
     }
   }
 
   @Test
-  public void testExpansion() throws Exception {
-    vars.put("<", "src1");
-    vars.put("OUTS", "out1 out2");
-    vars.put("@", "out1");
-    vars.put("^", "src1 src2 dep1 dep2");
-    vars.put("@D", "outdir");
-    vars.put("BINDIR", "bindir");
-
-    assertExpansionEquals("src1 src2", "$(SRCS)");
-    assertExpansionEquals("src1", "$<");
-    assertExpansionEquals("out1 out2", "$(OUTS)");
-    assertExpansionEquals("out1", "$(@)");
-    assertExpansionEquals("out1", "$@");
-    assertExpansionEquals("out1,", "$@,");
-
-    assertExpansionEquals("src1 src2 out1 out2", "$(SRCS) $(OUTS)");
-
-    assertExpansionEquals("cmd", "cmd");
-    assertExpansionEquals("cmd src1 src2,", "cmd $(SRCS),");
-    assertExpansionEquals("label1 src1 src2,", "label1 $(SRCS),");
-    assertExpansionEquals(":label1 src1 src2,", ":label1 $(SRCS),");
+  public void testVariableExpansion() throws Exception {
+    context.vars.put("SRCS", "src1 src2");
+    context.vars.put("<", "src1");
+    context.vars.put("OUTS", "out1 out2");
+    context.vars.put("@", "out1");
+    context.vars.put("^", "src1 src2 dep1 dep2");
+    context.vars.put("@D", "outdir");
+    context.vars.put("BINDIR", "bindir");
+
+    assertThat(expand("$(SRCS)")).isEqualTo("src1 src2");
+    assertThat(expand("$<")).isEqualTo("src1");
+    assertThat(expand("$(OUTS)")).isEqualTo("out1 out2");
+    assertThat(expand("$(@)")).isEqualTo("out1");
+    assertThat(expand("$@")).isEqualTo("out1");
+    assertThat(expand("$@,")).isEqualTo("out1,");
+
+    assertThat(expand("$(SRCS) $(OUTS)")).isEqualTo("src1 src2 out1 out2");
+
+    assertThat(expand("cmd")).isEqualTo("cmd");
+    assertThat(expand("cmd $(SRCS),")).isEqualTo("cmd src1 src2,");
+    assertThat(expand("label1 $(SRCS),")).isEqualTo("label1 src1 src2,");
+    assertThat(expand(":label1 $(SRCS),")).isEqualTo(":label1 src1 src2,");
+  }
+
+  @Test
+  public void testUndefinedVariableExpansion() throws Exception {
+    assertThat(expansionFailure("$(foo)"))
+        .hasMessageThat().isEqualTo("$(foo) not defined");
+  }
+
+  @Test
+  public void testFunctionExpansion() throws Exception {
+    context.functions.put("foo", (String p) -> "FOO(" + p + ")");
+    context.vars.put("bar", "bar");
+
+    assertThat(expand("$(foo baz)")).isEqualTo("FOO(baz)");
+    assertThat(expand("$(bar) $(foo baz)")).isEqualTo("bar FOO(baz)");
+    assertThat(expand("xyz$(foo baz)zyx")).isEqualTo("xyzFOO(baz)zyx");
+  }
+
+  @Test
+  public void testFunctionExpansionThrows() throws Exception {
+    try {
+      TemplateExpander.expand("$(foo baz)", new TemplateContext() {
+        @Override
+        public String lookupVariable(String name) throws ExpansionException {
+          throw new ExpansionException(name);
+        }
 
+        @Override
+        public String lookupFunction(String name, String param) throws ExpansionException {
+          throw new ExpansionException(name + "(" + param + ")");
+        }
+      });
+      fail();
+    } catch (ExpansionException e) {
+      assertThat(e).hasMessageThat().isEqualTo("foo(baz)");
+    }
+  }
+
+  @Test
+  public void testUndefinedFunctionExpansion() throws Exception {
     // Note: $(location x) is considered an undefined variable;
-    assertExpansionFails("$(location label1) not defined",
-                         "$(location label1), $(SRCS),");
+    assertThat(expansionFailure("$(location label1), $(SRCS),"))
+        .hasMessageThat().isEqualTo("$(location) not defined");
+    assertThat(expansionFailure("$(basename file)"))
+        .hasMessageThat().isEqualTo("$(basename) not defined");
   }
 
   @Test
   public void testRecursiveExpansion() throws Exception {
     // Expansion is recursive: $(recursive) -> $(SRCS) -> "src1 src2"
-    vars.put("recursive", "$(SRCS)");
-    assertExpansionEquals("src1 src2", "$(recursive)");
+    context.vars.put("SRCS", "src1 src2");
+    context.vars.put("recursive", "$(SRCS)");
+    assertThat(expand("$(recursive)")).isEqualTo("src1 src2");
+  }
 
+  @Test
+  public void testRecursiveExpansionDoesNotSpanExpansionBoundaries() throws Exception {
     // Recursion does not span expansion boundaries:
     // $(recur2a)$(recur2b) --> "$" + "(SRCS)"  --/--> "src1 src2"
-    vars.put("recur2a", "$$");
-    vars.put("recur2b", "(SRCS)");
-    assertExpansionEquals("$(SRCS)", "$(recur2a)$(recur2b)");
+    context.vars.put("SRCS", "src1 src2");
+    context.vars.put("recur2a", "$$");
+    context.vars.put("recur2b", "(SRCS)");
+    assertThat(expand("$(recur2a)$(recur2b)")).isEqualTo("$(SRCS)");
+  }
+
+  @Test
+  public void testSelfInfiniteExpansionFailsGracefully() throws Exception {
+    context.vars.put("infinite", "$(infinite)");
+    assertThat(expansionFailure("$(infinite)")).hasMessageThat()
+        .isEqualTo("potentially unbounded recursion during expansion of '$(infinite)'");
   }
 
   @Test
-  public void testInfiniteRecursionFailsGracefully() throws Exception {
-    vars.put("infinite", "$(infinite)");
-    assertExpansionFails("potentially unbounded recursion during expansion "
-                         + "of '$(infinite)'",
-                         "$(infinite)");
-
-    vars.put("black", "$(white)");
-    vars.put("white", "$(black)");
-    assertExpansionFails("potentially unbounded recursion during expansion "
-                         + "of '$(black)'",
-                         "$(white) is the new $(black)");
+  public void testMutuallyInfiniteExpansionFailsGracefully() throws Exception {
+    context.vars.put("black", "$(white)");
+    context.vars.put("white", "$(black)");
+    assertThat(expansionFailure("$(white) is the new $(black)")).hasMessageThat()
+        .isEqualTo("potentially unbounded recursion during expansion of '$(black)'");
   }
 
   @Test
   public void testErrors() throws Exception {
-    assertExpansionFails("unterminated variable reference", "$(SRCS");
-    assertExpansionFails("unterminated $", "$");
+    assertThat(expansionFailure("$(SRCS")).hasMessageThat()
+        .isEqualTo("unterminated variable reference");
+    assertThat(expansionFailure("$")).hasMessageThat().isEqualTo("unterminated $");
 
     String suffix = "instead for \"Make\" variables, or escape the '$' as '$$' if you intended "
         + "this for the shell";
-    assertExpansionFails("'$file' syntax is not supported; use '$(file)' " + suffix,
-                         "for file in a b c;do echo $file;done");
-    assertExpansionFails("'${file%:.*8}' syntax is not supported; use '$(file%:.*8)' " + suffix,
-                         "${file%:.*8}");
+    assertThat(expansionFailure("for file in a b c;do echo $file;done")).hasMessageThat()
+        .isEqualTo("'$file' syntax is not supported; use '$(file)' " + suffix);
+    assertThat(expansionFailure("${file%:.*8}")).hasMessageThat()
+        .isEqualTo("'${file%:.*8}' syntax is not supported; use '$(file%:.*8)' " + suffix);
   }
 
   @Test
-  public void testShellVariables() throws Exception {
-    assertExpansionEquals("for file in a b c;do echo $file;done",
-        "for file in a b c;do echo $$file;done");
-    assertExpansionEquals("${file%:.*8}", "$${file%:.*8}");
-    assertExpansionFails("$(basename file) not defined", "$(basename file)");
-    assertExpansionEquals("$(basename file)", "$$(basename file)");
+  public void testDollarDollar() throws Exception {
+    assertThat(expand("for file in a b c;do echo $$file;done"))
+        .isEqualTo("for file in a b c;do echo $file;done");
+    assertThat(expand("$${file%:.*8}")).isEqualTo("${file%:.*8}");
+    assertThat(expand("$$(basename file)")).isEqualTo("$(basename file)");
   }
 }
diff --git a/src/test/java/com/google/devtools/build/lib/bazel/rules/genrule/GenRuleCommandSubstitutionTest.java b/src/test/java/com/google/devtools/build/lib/bazel/rules/genrule/GenRuleCommandSubstitutionTest.java
index c4c9a66978..8b2b1d4eaa 100644
--- a/src/test/java/com/google/devtools/build/lib/bazel/rules/genrule/GenRuleCommandSubstitutionTest.java
+++ b/src/test/java/com/google/devtools/build/lib/bazel/rules/genrule/GenRuleCommandSubstitutionTest.java
@@ -105,10 +105,10 @@ public class GenRuleCommandSubstitutionTest extends BuildViewTestCase {
     assertExpansionFails("$(locationz) not defined", "//test");
 
     genrule("$(locationz )");
-    assertExpansionFails("$(locationz ) not defined", "//test");
+    assertExpansionFails("$(locationz) not defined", "//test");
 
     genrule("$(locationz foo )");
-    assertExpansionFails("$(locationz foo ) not defined", "//test");
+    assertExpansionFails("$(locationz) not defined", "//test");
   }
 
   @Test
@@ -233,10 +233,10 @@ public class GenRuleCommandSubstitutionTest extends BuildViewTestCase {
     assertExpansionFails("$(locationsz) not defined", "//test");
 
     genrule("$(locationsz )");
-    assertExpansionFails("$(locationsz ) not defined", "//test");
+    assertExpansionFails("$(locationsz) not defined", "//test");
 
     genrule("$(locationsz foo )");
-    assertExpansionFails("$(locationsz foo ) not defined", "//test");
+    assertExpansionFails("$(locationsz) not defined", "//test");
   }
 
   @Test
@@ -447,8 +447,8 @@ public class GenRuleCommandSubstitutionTest extends BuildViewTestCase {
     assertNoEvents();
 
     genrule("$(basename file)");
-    assertExpansionFails("$(basename file) not defined", "//test");
-    assertContainsEvent("$(basename file) not defined");
+    assertExpansionFails("$(basename) not defined", "//test");
+    assertContainsEvent("$(basename) not defined");
   }
 
   @Test
diff --git a/src/test/java/com/google/devtools/build/lib/skylark/SkylarkRuleImplementationFunctionsTest.java b/src/test/java/com/google/devtools/build/lib/skylark/SkylarkRuleImplementationFunctionsTest.java
index c4f5f5af30..bc4b40e197 100644
--- a/src/test/java/com/google/devtools/build/lib/skylark/SkylarkRuleImplementationFunctionsTest.java
+++ b/src/test/java/com/google/devtools/build/lib/skylark/SkylarkRuleImplementationFunctionsTest.java
@@ -172,7 +172,7 @@ public class SkylarkRuleImplementationFunctionsTest extends SkylarkTestCase {
       setupSkylarkFunction(line);
       fail();
     } catch (EvalException e) {
-      assertThat(e).hasMessage(errorMsg);
+      assertThat(e).hasMessageThat().isEqualTo(errorMsg);
     }
   }