Encapsulate function references

5 files changed, 84 insertions, 22 deletions

diff --git a/model-inference/src/main/java/ai/vespa/models/evaluation/FunctionReference.java b/model-inference/src/main/java/ai/vespa/models/evaluation/FunctionReference.java
new file mode 100644
index 00000000000..115303e0fdd
--- /dev/null
+++ b/model-inference/src/main/java/ai/vespa/models/evaluation/FunctionReference.java
@@ -0,0 +1,63 @@
+// Copyright 2018 Yahoo Holdings. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root.
+package ai.vespa.models.evaluation;
+
+import java.util.Optional;
+import java.util.regex.Matcher;
+import java.util.regex.Pattern;
+
+/**
+ * A reference to a function.
+ * The function may be
+ * - free: Callable from users of models, or
+ * - bound: Representing a specific invocation from another ranking expression.
+ * In bound functions, any arguments are replaced by the values supplied in the function invocation.
+ * Function references has a serial form (textual representation) used in ranking expressions received in ranking
+ * expression configurations.
+ *
+ * This is immutable.
+ *
+ * @author bratseth
+ */
+class FunctionReference {
+
+    private static final Pattern referencePattern =
+            Pattern.compile("rankingExpression\\(([a-zA-Z0-9_]+)(@[a-f0-9]+\\.[a-f0-9]+)?\\)(\\.rankingScript)?");
+
+    /** The name of the function referenced */
+    private final String name;
+
+    /** The id of the specific invocation of the function, or null if it is free */
+    private final String instance;
+
+    private FunctionReference(String name, String instance) {
+        this.name = name;
+        this.instance = instance;
+    }
+
+    /** Returns the name of the function referenced */
+    String functionName() { return name; }
+
+    boolean isFree() {
+        return instance == null;
+    }
+
+    String serialForm() {
+        return "rankingExpression(" + name + (instance != null ? instance : "") + ")";
+    }
+
+    @Override
+    public String toString() { return serialForm(); }
+
+    // TODO: Equals and hashcode
+
+    /** Returns a function reference from the given serial form, or empty if the string is not a valid reference */
+    static Optional<FunctionReference> fromSerial(String serialForm) {
+        Matcher expressionMatcher = referencePattern.matcher(serialForm);
+        if ( ! expressionMatcher.matches()) return Optional.empty();
+
+        String name = expressionMatcher.group(1);
+        String instance = expressionMatcher.group(2);
+        return Optional.of(new FunctionReference(name, instance));
+    }
+
+}
diff --git a/model-inference/src/main/java/ai/vespa/models/evaluation/LazyArrayContext.java b/model-inference/src/main/java/ai/vespa/models/evaluation/LazyArrayContext.java
index c38018dfcf5..29c934827bf 100644
--- a/model-inference/src/main/java/ai/vespa/models/evaluation/LazyArrayContext.java
+++ b/model-inference/src/main/java/ai/vespa/models/evaluation/LazyArrayContext.java
@@ -18,6 +18,7 @@ import java.util.Arrays;
 import java.util.LinkedHashSet;
 import java.util.Map;
 import java.util.Set;
+import java.util.stream.Collectors;
 
 /**
  * An array context supporting functions invocations implemented as lazy values.
diff --git a/model-inference/src/main/java/ai/vespa/models/evaluation/Model.java b/model-inference/src/main/java/ai/vespa/models/evaluation/Model.java
index 54ab356fd0c..5ce5a1754a9 100644
--- a/model-inference/src/main/java/ai/vespa/models/evaluation/Model.java
+++ b/model-inference/src/main/java/ai/vespa/models/evaluation/Model.java
@@ -4,6 +4,7 @@ package ai.vespa.models.evaluation;
 import com.google.common.collect.ImmutableList;
 import com.google.common.collect.ImmutableMap;
 import com.yahoo.searchlib.rankingexpression.ExpressionFunction;
+import com.yahoo.searchlib.rankingexpression.evaluation.ContextIndex;
 
 import java.util.Collection;
 import java.util.Collections;
@@ -37,25 +38,29 @@ public class Model {
         this.name = name;
         this.functions = ImmutableList.copyOf(functions);
 
-        ImmutableMap.Builder<String, ExpressionFunction> functionsBuilder = new ImmutableMap.Builder<>();
-        for (ExpressionFunction function : referencedFunctions)
-            functionsBuilder.put(function.getName(), optimize(function));
-        this.referencedFunctions = functionsBuilder.build();
-
         ImmutableMap.Builder<String, LazyArrayContext> contextBuilder = new ImmutableMap.Builder<>();
         for (ExpressionFunction function : functions) {
             try {
-                contextBuilder.put(function.getName(), new LazyArrayContext(function.getBody(), this.referencedFunctions, this));
+                contextBuilder.put(function.getName(),
+                                   new LazyArrayContext(function.getBody(),
+                                                        referencedFunctions.stream().collect(Collectors.toMap(e -> e.getName(), e -> e)),
+                                                        this));
             }
             catch (RuntimeException e) {
                 throw new IllegalArgumentException("Could not prepare an evaluation context for " + function, e);
             }
         }
         this.contextPrototypes = contextBuilder.build();
+
+        ImmutableMap.Builder<String, ExpressionFunction> functionsBuilder = new ImmutableMap.Builder<>();
+        for (ExpressionFunction function : referencedFunctions)
+            functionsBuilder.put(function.getName(), optimize(function,
+                                                              contextPrototypes.get(FunctionReference.fromSerial(function.getName()).get().serialForm())));
+        this.referencedFunctions = functionsBuilder.build();
     }
 
     /** Returns an optimized version of the given function */
