Return function types without evaluating

4 files changed, 25 insertions, 27 deletions

diff --git a/model-evaluation/src/main/java/ai/vespa/models/evaluation/LazyArrayContext.java b/model-evaluation/src/main/java/ai/vespa/models/evaluation/LazyArrayContext.java
index 78b30f0c873..4d1b5a97583 100644
--- a/model-evaluation/src/main/java/ai/vespa/models/evaluation/LazyArrayContext.java
+++ b/model-evaluation/src/main/java/ai/vespa/models/evaluation/LazyArrayContext.java
@@ -4,7 +4,6 @@ package ai.vespa.models.evaluation;
 import com.google.common.collect.ImmutableMap;
 import com.google.common.collect.ImmutableSet;
 import com.yahoo.searchlib.rankingexpression.ExpressionFunction;
-import com.yahoo.searchlib.rankingexpression.RankingExpression;
 import com.yahoo.searchlib.rankingexpression.Reference;
 import com.yahoo.searchlib.rankingexpression.evaluation.Context;
 import com.yahoo.searchlib.rankingexpression.evaluation.ContextIndex;
@@ -31,22 +30,22 @@ public final class LazyArrayContext extends Context implements ContextIndex {
 
     public final static Value defaultContextValue = DoubleValue.zero;
 
+    private final ExpressionFunction function;
+
     private final IndexedBindings indexedBindings;
 
-    private LazyArrayContext(IndexedBindings indexedBindings) {
+    private LazyArrayContext(ExpressionFunction function, IndexedBindings indexedBindings) {
+        this.function = function;
         this.indexedBindings = indexedBindings.copy(this);
     }
 
-    /**
-     * Create a fast lookup, lazy context for an expression.
-     *
-     * @param expression the expression to create a context for
-     */
-    LazyArrayContext(RankingExpression expression,
-                     Map<FunctionReference, ExpressionFunction> functions,
+    /** Create a fast lookup, lazy context for a function */
+    LazyArrayContext(ExpressionFunction function,
+                     Map<FunctionReference, ExpressionFunction> referencedFunctions,
                      List<Constant> constants,
                      Model model) {
-        this.indexedBindings = new IndexedBindings(expression, functions, constants, this, model);
+        this.function = function;
+        this.indexedBindings = new IndexedBindings(function, referencedFunctions, constants, this, model);
     }
 
     /**
@@ -76,7 +75,6 @@ public final class LazyArrayContext extends Context implements ContextIndex {
 
     @Override
     public TensorType getType(Reference reference) {
-        // TODO: Add type information so we do not need to evaluate to get this
         return get(requireIndexOf(reference.toString())).type();
     }
 
@@ -128,7 +126,7 @@ public final class LazyArrayContext extends Context implements ContextIndex {
      * in a different thread or for re-binding free variables.
      */
     LazyArrayContext copy() {
-        return new LazyArrayContext(indexedBindings);
+        return new LazyArrayContext(function, indexedBindings);
     }
 
     private static class IndexedBindings {
@@ -154,15 +152,15 @@ public final class LazyArrayContext extends Context implements ContextIndex {
          * Creates indexed bindings for the given expressions.
          * The given expression and functions may be inspected but cannot be stored.
          */
-        IndexedBindings(RankingExpression expression,
-                        Map<FunctionReference, ExpressionFunction> functions,
+        IndexedBindings(ExpressionFunction function,
+                        Map<FunctionReference, ExpressionFunction> referencedFunctions,
                         List<Constant> constants,
                         LazyArrayContext owner,
                         Model model) {
             // 1. Determine and prepare bind targets
             Set<String> bindTargets = new LinkedHashSet<>();
             Set<String> arguments = new LinkedHashSet<>(); // Arguments: Bind targets which need to be bound before invocation
-            extractBindTargets(expression.getRoot(), functions, bindTargets, arguments);
+            extractBindTargets(function.getBody().getRoot(), referencedFunctions, bindTargets, arguments);
 
             this.arguments = ImmutableSet.copyOf(arguments);
             values = new Value[bindTargets.size()];
@@ -183,10 +181,10 @@ public final class LazyArrayContext extends Context implements ContextIndex {
                     values[index] = new TensorValue(constant.value());
             }
 
-            for (Map.Entry<FunctionReference, ExpressionFunction> function : functions.entrySet()) {
-                Integer index = nameToIndex.get(function.getKey().serialForm());
+            for (Map.Entry<FunctionReference, ExpressionFunction> referencedFunction : referencedFunctions.entrySet()) {
+                Integer index = nameToIndex.get(referencedFunction.getKey().serialForm());
                 if (index != null) // Referenced in this, so bind it
-                    values[index] = new LazyValue(function.getKey(), owner, model);
+                    values[index] = new LazyValue(referencedFunction.getKey(), owner, model);
             }
         }
 
diff --git a/model-evaluation/src/main/java/ai/vespa/models/evaluation/LazyValue.java b/model-evaluation/src/main/java/ai/vespa/models/evaluation/LazyValue.java
index 4a1ee22d288..a7b536ba911 100644
--- a/model-evaluation/src/main/java/ai/vespa/models/evaluation/LazyValue.java
+++ b/model-evaluation/src/main/java/ai/vespa/models/evaluation/LazyValue.java
@@ -42,7 +42,7 @@ class LazyValue extends Value {
 
     @Override
     public TensorType type() {
-        return computedValue().type(); // TODO: Keep type information in this/ExpressionFunction to avoid computing here
+        return model.requireReferencedFunction(function).returnType().get();
     }
 
     @Override
diff --git a/model-evaluation/src/main/java/ai/vespa/models/evaluation/Model.java b/model-evaluation/src/main/java/ai/vespa/models/evaluation/Model.java
index fda1ae935ca..37c2d7961a8 100644
--- a/model-evaluation/src/main/java/ai/vespa/models/evaluation/Model.java
+++ b/model-evaluation/src/main/java/ai/vespa/models/evaluation/Model.java
@@ -55,7 +55,7 @@ public class Model {
         ImmutableMap.Builder<String, LazyArrayContext> contextBuilder = new ImmutableMap.Builder<>();
         for (Map.Entry<FunctionReference, ExpressionFunction> function : functions.entrySet()) {
             try {
-                LazyArrayContext context = new LazyArrayContext(function.getValue().getBody(), referencedFunctions, constants, this);
+                LazyArrayContext context = new LazyArrayContext(function.getValue(), referencedFunctions, constants, this);
                 contextBuilder.put(function.getValue().getName(), context);
                 for (String argument : context.arguments()) {
                     if (function.getValue().argumentTypes().get(argument) == null)
diff --git a/searchlib/src/main/java/com/yahoo/searchlib/rankingexpression/Reference.java b/searchlib/src/main/java/com/yahoo/searchlib/rankingexpression/Reference.java
index c42d9ecc37f..cd5f42ac05c 100644
--- a/searchlib/src/main/java/com/yahoo/searchlib/rankingexpression/Reference.java
+++ b/searchlib/src/main/java/com/yahoo/searchlib/rankingexpression/Reference.java
@@ -67,6 +67,13 @@ public class Reference extends TypeContext.Name {
     }
 
     /**
+     * Returns whether this is a simple identifier - no arguments or output
+     */
+    public boolean isIdentifier() {
+        return this.arguments.expressions().size() == 0 && output == null;
+    }
+
+    /**
      * A <i>simple feature reference</i> is a reference with a single identifier argument
      * (and an optional output).
      */
@@ -97,13 +104,6 @@ public class Reference extends TypeContext.Name {
         }
     }
 
-    /**
-     * Returns whether this is a simple identifier - no arguments or output
-     */
-    public boolean isIdentifier() {
-        return this.arguments.expressions().size() == 0 && output == null;
-    }
-
     public Reference withArguments(Arguments arguments) {
         return new Reference(name(), arguments, output);
     }