1 files changed, 83 insertions, 19 deletions

diff --git a/searchlib/src/main/javacc/RankingExpressionParser.jj b/searchlib/src/main/javacc/RankingExpressionParser.jj
index de3ad6b5d8c..e413e398183 100755
--- a/searchlib/src/main/javacc/RankingExpressionParser.jj
+++ b/searchlib/src/main/javacc/RankingExpressionParser.jj
@@ -6,7 +6,6 @@
  */
 options {
     CACHE_TOKENS       = true;
-    STATIC             = false;
     DEBUG_PARSER       = false;
     USER_TOKEN_MANAGER = false;
     ERROR_REPORTING    = true;
@@ -476,13 +475,14 @@ TensorFunctionNode tensorConcat() :
 TensorFunctionNode tensorGenerate() :
 {
     TensorType type;
+    List dimensionOrder = new ArrayList();
     TensorFunctionNode expression;
 }
 {
-    <TENSOR> type = tensorType()
+    <TENSOR> type = tensorType(dimensionOrder)
     (
         expression = tensorGenerateBody(type) |
-        expression = tensorValueBody(type)
+        expression = tensorValueBody(type, dimensionOrder)
     )
     { return expression; }
 }
@@ -501,7 +501,7 @@ TensorFunctionNode tensorRange() :
     TensorType type;
 }
 {
-    <RANGE> type = tensorType()
+    <RANGE> type = tensorType(null)
     { return new TensorFunctionNode(new Range(type)); }
 }
 
@@ -510,7 +510,7 @@ TensorFunctionNode tensorDiag() :
     TensorType type;
 }
 {
-    <DIAG> type = tensorType()
+    <DIAG> type = tensorType(null)
     { return new TensorFunctionNode(new Diag(type)); }
 }
 
@@ -519,7 +519,7 @@ TensorFunctionNode tensorRandom() :
     TensorType type;
 }
 {
-    <RANDOM> type = tensorType()
+    <RANDOM> type = tensorType(null)
     { return new TensorFunctionNode(new Random(type)); }
 }
 
@@ -619,7 +619,7 @@ Reduce.Aggregator tensorReduceAggregator() :
     { return Reduce.Aggregator.valueOf(token.image); }
 }
 
-TensorType tensorType() :
+TensorType tensorType(List dimensionOrder) :
 {
     TensorType.Builder builder;
     TensorType.Value valueType;
@@ -628,8 +628,8 @@ TensorType tensorType() :
     valueType = optionalTensorValueTypeParameter()
     { builder = new TensorType.Builder(valueType); }
     <LBRACE>
-    ( tensorTypeDimension(builder) ) ?
-    ( <COMMA> tensorTypeDimension(builder) ) *
+    ( tensorTypeDimension(builder, dimensionOrder) ) ?
+    ( <COMMA> tensorTypeDimension(builder, dimensionOrder) ) *
     <RBRACE>
     { return builder.build(); }
 }
@@ -643,13 +643,17 @@ TensorType.Value optionalTensorValueTypeParameter() :
     { return TensorType.Value.fromId(valueType); }
 }
 
