Propagate binding context to/from tensor functions

2 files changed, 43 insertions, 8 deletions

diff --git a/vespajlib/src/main/java/com/yahoo/tensor/functions/Generate.java b/vespajlib/src/main/java/com/yahoo/tensor/functions/Generate.java
index e5095178be7..ac6621ce78b 100644
--- a/vespajlib/src/main/java/com/yahoo/tensor/functions/Generate.java
+++ b/vespajlib/src/main/java/com/yahoo/tensor/functions/Generate.java
@@ -91,7 +91,7 @@ public class Generate<NAMETYPE extends Name> extends PrimitiveTensorFunction<NAM
     public Tensor evaluate(EvaluationContext<NAMETYPE> context) {
         Tensor.Builder builder = Tensor.Builder.of(type);
         IndexedTensor.Indexes indexes = IndexedTensor.Indexes.of(dimensionSizes(type));
-        GenerateContext generateContext = new GenerateContext(type, context);
+        GenerateEvaluationContext generateContext = new GenerateEvaluationContext(type, context);
         for (int i = 0; i < indexes.size(); i++) {
             indexes.next();
             builder.cell(generateContext.apply(indexes), indexes.indexesForReading());
@@ -113,7 +113,7 @@ public class Generate<NAMETYPE extends Name> extends PrimitiveTensorFunction<NAM
         if (freeGenerator != null)
             return freeGenerator.toString();
         else
-            return boundGenerator.toString(context);
+            return boundGenerator.toString(new GenerateToStringContext(context));
     }
 
     /**
@@ -121,19 +121,18 @@ public class Generate<NAMETYPE extends Name> extends PrimitiveTensorFunction<NAM
      * This returns all the current index values as variables and falls back to delivering from the given
      * evaluation context.
      */
-    private class GenerateContext implements EvaluationContext<NAMETYPE> {
+    private class GenerateEvaluationContext implements EvaluationContext<NAMETYPE> {
 
         private final TensorType type;
         private final EvaluationContext<NAMETYPE> context;
 
         private IndexedTensor.Indexes indexes;
 
-        GenerateContext(TensorType type, EvaluationContext<NAMETYPE> context) {
+        GenerateEvaluationContext(TensorType type, EvaluationContext<NAMETYPE> context) {
             this.type = type;
             this.context = context;
         }
 
-        @SuppressWarnings("unchecked")
         double apply(IndexedTensor.Indexes indexes) {
             if (freeGenerator != null) {
                 return freeGenerator.apply(indexes.toList());
@@ -173,4 +172,26 @@ public class Generate<NAMETYPE extends Name> extends PrimitiveTensorFunction<NAM
 
     }
 
+    /** A context which adds the bindings of the generate dimension names to the given context. */
+    private class GenerateToStringContext implements ToStringContext {
+
+        private final ToStringContext context;
+
+        public GenerateToStringContext(ToStringContext context) {
+            this.context = context;
+        }
+
+        @Override
+        public String getBinding(String identifier) {
+            if (type.dimension(identifier).isPresent())
+                return identifier; // dimension names are bound but not substituted in the generate context
+            else
+                return context.getBinding(identifier);
+        }
+
+        @Override
+        public ToStringContext wrapped() { return context; }
+
+    }
+
 }
diff --git a/vespajlib/src/main/java/com/yahoo/tensor/functions/ToStringContext.java b/vespajlib/src/main/java/com/yahoo/tensor/functions/ToStringContext.java
index cb7f376c365..c09631d36d7 100644
--- a/vespajlib/src/main/java/com/yahoo/tensor/functions/ToStringContext.java
+++ b/vespajlib/src/main/java/com/yahoo/tensor/functions/ToStringContext.java
@@ -3,13 +3,27 @@ package com.yahoo.tensor.functions;
 
 /**
  * A context which is passed down to all nested functions when returning a string representation.
- * The default implementation is empty as this library does not in itself have any need for a
- * context.
  *
  * @author bratseth
  */
 public interface ToStringContext {
 
-    static ToStringContext empty() { return new ToStringContext() {}; }
+    static ToStringContext empty() { return new EmptyStringContext(); }
+
+    /** Returns the name an identifier is bound to, or null if not bound in this context */
+    String getBinding(String name);
+
+    /** Returns another context this wraps, or null if none is wrapped */
+    ToStringContext wrapped();
+
+    class EmptyStringContext implements ToStringContext {
+
+        @Override
+        public String getBinding(String name) { return null; }
+
+        @Override
+        public ToStringContext wrapped() { return null; }
+
+    }
 
 }