Add support for shortcut properties of non-primitive type

java-client/src/main/java/co/elastic/clients/elasticsearch/_types/query_dsl/FunctionScoreQuery.java

@@ -359,6 +359,8 @@ protected static void setupFunctionScoreQueryDeserializer(ObjectDeserializer<Fun
 		op.add(Builder::query, Query._DESERIALIZER, "query");
 		op.add(Builder::scoreMode, FunctionScoreMode._DESERIALIZER, "score_mode");
 
+		op.shortcutProperty("functions", true);
+
 	}
 
 }

java-client/src/main/java/co/elastic/clients/elasticsearch/_types/query_dsl/FuzzyQuery.java

@@ -389,7 +389,7 @@ protected static void setupFuzzyQueryDeserializer(ObjectDeserializer<FuzzyQuery.
 		op.add(Builder::value, FieldValue._DESERIALIZER, "value");
 
 		op.setKey(Builder::field, JsonpDeserializer.stringDeserializer());
-		op.shortcutProperty("value");
+		op.shortcutProperty("value", true);
 
 	}
 

java-client/src/main/java/co/elastic/clients/elasticsearch/_types/query_dsl/MatchQuery.java

@@ -671,7 +671,7 @@ protected static void setupMatchQueryDeserializer(ObjectDeserializer<MatchQuery.
 		op.add(Builder::zeroTermsQuery, ZeroTermsQuery._DESERIALIZER, "zero_terms_query");
 
 		op.setKey(Builder::field, JsonpDeserializer.stringDeserializer());
-		op.shortcutProperty("query");
+		op.shortcutProperty("query", true);
 
 	}
 

java-client/src/main/java/co/elastic/clients/elasticsearch/_types/query_dsl/TermQuery.java

@@ -294,7 +294,7 @@ protected static void setupTermQueryDeserializer(ObjectDeserializer<TermQuery.Bu
 		op.add(Builder::caseInsensitive, JsonpDeserializer.booleanDeserializer(), "case_insensitive");
 
 		op.setKey(Builder::field, JsonpDeserializer.stringDeserializer());
-		op.shortcutProperty("value");
+		op.shortcutProperty("value", true);
 
 	}
 

java-client/src/main/java/co/elastic/clients/elasticsearch/core/rank_eval/RankEvalQuery.java

@@ -187,7 +187,7 @@ protected static void setupRankEvalQueryDeserializer(ObjectDeserializer<RankEval
 		op.add(Builder::query, Query._DESERIALIZER, "query");
 		op.add(Builder::size, JsonpDeserializer.integerDeserializer(), "size");
 
-		op.shortcutProperty("query");
+		op.shortcutProperty("query", true);
 
 	}
 

java-client/src/main/java/co/elastic/clients/elasticsearch/core/search/CompletionContext.java

@@ -362,7 +362,7 @@ protected static void setupCompletionContextDeserializer(ObjectDeserializer<Comp
 		op.add(Builder::precision, GeoHashPrecision._DESERIALIZER, "precision");
 		op.add(Builder::prefix, JsonpDeserializer.booleanDeserializer(), "prefix");
 
-		op.shortcutProperty("context");
+		op.shortcutProperty("context", true);
 
 	}
 

java-client/src/main/java/co/elastic/clients/json/JsonpUtils.java

@@ -253,7 +253,7 @@ public static Map.Entry<String, JsonParser> lookAheadFieldValue(
                 throw new JsonpMappingException("Property '" + name + "' not found", location);
             }
 
-            JsonParser newParser = objectParser(object, mapper);
+            JsonParser newParser = jsonValueParser(object, mapper);
 
             // Pin location to the start of the look ahead, as the new parser will return locations in its own buffer
             newParser = new DelegatingJsonParser(newParser) {
@@ -272,16 +272,60 @@ public String toString() {
         }
     }
 