-void tensorTypeDimension(TensorType.Builder builder) :
+void tensorTypeDimension(TensorType.Builder builder, List dimensionOrder) :
 {
     String name;
     int size;
 }
 {
     name = identifier()
+    {   // Keep track of the order in which dimensions are written, if necessary
+        if (dimensionOrder != null)
+            dimensionOrder.add(name);
+    }
     (
       ( <LCURLY> <RCURLY> { builder.mapped(name); } ) |
       LOOKAHEAD(2) ( <LSQUARE> <RSQUARE> { builder.indexed(name); } ) |
@@ -832,15 +836,16 @@ Value primitiveValue() :
     { return Value.parse(sign + token.image); }
 }
 
-TensorFunctionNode tensorValueBody(TensorType type) :
+TensorFunctionNode tensorValueBody(TensorType type, List dimensionOrder) :
 {
     DynamicTensor dynamicTensor;
 }
 {
     <COLON>
     (
+        LOOKAHEAD(2) dynamicTensor = mixedTensorValueBody(type, dimensionOrder) |
         dynamicTensor = mappedTensorValueBody(type) |
-        dynamicTensor = indexedTensorValueBody(type)
+        dynamicTensor = indexedTensorValueBody(type, dimensionOrder)
     )
     { return new TensorFunctionNode(dynamicTensor); }
 }
@@ -851,23 +856,82 @@ DynamicTensor mappedTensorValueBody(TensorType type) :
 }
 {
     <LCURLY>
-        (         tensorCell(type, cells))*
+        [         tensorCell(type, cells)]
         ( <COMMA> tensorCell(type, cells))*
     <RCURLY>
     { return DynamicTensor.from(type, TensorFunctionNode.wrapScalars(cells)); }
 }
 
-DynamicTensor indexedTensorValueBody(TensorType type) :
+DynamicTensor mixedTensorValueBody(TensorType type, List dimensionOrder) :
+{
+    java.util.Map cells = new LinkedHashMap();
+}
+{
+    <LCURLY>
+        keyValueOrMixedBlock(type, dimensionOrder, cells)
+        ( <COMMA> keyValueOrMixedBlock(type, dimensionOrder, cells))*
+    <RCURLY>
+    { return DynamicTensor.from(type, cells); }
+}
+
+DynamicTensor indexedTensorValueBody(TensorType type, List dimensionOrder) :
+{
+    List cells;
+}
+{
+    cells = indexedTensorCells()
+    { return DynamicTensor.from(type, TensorFunctionNode.wrapScalars(type, dimensionOrder, cells)); }
+}
+
+void keyValueOrMixedBlock(TensorType type, List dimensionOrder, java.util.Map cellMap) : {}
+{
+    LOOKAHEAD(3) mixedBlock(type, dimensionOrder, cellMap) | keyValue(type, cellMap)
+}
+
+void keyValue(TensorType type, java.util.Map cellMap) :
+{
+    String label;
+    ExpressionNode value;
+}
+{
+    label = tag() <COLON> value = expression()
+    { cellMap.put(TensorAddress.ofLabels(label), TensorFunctionNode.wrapScalar(value)); }
+}
+
+void mixedBlock(TensorType type, List dimensionOrder, java.util.Map cellMap) :
+{
+    String label;
+    List cells;
+}
+{
+    label = tag() <COLON> cells = indexedTensorCells()
+    { TensorFunctionNode.wrapScalarBlock(type, dimensionOrder, label, cells, cellMap); }
+}
+
+List indexedTensorCells() :
 {
     List cells = new ArrayList();
+}
+{
+    <LSQUARE> indexedTensorCellSubspaceList(cells) <RSQUARE>
+    { return cells; }
+}
+
+void indexedTensorCellSubspaceList(List cells) :
+{
+}
+{
+    indexedTensorCellSubspace(cells) ( LOOKAHEAD(2) <COMMA> indexedTensorCellSubspace(cells) )*
+}
+
+void indexedTensorCellSubspace(List cells) :
+{
     ExpressionNode value;
 }
 {
-    <LSQUARE> // TODO: Parse inner square brackets properly
-        ( (<LSQUARE>)* value = expression() (<RSQUARE>)* { cells.add(value); }  )*
-        ( <COMMA> (<LSQUARE>)* value = expression() (<RSQUARE>)* { cells.add(value); } )*
-//    <RSQUARE>
-    { return DynamicTensor.from(type, TensorFunctionNode.wrapScalars(cells)); }
+    ( <LSQUARE> indexedTensorCellSubspaceList(cells) <RSQUARE> )
+    |
+    ( value = expression() { cells.add(value); } )
 }
 
 void tensorCell(TensorType type, java.util.Map cells) :