Add `doc` fields to script factory (#60992)

This causes painless to generate the `docFields` method on every script
factory. It returns a `List<String>` containing all of the constant
keys passed to the `doc` variable's `get` method.

modules/lang-painless/src/main/java/org/elasticsearch/painless/PainlessScriptEngine.java

@@ -299,16 +299,15 @@ public final class PainlessScriptEngine implements ScriptEngine {
         constructor.endMethod();
 
         Method reflect = null;
+        Method docFieldsReflect = null;
 
         for (Method method : context.factoryClazz.getMethods()) {
             if ("newInstance".equals(method.getName())) {
                 reflect = method;
-
-                break;
             } else if ("newFactory".equals(method.getName())) {
                 reflect = method;
-
-                break;
+            } else if ("docFields".equals(method.getName())) {
+                docFieldsReflect = method;
             }
         }
 
@@ -341,6 +340,32 @@ public final class PainlessScriptEngine implements ScriptEngine {
         deterAdapter.returnValue();
         deterAdapter.endMethod();
 
+        if (docFieldsReflect != null) {
+            if (false == docFieldsReflect.getReturnType().equals(List.class)) {
+                throw new IllegalArgumentException("doc_fields must return a List");
+            }
+            if (docFieldsReflect.getParameterCount() != 0) {
+                throw new IllegalArgumentException("doc_fields may not take parameters");
+            }
+            org.objectweb.asm.commons.Method docFields = new org.objectweb.asm.commons.Method(docFieldsReflect.getName(),
+                MethodType.methodType(List.class).toMethodDescriptorString());
+            GeneratorAdapter docAdapter = new GeneratorAdapter(Opcodes.ASM5, docFields,
+                writer.visitMethod(Opcodes.ACC_PUBLIC, docFieldsReflect.getName(), docFields.getDescriptor(), null, null));
+            docAdapter.visitCode();
+            docAdapter.newInstance(WriterConstants.ARRAY_LIST_TYPE);
+            docAdapter.dup();
+            docAdapter.push(scriptScope.docFields().size());
+            docAdapter.invokeConstructor(WriterConstants.ARRAY_LIST_TYPE, WriterConstants.ARRAY_LIST_CTOR_WITH_SIZE);
+            for (int i = 0; i < scriptScope.docFields().size(); i++) {
+                docAdapter.dup();
+                docAdapter.push(scriptScope.docFields().get(i));
+                docAdapter.invokeInterface(WriterConstants.LIST_TYPE, WriterConstants.LIST_ADD);
+                docAdapter.pop(); // Don't want the result of calling add
+            }
+            docAdapter.returnValue();
+            docAdapter.endMethod();
+        }
+
         writer.visitEnd();
         Class<?> factory = loader.defineFactory(className.replace('/', '.'), writer.toByteArray());
 

modules/lang-painless/src/main/java/org/elasticsearch/painless/WriterConstants.java

@@ -30,8 +30,10 @@ import java.lang.invoke.MethodHandle;
 import java.lang.invoke.MethodHandles;
 import java.lang.invoke.MethodType;
 import java.time.ZonedDateTime;
+import java.util.ArrayList;
 import java.util.Collection;
 import java.util.Iterator;
+import java.util.List;
 import java.util.Objects;
 import java.util.regex.Matcher;
 import java.util.regex.Pattern;
@@ -182,6 +184,12 @@ public final class WriterConstants {
     public static final Type COLLECTION_TYPE = Type.getType(Collection.class);
     public static final Method COLLECTION_SIZE = getAsmMethod(int.class, "size");
 
+    public static final Type LIST_TYPE = Type.getType(List.class);
+    public static final Method LIST_ADD = getAsmMethod(boolean.class, "add", Object.class);
+
+    public static final Type ARRAY_LIST_TYPE = Type.getType(ArrayList.class);
+    public static final Method ARRAY_LIST_CTOR_WITH_SIZE = getAsmMethod(void.class, CTOR_METHOD_NAME, int.class);
+
     private static Method getAsmMethod(final Class<?> rtype, final String name, final Class<?>... ptypes) {
         return new Method(name, MethodType.methodType(rtype, ptypes).toMethodDescriptorString());
     }

modules/lang-painless/src/test/java/org/elasticsearch/painless/FactoryTests.java

@@ -28,6 +28,8 @@ import java.util.Collections;
 import java.util.List;
 import java.util.Map;
 
+import static org.hamcrest.Matchers.equalTo;
+
 public class FactoryTests extends ScriptTestCase {
 
     @Override
@@ -40,6 +42,7 @@ public class FactoryTests extends ScriptTestCase {
         contexts.put(TemplateScript.CONTEXT, Whitelist.BASE_WHITELISTS);
         contexts.put(VoidReturnTestScript.CONTEXT, Whitelist.BASE_WHITELISTS);
         contexts.put(FactoryTestConverterScript.CONTEXT, Whitelist.BASE_WHITELISTS);
+        contexts.put(DocFieldsTestScript.CONTEXT, Whitelist.BASE_WHITELISTS);
 
         return contexts;
     }
@@ -183,6 +186,8 @@ public class FactoryTests extends ScriptTestCase {
         FactoryTestScript script = factory.newInstance(Collections.singletonMap("test", 2));
         assertEquals(4, script.execute(2));
         assertEquals(5, script.execute(3));
+        // The factory interface doesn't define `docFields` so we don't generate it.
+        expectThrows(NoSuchMethodException.class, () -> factory.getClass().getMethod("docFields"));
         script = factory.newInstance(Collections.singletonMap("test", 3));
         assertEquals(5, script.execute(2));
         assertEquals(2, script.execute(-1));
@@ -369,4 +374,31 @@ public class FactoryTests extends ScriptTestCase {
                 FactoryTestConverterScript.CONTEXT, Collections.emptyMap()));
         assertEquals(cce.getMessage(), "Cannot cast from [boolean] to [long[]].");
     }
+
+    public abstract static class DocFieldsTestScript {
+        public static final ScriptContext<DocFieldsTestScript.Factory> CONTEXT = new ScriptContext<>(
+            "test",
+            DocFieldsTestScript.Factory.class
+        );
+
+        public interface Factory {
+            DocFieldsTestScript newInstance();
+
+            List<String> docFields();
+        }
+
+        public static final String[] PARAMETERS = new String[] {};
+
+        public abstract String execute();
+
+        public final Map<String, String> getDoc() {
+            return Map.of("cat", "meow", "dog", "woof");
+        }
+    }
+
+    public void testDocFields() {
+        DocFieldsTestScript.Factory f = scriptEngine.compile("test", "doc['cat'] + doc['dog']", DocFieldsTestScript.CONTEXT, Map.of());
+        assertThat(f.docFields(), equalTo(List.of("cat", "dog")));
+        assertThat(f.newInstance().execute(), equalTo("meowwoof"));
+    }
 }