Move types to ExpressionFunction

1 files changed, 40 insertions, 4 deletions

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);
     }
 
     /**