Prefix index variables by %

3 files changed, 33 insertions, 13 deletions

diff --git a/searchlib/src/main/javacc/RankingExpressionParser.jj b/searchlib/src/main/javacc/RankingExpressionParser.jj
index c7870182939..1153053ed09 100755
--- a/searchlib/src/main/javacc/RankingExpressionParser.jj
+++ b/searchlib/src/main/javacc/RankingExpressionParser.jj
@@ -68,7 +68,7 @@ TOKEN :
     <MOD: "%"> |
     <POWOP: "^"> |
 
-    <DOLLAR: "$"> |
+    <DOLLAR:"$"> |
     <COMMA: ","> |
     <COLON: ":"> |
 
@@ -244,6 +244,7 @@ ExpressionNode value() :
                    LOOKAHEAD(4) value = function()                            |
                    value = feature()                                          |
                    value = legacyQueryFeature()                               |
+                   value = indexVariable()                                    |
                    ( <LBRACE> value = expression() <RBRACE> { value = new EmbracedNode(value); } ) )
 
       )
@@ -290,6 +291,17 @@ ReferenceNode legacyQueryFeature() :
     }
 }
 
+ReferenceNode indexVariable() :
+{
+    String name;
+}
+{
+     ( <MOD> name = identifier() )
+     {
+         return new ReferenceNode("%" + name);
+     }
+}
+
 String outs() :
 {
     StringBuilder ret = new StringBuilder();
diff --git a/searchlib/src/test/java/com/yahoo/searchlib/rankingexpression/evaluation/EvaluationTestCase.java b/searchlib/src/test/java/com/yahoo/searchlib/rankingexpression/evaluation/EvaluationTestCase.java
index 99047aeb79d..f89d5358922 100644
--- a/searchlib/src/test/java/com/yahoo/searchlib/rankingexpression/evaluation/EvaluationTestCase.java
+++ b/searchlib/src/test/java/com/yahoo/searchlib/rankingexpression/evaluation/EvaluationTestCase.java
@@ -361,16 +361,21 @@ public class EvaluationTestCase {
         tester.assertEvaluates("tensor(j[3]):[3, 3, 3]",
                                "tensor(j[3])(tensor0[2])",
                                "tensor(values[5]):[1, 2, 3, 4, 5]");
+        // tensor result dimensions are given from argument dimensions, not the resulting values
+        tester.assertEvaluates("tensor(x{}):{}", "tensor0 * tensor1", "{ {x:0}:1 }", "tensor(x{}):{ {x:1}:1 }");
+        tester.assertEvaluates("tensor(x{},y{}):{}", "tensor0 * tensor1", "{ {x:0}:1 }", "tensor(x{},y{}):{ {x:1,y:0}:1, {x:2,y:1}:1 }");
+    }
+
+    @Test
+    public void testTensorReshape() {
+        EvaluationTester tester = new EvaluationTester();
         tester.assertEvaluates("tensor(j[3]):[5, 4, 3]",
-                               "tensor(j[3])(tensor0[4-j])",
+                               "tensor(j[3])(tensor0[4-%j])",
                                "tensor(values[5]):[1, 2, 3, 4, 5]");
         tester.assertEvaluates("tensor(j[2]):[6, 5]",
-                               "tensor(j[2])(tensor0{key:bar,i:2-j})",
+                               "tensor(j[2])(tensor0{key:bar,i:2-%j})",
                                "tensor(key{},i[5]):{{key:foo,i:0}:1,{key:foo,i:1}:2,{key:foo,i:2}:2,{key:bar,i:0}:4,{key:bar,i:1}:5,{key:bar,i:2}:6}");
 
-        // tensor result dimensions are given from argument dimensions, not the resulting values
-        tester.assertEvaluates("tensor(x{}):{}", "tensor0 * tensor1", "{ {x:0}:1 }", "tensor(x{}):{ {x:1}:1 }");
-        tester.assertEvaluates("tensor(x{},y{}):{}", "tensor0 * tensor1", "{ {x:0}:1 }", "tensor(x{},y{}):{ {x:1,y:0}:1, {x:2,y:1}:1 }");
     }
 
     @Test
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 aaed607aaa1..d3a2dc9169c 100644
--- a/vespajlib/src/main/java/com/yahoo/tensor/functions/Generate.java
+++ b/vespajlib/src/main/java/com/yahoo/tensor/functions/Generate.java
@@ -145,17 +145,15 @@ public class Generate<NAMETYPE extends TypeContext.Name>  extends PrimitiveTenso
 
         @Override
         public Tensor getTensor(String name) {
-            Optional<Integer> index = type.indexOfDimension(name);
-            if (index.isPresent()) // this is the name of a dimension
-                return Tensor.from(indexes.indexesForReading()[index.get()]);
+            if (name.startsWith("%")) // this is the name of a dimension
+                return Tensor.from(indexes.indexesForReading()[dimensionIndexFrom(name)]);
             else
                 return context.getTensor(name);
         }
 
         @Override
         public TensorType getType(NAMETYPE name) {
-            Optional<Integer> index = type.indexOfDimension(name.name());
-            if (index.isPresent()) // this is the name of a dimension
+            if (name.name().startsWith("%")) // this is the name of a dimension
                 return TensorType.empty;
             else
                 return context.getType(name);
@@ -163,13 +161,18 @@ public class Generate<NAMETYPE extends TypeContext.Name>  extends PrimitiveTenso
 
         @Override
         public TensorType getType(String name) {
-            Optional<Integer> index = type.indexOfDimension(name);
-            if (index.isPresent()) // this is the name of a dimension
+            if (name.startsWith("%")) // this is the name of a dimension
                 return TensorType.empty;
             else
                 return context.getType(name);
         }
 
+        private int dimensionIndexFrom(String nameWithPercentage) {
+            String name = nameWithPercentage.substring(1);
+            return type.indexOfDimension(name)
+                       .orElseThrow(() -> new IllegalArgumentException("No dimension '" + name + "' in " + type));
+        }
+
     }
 
 }