-    private ExpressionFunction optimize(ExpressionFunction function) {
+    private ExpressionFunction optimize(ExpressionFunction function, ContextIndex context) {
         return function; // TODO
     }
 
diff --git a/model-inference/src/main/java/ai/vespa/models/evaluation/RankProfilesConfigImporter.java b/model-inference/src/main/java/ai/vespa/models/evaluation/RankProfilesConfigImporter.java
index bd0453f2826..8fe2727b4df 100644
--- a/model-inference/src/main/java/ai/vespa/models/evaluation/RankProfilesConfigImporter.java
+++ b/model-inference/src/main/java/ai/vespa/models/evaluation/RankProfilesConfigImporter.java
@@ -10,8 +10,7 @@ import java.util.ArrayList;
 import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
-import java.util.regex.Matcher;
-import java.util.regex.Pattern;
+import java.util.Optional;
 
 /**
  * Converts RankProfilesConfig instances to RankingExpressions for evaluation
@@ -20,9 +19,6 @@ import java.util.regex.Pattern;
  */
 class RankProfilesConfigImporter {
 
-    private static final Pattern expressionPattern =
-            Pattern.compile("rankingExpression\\(([a-zA-Z0-9_]+)(@[a-f0-9]+\\.[a-f0-9]+)?\\)\\.rankingScript");
-
     /**
      * Returns a map of the models contained in this config, indexed on name.
      * The map is modifiable and owned by the caller.
@@ -47,19 +43,16 @@ class RankProfilesConfigImporter {
         ExpressionFunction firstPhase = null;
         ExpressionFunction secondPhase = null;
         for (RankProfilesConfig.Rankprofile.Fef.Property property : profile.fef().property()) {
-            Matcher expressionMatcher = expressionPattern.matcher(property.name());
-            if ( expressionMatcher.matches()) {
-                String name = expressionMatcher.group(1);
-                String instance = expressionMatcher.group(2);
+            Optional<FunctionReference> reference = FunctionReference.fromSerial(property.name());
+            if ( reference.isPresent()) {
                 List<String> arguments = new ArrayList<>(); // TODO: Arguments?
-                RankingExpression expression = new RankingExpression(name, property.value());
+                RankingExpression expression = new RankingExpression(reference.get().functionName(), property.value());
 
-                if (instance == null) // free function;  make available in model under configured name
-                    functions.add(new ExpressionFunction(name, arguments, expression)); //
+                if (reference.get().isFree()) // make available in model under configured name
+                    functions.add(new ExpressionFunction(reference.get().functionName(), arguments, expression)); //
 
                 // Make all functions, bound or not available under the name they are referenced by in expressions
-                boundFunctions.add(new ExpressionFunction("rankingExpression(" + name + (instance != null ? instance : "") + ")",
-                                                          arguments, expression));
+                boundFunctions.add(new ExpressionFunction(reference.get().serialForm(), arguments, expression));
             }
             else if (property.name().equals("vespa.rank.firstphase")) { // Include in addition to macros
                 firstPhase = new ExpressionFunction("firstphase", new ArrayList<>(),
diff --git a/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/maintenance/NodeRebooter.java b/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/maintenance/NodeRebooter.java
index cc4834d97a6..6b538946e09 100644
--- a/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/maintenance/NodeRebooter.java
+++ b/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/maintenance/NodeRebooter.java
@@ -35,7 +35,7 @@ public class NodeRebooter extends Maintainer {
 
     @Override
     protected void maintain() {
-        // Reboot candidates: Nodes in long-term states, which we know an safely orchestrate a reboot
+        // Reboot candidates: Nodes in long-term states, which we know can safely orchestrate a reboot
         List<Node> rebootCandidates = nodeRepository().getNodes(NodeType.tenant, Node.State.active, Node.State.ready);
         rebootCandidates.addAll(nodeRepository().getNodes(NodeType.proxy, Node.State.active, Node.State.ready));
         rebootCandidates.addAll(nodeRepository().getNodes(NodeType.host, Node.State.active, Node.State.ready));