Verify macros

4 files changed, 44 insertions, 13 deletions

diff --git a/vespajlib/src/main/java/com/yahoo/tensor/evaluation/EvaluationContext.java b/vespajlib/src/main/java/com/yahoo/tensor/evaluation/EvaluationContext.java
index e18b77a0434..3fb94f1251b 100644
--- a/vespajlib/src/main/java/com/yahoo/tensor/evaluation/EvaluationContext.java
+++ b/vespajlib/src/main/java/com/yahoo/tensor/evaluation/EvaluationContext.java
@@ -3,7 +3,6 @@ package com.yahoo.tensor.evaluation;
 
 import com.google.common.annotations.Beta;
 import com.yahoo.tensor.Tensor;
-import com.yahoo.tensor.TensorType;
 
 /**
  * An evaluation context which is passed down to all nested functions during evaluation.
@@ -11,15 +10,7 @@ import com.yahoo.tensor.TensorType;
  * @author bratseth
  */
 @Beta
-public interface EvaluationContext {
-
-    /**
-     * Returns tye type of the tensor with this name.
-     *
-     * @return returns the type of the tensor which will be returned by calling getTensor(name)
-     *         or null if getTensor will return null.
-     */
-    TensorType getTensorType(String name);
+public interface EvaluationContext extends TypeContext {
 
     /** Returns the tensor bound to this name, or null if none */
     Tensor getTensor(String name);
diff --git a/vespajlib/src/main/java/com/yahoo/tensor/evaluation/MapEvaluationContext.java b/vespajlib/src/main/java/com/yahoo/tensor/evaluation/MapEvaluationContext.java
index 6bdfe8f19b6..9fe6b7d053f 100644
--- a/vespajlib/src/main/java/com/yahoo/tensor/evaluation/MapEvaluationContext.java
+++ b/vespajlib/src/main/java/com/yahoo/tensor/evaluation/MapEvaluationContext.java
@@ -20,7 +20,7 @@ public class MapEvaluationContext implements EvaluationContext {
     public void put(String name, Tensor tensor) { bindings.put(name, tensor); }
 
     @Override
-    public TensorType getTensorType(String name) {
+    public TensorType getType(String name) {
         Tensor tensor = bindings.get(name);
         if (tensor == null) return null;
         return tensor.type();
diff --git a/vespajlib/src/main/java/com/yahoo/tensor/evaluation/TypeContext.java b/vespajlib/src/main/java/com/yahoo/tensor/evaluation/TypeContext.java
new file mode 100644
index 00000000000..9b2e81f0b6d
--- /dev/null
+++ b/vespajlib/src/main/java/com/yahoo/tensor/evaluation/TypeContext.java
@@ -0,0 +1,18 @@
+package com.yahoo.tensor.evaluation;// Copyright 2018 Yahoo Holdings. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root.
+
+import com.yahoo.tensor.TensorType;
+
+/**
+ * @author bratseth
+ */
+public interface TypeContext {
+
+    /**
+     * Returns tye type of the tensor with this name.
+     *
+     * @return returns the type of the tensor which will be returned by calling getTensor(name)
+     *         or null if getTensor will return null.
+     */
+    TensorType getType(String name);
+
+}
diff --git a/vespajlib/src/main/java/com/yahoo/tensor/evaluation/VariableTensor.java b/vespajlib/src/main/java/com/yahoo/tensor/evaluation/VariableTensor.java
index 6c149724aca..34beb465d4c 100644
--- a/vespajlib/src/main/java/com/yahoo/tensor/evaluation/VariableTensor.java
+++ b/vespajlib/src/main/java/com/yahoo/tensor/evaluation/VariableTensor.java
@@ -10,6 +10,7 @@ import com.yahoo.tensor.functions.ToStringContext;
 
 import java.util.Collections;
 import java.util.List;
+import java.util.Optional;
 
 /**
  * A tensor variable name which resolves to a tensor in the context at evaluation time
@@ -20,9 +21,17 @@ import java.util.List;
 public class VariableTensor extends PrimitiveTensorFunction {
 
     private final String name;
+    private final Optional<TensorType> requiredType;
 
     public VariableTensor(String name) {
         this.name = name;
+        this.requiredType = Optional.empty();
+    }
+
+    /** A variable tensor which must be compatible with the given type */
+    public VariableTensor(String name, TensorType requiredType) {
+        this.name = name;
+        this.requiredType = Optional.of(requiredType);
     }
 
     @Override
@@ -35,11 +44,19 @@ public class VariableTensor extends PrimitiveTensorFunction {
     public PrimitiveTensorFunction toPrimitive() { return this; }
 
     @Override
-    public TensorType type(EvaluationContext context) { return context.getTensorType(name); }
+    public TensorType type(TypeContext context) {
+        TensorType givenType = context.getType(name);
+        if (givenType == null) return null;
+        verifyType(givenType);
+        return givenType;
+    }
 
     @Override
     public Tensor evaluate(EvaluationContext context) {
-        return context.getTensor(name);
+        Tensor tensor = context.getTensor(name);
+        if (tensor == null) return null;
+        verifyType(tensor.type());
+        return tensor;
     }
 
     @Override
@@ -47,4 +64,9 @@ public class VariableTensor extends PrimitiveTensorFunction {
         return name;
     }
 
+    private void verifyType(TensorType givenType) {
+        if (requiredType.isPresent() && ! givenType.isAssignableTo(requiredType.get()))
+            throw new IllegalArgumentException("Variable '" + name + "' must be compatible with " +
+                                               requiredType.get() + " but was " + givenType);
+    }
 }