+    /**
+     * In union types, find the variant to be used by looking up property names in the JSON stream until we find one that
+     * uniquely identifies the variant.
+     *
+     * @param <Variant> the type of variant descriptors used by the caller.
+     * @param variants a map of variant descriptors, keyed by the property name that uniquely identifies the variant.
+     * @return a pair containing the variant descriptor (or {@code null} if not found), and a parser to be used to read the JSON object.
+     */
+
+    public static <Variant> Map.Entry<Variant, JsonParser> findVariant(
+        Map<String, Variant> variants, JsonParser parser, JsonpMapper mapper
+    ) {
+        if (parser instanceof LookAheadJsonParser) {
+            return ((LookAheadJsonParser) parser).findVariant(variants);
+        } else {
+            // If it's an object, find matching field names
+            Variant variant = null;
+            JsonValue value = parser.getValue();
+
+            if (value instanceof JsonObject) {
+                for (String field: value.asJsonObject().keySet()) {
+                    variant = variants.get(field);
+                    if (variant != null) {
+                        break;
+                    }
+                }
+            }
+
+            // Traverse the object we have inspected
+            parser = JsonpUtils.jsonValueParser(value, mapper);
+            return new AbstractMap.SimpleImmutableEntry<>(variant, parser);
+        }
+    }
+
     /**
      * Create a parser that traverses a JSON object
+     *
+     * @deprecated use {@link #jsonValueParser(JsonValue, JsonpMapper)}
      */
+    @Deprecated
     public static JsonParser objectParser(JsonObject object, JsonpMapper mapper) {
+        return jsonValueParser(object, mapper);
+    }
+
+    /**
+     * Create a parser that traverses a JSON value
+     */
+    public static JsonParser jsonValueParser(JsonValue value, JsonpMapper mapper) {
         // FIXME: we should have used createParser(object), but this doesn't work as it creates a
         // org.glassfish.json.JsonStructureParser that doesn't implement the JsonP 1.0.1 features, in particular
         // parser.getObject(). So deserializing recursive internally-tagged union would fail with UnsupportedOperationException
         // While glassfish has this issue or until we write our own, we roundtrip through a string.
 
-        String strObject = object.toString();
+        String strObject = value.toString();
         return mapper.jsonProvider().createParser(new StringReader(strObject));
     }
 

java-client/src/main/java/co/elastic/clients/json/ObjectDeserializer.java

