Move types to ExpressionFunction

4 files changed, 51 insertions, 21 deletions

diff --git a/config-model/src/main/java/com/yahoo/searchdefinition/RankProfile.java b/config-model/src/main/java/com/yahoo/searchdefinition/RankProfile.java
index 7d8860204f2..937151c0d3a 100644
--- a/config-model/src/main/java/com/yahoo/searchdefinition/RankProfile.java
+++ b/config-model/src/main/java/com/yahoo/searchdefinition/RankProfile.java
@@ -534,7 +534,7 @@ public class RankProfile implements Serializable, Cloneable {
 
     /** Adds a function and returns it */
     public RankingExpressionFunction addFunction(ExpressionFunction function, boolean inline) {
-        RankingExpressionFunction rankingExpressionFunction = new RankingExpressionFunction(function, inline, Optional.empty());
+        RankingExpressionFunction rankingExpressionFunction = new RankingExpressionFunction(function, inline);
         functions.put(function.getName(), rankingExpressionFunction);
         return rankingExpressionFunction;
     }
@@ -694,7 +694,7 @@ public class RankProfile implements Serializable, Cloneable {
         for (Map.Entry<String, RankingExpressionFunction> entry : functions.entrySet()) {
             RankingExpressionFunction rankingExpressionFunction = entry.getValue();
             RankingExpression compiled = compile(rankingExpressionFunction.function().getBody(), queryProfiles, importedModels, getConstants(), inlineFunctions, expressionTransforms);
-            compiledFunctions.put(entry.getKey(), rankingExpressionFunction.withBody(compiled));
+            compiledFunctions.put(entry.getKey(), rankingExpressionFunction.withExpression(compiled));
         }
         return compiledFunctions;
     }
@@ -897,21 +897,18 @@ public class RankProfile implements Serializable, Cloneable {
     /** A function in a rank profile */
     public static class RankingExpressionFunction {
 
-        private final ExpressionFunction function;
+        private ExpressionFunction function;
 
         /** True if this should be inlined into calling expressions. Useful for very cheap functions. */
         private final boolean inline;
 
-        private Optional<TensorType> type;
-
-        public RankingExpressionFunction(ExpressionFunction function, boolean inline, Optional<TensorType> type) {
+        public RankingExpressionFunction(ExpressionFunction function, boolean inline) {
             this.function = function;
             this.inline = inline;
-            this.type = type;
         }
 
-        public void setType(TensorType type) {
-            this.type = Optional.of(type);
+        public void setReturnType(TensorType type) {
+            this.function = function.withReturnType(type);
         }
 
         public ExpressionFunction function() { return function; }
@@ -920,11 +917,8 @@ public class RankProfile implements Serializable, Cloneable {
             return inline && function.arguments().isEmpty(); // only inline no-arg functions;
         }
 
-        /** Returns the type this produces, or empty if not resolved */
-        public Optional<TensorType> type() { return type; }
-
-        public RankingExpressionFunction withBody(RankingExpression expression) {
-            return new RankingExpressionFunction(function.withBody(expression), inline, type);
+        public RankingExpressionFunction withExpression(RankingExpression expression) {
+            return new RankingExpressionFunction(function.withBody(expression), inline);
         }
 
         @Override
diff --git a/config-model/src/main/java/com/yahoo/searchdefinition/derived/RawRankProfile.java b/config-model/src/main/java/com/yahoo/searchdefinition/derived/RawRankProfile.java
index c3a2673ef2b..f575d1178c0 100644
--- a/config-model/src/main/java/com/yahoo/searchdefinition/derived/RawRankProfile.java
+++ b/config-model/src/main/java/com/yahoo/searchdefinition/derived/RawRankProfile.java
@@ -208,8 +208,8 @@ public class RawRankProfile implements RankProfilesConfig.Producer {
                 String expressionString = e.getValue().function().getBody().getRoot().toString(new StringBuilder(), context, null, null).toString();
                 context.addFunctionSerialization(RankingExpression.propertyName(e.getKey()), expressionString);
 
-                if (e.getValue().type().isPresent())
-                    context.addFunctionTypeSerialization(RankingExpression.propertyTypeName(e.getKey()), e.getValue().type().get().toString());
+                if (e.getValue().function().returnType().isPresent())
+                    context.addFunctionTypeSerialization(RankingExpression.propertyTypeName(e.getKey()), e.getValue().function().returnType().get().toString());
                 else if (e.getValue().function().arguments().isEmpty())
                     throw new IllegalStateException("Type of function '" + e.getKey() + "' is not resolved");
             }
diff --git a/config-model/src/main/java/com/yahoo/searchdefinition/processing/RankingExpressionTypeResolver.java b/config-model/src/main/java/com/yahoo/searchdefinition/processing/RankingExpressionTypeResolver.java
index f502af922e9..4c8b5910b78 100644
--- a/config-model/src/main/java/com/yahoo/searchdefinition/processing/RankingExpressionTypeResolver.java
+++ b/config-model/src/main/java/com/yahoo/searchdefinition/processing/RankingExpressionTypeResolver.java
@@ -64,7 +64,7 @@ public class RankingExpressionTypeResolver extends Processor {
             TensorType type = resolveType(function.getValue().function().getBody(),
                                           "function '" + function.getKey() + "'",
                                           context);
-            function.getValue().setType(type);
+            function.getValue().setReturnType(type);
         }
 
         if (validate) {
diff --git a/searchlib/src/main/java/com/yahoo/searchlib/rankingexpression/ExpressionFunction.java b/searchlib/src/main/java/com/yahoo/searchlib/rankingexpression/ExpressionFunction.java
index da34ab8822d..63d3f9df663 100755
--- a/searchlib/src/main/java/com/yahoo/searchlib/rankingexpression/ExpressionFunction.java
+++ b/searchlib/src/main/java/com/yahoo/searchlib/rankingexpression/ExpressionFunction.java
@@ -2,8 +2,11 @@
 package com.yahoo.searchlib.rankingexpression;
 
 import com.google.common.collect.ImmutableList;
+import com.google.common.collect.ImmutableMap;
+import com.yahoo.log.event.Collection;
 import com.yahoo.searchlib.rankingexpression.rule.ExpressionNode;
 import com.yahoo.searchlib.rankingexpression.rule.SerializationContext;
+import com.yahoo.tensor.TensorType;
 import com.yahoo.text.Utf8;
 
 import java.security.MessageDigest;
@@ -13,9 +16,14 @@ import java.util.Deque;
 import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
+import java.util.Objects;
+import java.util.Optional;
 
 /**
- * A function defined by a ranking expression
+ * A function defined by a ranking expression, optionally containing type information
+ * for inputs and outputs.
+ *
+ * Immutable, but note that ranking expressions are *not* immutable.
  *
  * @author Simon Thoresen Hult
  * @author bratseth
@@ -24,8 +32,13 @@ public class ExpressionFunction {
 
     private final String name;
     private final ImmutableList<String> arguments;
+
+    /** Types of the inputs, if known. The keys here is any subset (including empty and identity) of the argument list */
+    private final ImmutableMap<String, TensorType> argumentTypes;
     private final RankingExpression body;
 
+    private final Optional<TensorType> returnType;
+
     /**
      * Constructs a new function with no arguments
      *
@@ -44,9 +57,18 @@ public class ExpressionFunction {
      * @param body the ranking expression that defines this function
      */
     public ExpressionFunction(String name, List<String> arguments, RankingExpression body) {
-        this.name = name;
+        this(name, arguments, body, ImmutableMap.of(), Optional.empty());
+    }
+
+    private ExpressionFunction(String name, List<String> arguments, RankingExpression body,
+                               Map<String, TensorType> argumentTypes, Optional<TensorType> returnType) {
+        this.name = Objects.requireNonNull(name, "name cannot be null");
         this.arguments = arguments==null ? ImmutableList.of() : ImmutableList.copyOf(arguments);
-        this.body = body;
+        this.body = Objects.requireNonNull(body, "body cannot be null");
+        if ( ! this.arguments.containsAll(argumentTypes.keySet()))
+            throw new IllegalArgumentException("Argument type keys must be a subset of the argument keys");
+        this.argumentTypes = ImmutableMap.copyOf(argumentTypes);
+        this.returnType = Objects.requireNonNull(returnType, "returnType cannot be null");
     }
 
     public String getName() { return name; }
@@ -56,9 +78,23 @@ public class ExpressionFunction {
 
     public RankingExpression getBody() { return body; }
 
+    /** Returns the types of the arguments of this, if specified. The keys of this may be any subset of the arguments */
+    public Map<String, TensorType> argumentTypes() { return argumentTypes; }
+
+    /** Returns the return type of this, or empty if not specified */
+    public Optional<TensorType> returnType() { return returnType; }
+
     /** Returns a copy of this with the body changed to the given value */
     public ExpressionFunction withBody(RankingExpression body) {
-        return new ExpressionFunction(name, arguments, body);
+        return new ExpressionFunction(name, arguments, body, argumentTypes, returnType);
+    }
+
+    public ExpressionFunction withReturnType(TensorType returnType) {
+        return new ExpressionFunction(name, arguments, body, argumentTypes, Optional.of(returnType));
+    }
+
+    public ExpressionFunction withArgumentTypes(Map<String, TensorType> argumentTypes) {
+        return new ExpressionFunction(name, arguments, body, argumentTypes, returnType);
     }
 
     /**