Fix parsing of recursive "match" update syntax

3 files changed, 89 insertions, 97 deletions

diff --git a/document/src/main/java/com/yahoo/document/json/TokenBuffer.java b/document/src/main/java/com/yahoo/document/json/TokenBuffer.java
index c1ac239d5f0..dec84e46b77 100644
--- a/document/src/main/java/com/yahoo/document/json/TokenBuffer.java
+++ b/document/src/main/java/com/yahoo/document/json/TokenBuffer.java
@@ -150,34 +150,6 @@ public class TokenBuffer {
         return nesting;
     }
 
-    public Token prefetchScalar(String name) {
-        int localNesting = nesting();
-        int nestingBarrier = localNesting;
-        Token toReturn = null;
-        Iterator<Token> i;
-
-        if (name.equals(currentName()) && current().isScalarValue()) {
-            toReturn = tokens.get(position);
-        } else {
-            i = rest().iterator();
-            i.next(); // just ignore the first value, as we know it's not what
-                      // we're looking for, and it's nesting effect is already
-                      // included
-            while (i.hasNext()) {
-                Token t = i.next();
-                if (localNesting == nestingBarrier && name.equals(t.name) && t.token.isScalarValue()) {
-                    toReturn = t;
-                    break;
-                }
-                localNesting += nestingOffset(t.token);
-                if (localNesting < nestingBarrier) {
-                    break;
-                }
-            }
-        }
-        return toReturn;
-    }
-
     public void skipToRelativeNesting(int relativeNesting) {
         int initialNesting = nesting();
         do {
diff --git a/document/src/main/java/com/yahoo/document/json/readers/MapReader.java b/document/src/main/java/com/yahoo/document/json/readers/MapReader.java
index a373fa56fd9..c81d296db39 100644
--- a/document/src/main/java/com/yahoo/document/json/readers/MapReader.java
+++ b/document/src/main/java/com/yahoo/document/json/readers/MapReader.java
@@ -90,44 +90,37 @@ public class MapReader {
     @SuppressWarnings({ "rawtypes", "unchecked" })
     public static ValueUpdate createMapUpdate(TokenBuffer buffer,
                                               DataType currentLevel,
-                                              FieldValue keyParent,
-                                              FieldValue topLevelKey,
                                               boolean ignoreUndefinedFields) {
-        TokenBuffer.Token element = buffer.prefetchScalar(UPDATE_ELEMENT);
-        if (UPDATE_ELEMENT.equals(buffer.currentName())) {
+        if ( ! JsonToken.START_OBJECT.equals(buffer.current()))
+            throw new IllegalArgumentException("Expected object for match update, got " + buffer.current());
+        buffer.next();
+
+        FieldValue key = null;
+        ValueUpdate update;
+        boolean elementFirst = UPDATE_ELEMENT.equals(buffer.currentName());
+        if (elementFirst) {
+            key = keyTypeForMapUpdate(buffer.currentText(), currentLevel);
             buffer.next();
         }
 
-        FieldValue key = keyTypeForMapUpdate(element, currentLevel);
-        if (keyParent != null) {
-            ((CollectionFieldValue) keyParent).add(key);
-        }
-        // structure is: [(match + element)*, (element + action)]
-        // match will always have element, and either match or action
-        if (!UPDATE_MATCH.equals(buffer.currentName())) {
-            // we have reached an action...
-            if (topLevelKey == null) {
-                return ValueUpdate.createMap(key, readSingleUpdate(buffer, valueTypeForMapUpdate(currentLevel), buffer.currentName(), ignoreUndefinedFields));
-            } else {
-                return ValueUpdate.createMap(topLevelKey, readSingleUpdate(buffer, valueTypeForMapUpdate(currentLevel), buffer.currentName(), ignoreUndefinedFields));
-            }
-        } else {
-            // next level of matching
-            if (topLevelKey == null) {
-                return createMapUpdate(buffer, valueTypeForMapUpdate(currentLevel), key, key, ignoreUndefinedFields);
-            } else {
-                return createMapUpdate(buffer, valueTypeForMapUpdate(currentLevel), key, topLevelKey, ignoreUndefinedFields);
-            }
+        update = UPDATE_MATCH.equals(buffer.currentName()) ? createMapUpdate(buffer, valueTypeForMapUpdate(currentLevel), ignoreUndefinedFields)
+                                                           : readSingleUpdate(buffer, valueTypeForMapUpdate(currentLevel), buffer.currentName(), ignoreUndefinedFields);
+        buffer.next();
+
+        if ( ! elementFirst) {
+            key = keyTypeForMapUpdate(buffer.currentText(), currentLevel);
+            buffer.next();
         }
+
+        if ( ! JsonToken.END_OBJECT.equals(buffer.current()))
+            throw new IllegalArgumentException("Expected object end for match update, got " + buffer.current());
+
+        return ValueUpdate.createMap(key, update);
     }
 
     @SuppressWarnings("rawtypes")
     public static ValueUpdate createMapUpdate(TokenBuffer buffer, Field field, boolean ignoreUndefinedFields) {
-        buffer.next();
-        MapValueUpdate m = (MapValueUpdate) MapReader.createMapUpdate(buffer, field.getDataType(), null, null, ignoreUndefinedFields);
-        buffer.next();
-        // must generate the field value in parallel with the actual
-        return m;
+        return MapReader.createMapUpdate(buffer, field.getDataType(), ignoreUndefinedFields);
     }
 
     private static DataType valueTypeForMapUpdate(DataType parentType) {
@@ -142,14 +135,14 @@ public class MapReader {
         }
     }
 
-    private static FieldValue keyTypeForMapUpdate(TokenBuffer.Token element, DataType expectedType) {
+    private static FieldValue keyTypeForMapUpdate(String elementText, DataType expectedType) {
         FieldValue v;
         if (expectedType instanceof ArrayDataType) {
-            v = new IntegerFieldValue(Integer.valueOf(element.text));
+            v = new IntegerFieldValue(Integer.valueOf(elementText));
         } else if (expectedType instanceof WeightedSetDataType) {
-            v = ((WeightedSetDataType) expectedType).getNestedType().createFieldValue(element.text);
+            v = ((WeightedSetDataType) expectedType).getNestedType().createFieldValue(elementText);
         } else if (expectedType instanceof MapDataType) {
-            v = ((MapDataType) expectedType).getKeyType().createFieldValue(element.text);
+            v = ((MapDataType) expectedType).getKeyType().createFieldValue(elementText);
         } else {
             throw new IllegalArgumentException("Container type " + expectedType + " not supported for match update.");
         }
diff --git a/document/src/test/java/com/yahoo/document/json/JsonReaderTestCase.java b/document/src/test/java/com/yahoo/document/json/JsonReaderTestCase.java
index 839ce59fdd3..5e955ad82ac 100644
--- a/document/src/test/java/com/yahoo/document/json/JsonReaderTestCase.java
+++ b/document/src/test/java/com/yahoo/document/json/JsonReaderTestCase.java
@@ -770,47 +770,74 @@ public class JsonReaderTestCase {
     }
 
     @Test
-    public void testNestedArrayMatch() {
-        // Supposed to work?
-        assertThrows(ClassCastException.class,
-                     () -> parseUpdate("""
-                                       {
-                                         "update": "id:unittest:testArrayOfArrayOfInt::whee",
-                                         "fields": {
-                                           "arrayOfArrayOfInt": {
-                                             "match": {
-                                               "element": 1,
-                                               "match": {
-                                                 "element": 2,
-                                                 "assign": 3
-                                               }
-                                             }
-                                           }
-                                         }
-                                       }
-                                       """));
+    public void testNestedArrayMatch() throws IOException {
+        DocumentUpdate nested = parseUpdate("""
+                                            {
+                                              "update": "id:unittest:testArrayOfArrayOfInt::whee",
+                                              "fields": {
+                                                "arrayOfArrayOfInt": {
+                                                  "match": {
+                                                    "element": 1,
+                                                    "match": {
+                                                      "element": 2,
+                                                      "assign": 3
+                                                    }
+                                                  }
+                                                }
+                                              }
+                                            }
+                                            """);
+
+        DocumentUpdate equivalent = parseUpdate("""
+                                                {
+                                                  "update": "id:unittest:testArrayOfArrayOfInt::whee",
+                                                  "fields": {
+                                                    "arrayOfArrayOfInt": {
+                                                      "match": {
+                                                        "match": {
+                                                          "assign": 3,
+                                                          "element": 2
+                                                        },
+                                                        "element": 1
+                                                      }
+                                                    }
+                                                  }
+                                                }
+                                                """);
+
+        assertEquals(nested, equivalent);
+        assertEquals(1, nested.fieldUpdates().size());
+        FieldUpdate fu = nested.fieldUpdates().iterator().next();
+        assertEquals(1, fu.getValueUpdates().size());
+        MapValueUpdate mvu = (MapValueUpdate) fu.getValueUpdate(0);
+        assertEquals(new IntegerFieldValue(1), mvu.getValue());
+        MapValueUpdate nvu = (MapValueUpdate) mvu.getUpdate();
+        assertEquals(new IntegerFieldValue(2), nvu.getValue());
+        AssignValueUpdate avu = (AssignValueUpdate) nvu.getUpdate();
+        assertEquals(new IntegerFieldValue(3), avu.getValue());
     }
 
     @Test
     public void testMatchCannotUpdateNestedFields() {
         // Should this work? It doesn't.
-        assertThrows(ClassCastException.class,
-                     () -> parseUpdate("""
-                                       {
-                                         "update": "id:unittest:testMapStringToArrayOfInt::whee",
-                                         "fields": {
-                                           "actualMapStringToArrayOfInt": {
-                                             "match": {
-                                               "element": "bamse",
-                                               "match": {
-                                                 "element": 1,
-                                                 "assign": 4
-                                               }
-                                             }
-                                           }
-                                         }
-                                       }
-                                       """));
+        assertEquals("Field type Map<string,Array<int>> not supported.",
+                     assertThrows(UnsupportedOperationException.class,
+                                  () -> parseUpdate("""
+                                                    {
+                                                      "update": "id:unittest:testMapStringToArrayOfInt::whee",
+                                                      "fields": {
+                                                        "actualMapStringToArrayOfInt": {
+                                                          "match": {
+                                                            "element": "bamse",
+                                                            "match": {
+                                                              "element": 1,
+                                                              "assign": 4
+                                                            }
+                                                          }
+                                                        }
+                                                      }
+                                                    }
+                                                    """)).getMessage());
     }
 
     @Test