@@ -106,7 +106,6 @@ public EnumSet<Event> acceptedEvents() {
 
     //---------------------------------------------------------------------------------------------
     private static final EnumSet<Event> EventSetObject = EnumSet.of(Event.START_OBJECT, Event.KEY_NAME);
-    private static final EnumSet<Event> EventSetObjectAndString = EnumSet.of(Event.START_OBJECT, Event.VALUE_STRING, Event.KEY_NAME);
 
     private EnumSet<Event> acceptedEvents = EventSetObject; // May be changed in `shortcutProperty()`
     private final Supplier<ObjectType> constructor;
@@ -115,6 +114,7 @@ public EnumSet<Event> acceptedEvents() {
     private String typeProperty;
     private String defaultType;
     private FieldDeserializer<ObjectType> shortcutProperty;
+    private boolean shortcutIsObject;
     private QuadConsumer<ObjectType, String, JsonParser, JsonpMapper> unknownFieldHandler;
 
     public ObjectDeserializer(Supplier<ObjectType> constructor) {
@@ -133,6 +133,10 @@ public Set<String> fieldNames() {
         return this.shortcutProperty == null ? null : this.shortcutProperty.name;
     }
 
+    public boolean shortcutIsObject() {
+        return this.shortcutIsObject;
+    }
+
     @Override
     public EnumSet<Event> nativeEvents() {
         // May also return string if we have a shortcut property. This is needed to identify ambiguous unions.
@@ -145,33 +149,51 @@ public EnumSet<Event> acceptedEvents() {
     }
 
     public ObjectType deserialize(JsonParser parser, JsonpMapper mapper, Event event) {
-        return deserialize(constructor.get(), parser, mapper, event);
-    }
-
-    public ObjectType deserialize(ObjectType value, JsonParser parser, JsonpMapper mapper, Event event) {
         if (event == Event.VALUE_NULL) {
             return null;
         }
 
-        String keyName = null;
-        String fieldName = null;
+        ObjectType value = constructor.get();
+        deserialize(value, parser, mapper, event);
+        return value;
+    }
 
-        try {
+    public void deserialize(ObjectType value, JsonParser parser, JsonpMapper mapper, Event event) {
+        // Note: method is public as it's called by `withJson` to augment an already created object
 
-            if (singleKey != null) {
-                // There's a wrapping property whose name is the key value
-                if (event == Event.START_OBJECT) {
-                    event = JsonpUtils.expectNextEvent(parser, Event.KEY_NAME);
-                }
+        if (singleKey == null) {
+            // Nominal case
+            deserializeInner(value, parser, mapper, event);
+
+        } else {
+            // Single key dictionary: there's a wrapping property whose name is the key value
+            if (event == Event.START_OBJECT) {
+                event = JsonpUtils.expectNextEvent(parser, Event.KEY_NAME);
+            }
+
+            String keyName = parser.getString();
+            try {
                 singleKey.deserialize(parser, mapper, null, value, event);
                 event = parser.next();
+                deserializeInner(value, parser, mapper, event);
+            } catch (Exception e) {
+                throw JsonpMappingException.from(e, value, keyName, parser);
             }
 
-            if (shortcutProperty != null && event != Event.START_OBJECT && event != Event.KEY_NAME) {
-                // This is the shortcut property (should be a value event, this will be checked by its deserializer)
-                shortcutProperty.deserialize(parser, mapper, shortcutProperty.name, value, event);
+            JsonpUtils.expectNextEvent(parser, Event.END_OBJECT);
+        }
+    }
+
+    private void deserializeInner(ObjectType value, JsonParser parser, JsonpMapper mapper, Event event) {
+        String fieldName = null;
 
-            } else if (typeProperty == null) {
+        try {
+            if ((parser = deserializeWithShortcut(value, parser, mapper, event)) == null) {
+                // We found the shortcut form
+                return;
+            }
+
+            if (typeProperty == null) {
                 if (event != Event.START_OBJECT && event != Event.KEY_NAME) {
                     // Report we're waiting for a start_object, since this is the most common beginning for object parser
                     JsonpUtils.expectEvent(parser, Event.START_OBJECT, event);
@@ -209,16 +231,52 @@ public ObjectType deserialize(ObjectType value, JsonParser parser, JsonpMapper m
                     fieldDeserializer.deserialize(innerParser, mapper, variant, value);
                 }
             }
+        } catch (Exception e) {
+            // Add field name if present
+            throw JsonpMappingException.from(e, value, fieldName, parser);
+        }
+    }
 
-            if (singleKey != null) {
-                JsonpUtils.expectNextEvent(parser, Event.END_OBJECT);
+    /**
+     * Try to deserialize the value with its shortcut property, if any.
+     *
+     * @return {@code null} if the shortcut form was found, and otherwise a parser that should be used to parse the
+     *         non-shortcut form. It may be different from the orginal parser if look-ahead was needed.
+     */
+    private JsonParser deserializeWithShortcut(ObjectType value, JsonParser parser, JsonpMapper mapper, Event event) {
+        if (shortcutProperty != null) {
+            if (!shortcutIsObject) {
+                if (event != Event.START_OBJECT && event != Event.KEY_NAME) {
+                    // This is the shortcut property (should be a value or array event, this will be checked by its deserializer)
+                    shortcutProperty.deserialize(parser, mapper, shortcutProperty.name, value, event);
+                    return null;
+                }
+            } else {
+                // Fast path: we don't need to look ahead if the current event isn't an object
+                if (event != Event.START_OBJECT) {
+                    shortcutProperty.deserialize(parser, mapper, shortcutProperty.name, value, event);
+                    return null;
+                }
+
+                // Look ahead: does the shortcut property exist? If yes, the shortcut is used
+                Map.Entry<Object, JsonParser> shortcut = JsonpUtils.findVariant(
+                    Collections.singletonMap(shortcutProperty.name, Boolean.TRUE /* arbitrary non-null value */),
+                    parser, mapper
+                );
+
+                // Parse the buffered events
+                parser = shortcut.getValue();
+                event = parser.next();
+
+                // If shortcut property was not found, this is a shortcut. Otherwise, keep deserializing as usual
+                if (shortcut.getKey() == null) {
+                    shortcutProperty.deserialize(parser, mapper, shortcutProperty.name, value, event);
+                    return null;
+                }
             }
-        } catch (Exception e) {
-            // Add key name (for single key dicts) and field name if present
-            throw JsonpMappingException.from(e, value, fieldName, parser).prepend(value, keyName);
         }
 
-        return value;
+        return parser;
     }
 
     protected void parseUnknownField(JsonParser parser, JsonpMapper mapper, String fieldName, ObjectType object) {
@@ -249,14 +307,18 @@ public void ignore(String name) {
     }
 
     public void shortcutProperty(String name) {
+        shortcutProperty(name, false);
+    }
+
+    public void shortcutProperty(String name, boolean isObject) {
         this.shortcutProperty = this.fieldDeserializers.get(name);
         if (this.shortcutProperty == null) {
             throw new NoSuchElementException("No deserializer was setup for '" + name + "'");
         }
 
-        //acceptedEvents = EnumSet.copyOf(acceptedEvents);
-        //acceptedEvents.addAll(shortcutProperty.acceptedEvents());
-        acceptedEvents = EventSetObjectAndString;
+        acceptedEvents = EnumSet.copyOf(acceptedEvents);
+        acceptedEvents.addAll(shortcutProperty.acceptedEvents());
+        this.shortcutIsObject = isObject;
     }
 
     //----- Object types

java-client/src/main/java/co/elastic/clients/json/UnionDeserializer.java

@@ -20,7 +20,6 @@
 package co.elastic.clients.json;
 
 import co.elastic.clients.util.ObjectBuilder;
-import jakarta.json.JsonObject;
 import jakarta.json.stream.JsonLocation;
 import jakarta.json.stream.JsonParser;
 import jakarta.json.stream.JsonParser.Event;
@@ -265,28 +264,12 @@ public Union deserialize(JsonParser parser, JsonpMapper mapper, Event event) {
         JsonLocation location = parser.getLocation();
 
         if (member == null && event == Event.START_OBJECT && !objectMembers.isEmpty()) {
-            if (parser instanceof LookAheadJsonParser) {
-                Map.Entry<EventHandler<Union, Kind, Member>, JsonParser> memberAndParser =
-                    ((LookAheadJsonParser) parser).findVariant(objectMembers);
+            Map.Entry<EventHandler<Union, Kind, Member>, JsonParser> memberAndParser =
+                JsonpUtils.findVariant(objectMembers, parser, mapper);
 
-                member = memberAndParser.getKey();
-                // Parse the buffered parser
-                parser = memberAndParser.getValue();
-
-            } else {
-                // Parse as an object to find matching field names
-                JsonObject object = parser.getObject();
-
-                for (String field: object.keySet()) {
-                    member = objectMembers.get(field);
-                    if (member != null) {
-                        break;
-                    }
-                }
-
-                // Traverse the object we have inspected
-                parser = JsonpUtils.objectParser(object, mapper);
-            }
+            member = memberAndParser.getKey();
+            // Parse the buffered parser
+            parser = memberAndParser.getValue();
 
             if (member == null) {
                 member = fallbackObjectMember;

java-client/src/main/java/co/elastic/clients/json/jackson/JacksonJsonpParser.java

@@ -372,29 +372,34 @@ public <Variant> Map.Entry<Variant, JsonParser> findVariant(Map<String, Variant>
         TokenBuffer tb = new TokenBuffer(parser, null);
 
         try {
-            // The resulting parser must contain the full object, including START_EVENT
-            tb.copyCurrentEvent(parser);
-            while (parser.nextToken() != JsonToken.END_OBJECT) {
-
-                expectEvent(JsonToken.FIELD_NAME);
-                String fieldName = parser.getCurrentName();
-
-                Variant variant = variants.get(fieldName);
-                if (variant != null) {
-                    tb.copyCurrentEvent(parser);
-                    return new AbstractMap.SimpleImmutableEntry<>(
-                        variant,
-                        new JacksonJsonpParser(
-                            JsonParserSequence.createFlattened(false, tb.asParser(), parser),
-                            mapper
-                        )
-                    );
-                } else {
-                    tb.copyCurrentStructure(parser);
+            if (parser.currentToken() != JsonToken.START_OBJECT) {
+                // Primitive value or array
+                tb.copyCurrentStructure(parser);
+            } else {
+                // The resulting parser must contain the full object, including START_EVENT
+                tb.copyCurrentEvent(parser);
+                while (parser.nextToken() != JsonToken.END_OBJECT) {
+
+                    expectEvent(JsonToken.FIELD_NAME);
+                    String fieldName = parser.getCurrentName();
+
+                    Variant variant = variants.get(fieldName);
+                    if (variant != null) {
+                        tb.copyCurrentEvent(parser);
+                        return new AbstractMap.SimpleImmutableEntry<>(
+                            variant,
+                            new JacksonJsonpParser(
+                                JsonParserSequence.createFlattened(false, tb.asParser(), parser),
+                                mapper
+                            )
+                        );
+                    } else {
+                        tb.copyCurrentStructure(parser);
+                    }
                 }
+                // Copy ending END_OBJECT
+                tb.copyCurrentEvent(parser);
             }
-            // Copy ending END_OBJECT
-            tb.copyCurrentEvent(parser);
         } catch (IOException e) {
             throw JacksonUtils.convertException(e);
         }