Include argument type information in functions

1 files changed, 7 insertions, 1 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 f6502a9801d..787b857839d 100755
--- a/searchlib/src/main/java/com/yahoo/searchlib/rankingexpression/ExpressionFunction.java
+++ b/searchlib/src/main/java/com/yahoo/searchlib/rankingexpression/ExpressionFunction.java
@@ -11,6 +11,7 @@ import com.yahoo.text.Utf8;
 
 import java.security.MessageDigest;
 import java.security.NoSuchAlgorithmException;
+import java.util.ArrayList;
 import java.util.Collections;
 import java.util.Deque;
 import java.util.HashMap;
@@ -97,7 +98,12 @@ public class ExpressionFunction {
         return new ExpressionFunction(name, arguments, body, argumentTypes, Optional.of(returnType));
     }
 
-    public ExpressionFunction withArgumentTypes(Map<String, TensorType> argumentTypes) {
+    /** Returns a copy of this with the given argument and argument type added */
+    public ExpressionFunction withArgument(String argument, TensorType type) {
+        List<String> arguments = new ArrayList<>(this.arguments);
+        arguments.add(argument);
+        Map<String, TensorType> argumentTypes = new HashMap<>(this.argumentTypes);
+        argumentTypes.put(argument, type);
         return new ExpressionFunction(name, arguments, body, argumentTypes, returnType);
     }