1 files changed, 5 insertions, 9 deletions
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
index d2888200e0..96bb24ffe8 100644
--- a/src/main/java/com/google/devtools/build/lib/analysis/LocationTemplateContext.java
+++ b/src/main/java/com/google/devtools/build/lib/analysis/LocationTemplateContext.java
@@ -20,7 +20,6 @@ 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;
@@ -49,8 +48,7 @@ import javax.annotation.Nullable;
  */
 final class LocationTemplateContext implements TemplateContext {
   private final TemplateContext delegate;
-  private final Function<String, String> locationFunction;
-  private final Function<String, String> locationsFunction;
+  private final ImmutableMap<String, Function<String, String>> functions;
 
   private LocationTemplateContext(
       TemplateContext delegate,
@@ -58,8 +56,7 @@ final class LocationTemplateContext implements TemplateContext {
       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);
+    this.functions = LocationExpander.allLocationFunctions(root, locationMap, execPaths);
   }
 
   private LocationTemplateContext(
@@ -98,10 +95,9 @@ final class LocationTemplateContext implements TemplateContext {
   @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);
+      Function<String, String> f = functions.get(name);
+      if (f != null) {
+        return f.apply(param);
       }
     } catch (IllegalStateException e) {
       throw new ExpansionException(e.getMessage(), e);