Painless: Use LocalMethod Map For Lookup at Runtime (#32599)

This modifies Def to use a Map<String, LocalMethod> to look up user-defined methods at runtime 
instead of writing constant methodhandles to do the reverse lookup. This creates a consistency 
between how LocalMethods are looked up at compile-time and run-time. This consistency will allow 
this code to be more maintainable moving forward. This will also allow FunctionReference to be 
cleaned up in a follow up PR.

Author: jdconrad

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

@@ -20,6 +20,7 @@
 package org.elasticsearch.painless;
 
 import org.elasticsearch.bootstrap.BootstrapInfo;
+import org.elasticsearch.painless.Locals.LocalMethod;
 import org.elasticsearch.painless.antlr.Walker;
 import org.elasticsearch.painless.lookup.PainlessLookup;
 import org.elasticsearch.painless.node.SSource;
@@ -32,6 +33,7 @@
 import java.security.CodeSource;
 import java.security.SecureClassLoader;
 import java.security.cert.Certificate;
+import java.util.Map;
 import java.util.concurrent.atomic.AtomicInteger;
 
 import static org.elasticsearch.painless.WriterConstants.CLASS_NAME;
@@ -200,7 +202,7 @@ Constructor<?> compile(Loader loader, MainMethodReserved reserved, String name,
         ScriptClassInfo scriptClassInfo = new ScriptClassInfo(painlessLookup, scriptClass);
         SSource root = Walker.buildPainlessTree(scriptClassInfo, reserved, name, source, settings, painlessLookup,
                 null);
-        root.analyze(painlessLookup);
+        Map<String, LocalMethod> localMethods = root.analyze(painlessLookup);
         root.write();
 
         try {
@@ -209,6 +211,7 @@ Constructor<?> compile(Loader loader, MainMethodReserved reserved, String name,
             clazz.getField("$SOURCE").set(null, source);
             clazz.getField("$STATEMENTS").set(null, root.getStatements());
             clazz.getField("$DEFINITION").set(null, painlessLookup);
+            clazz.getField("$LOCALS").set(null, localMethods);
 
             return clazz.getConstructors()[0];
         } catch (Exception exception) { // Catch everything to let the user know this is something caused internally.

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

@@ -19,6 +19,7 @@
 
 package org.elasticsearch.painless;
 
+import org.elasticsearch.painless.Locals.LocalMethod;
 import org.elasticsearch.painless.lookup.PainlessClass;
 import org.elasticsearch.painless.lookup.PainlessLookup;
 import org.elasticsearch.painless.lookup.PainlessLookupUtility;
@@ -232,8 +233,10 @@ static PainlessMethod lookupMethodInternal(PainlessLookup painlessLookup, Class<
      * @throws IllegalArgumentException if no matching whitelisted method was found.
      * @throws Throwable if a method reference cannot be converted to an functional interface
      */
-    static MethodHandle lookupMethod(PainlessLookup painlessLookup, MethodHandles.Lookup methodHandlesLookup, MethodType callSiteType,
-                                     Class<?> receiverClass, String name, Object args[]) throws Throwable {
+    static MethodHandle lookupMethod(PainlessLookup painlessLookup, Map<String, LocalMethod> localMethods,
+            MethodHandles.Lookup methodHandlesLookup, MethodType callSiteType, Class<?> receiverClass, String name, Object args[])
+            throws Throwable {
+
          String recipeString = (String) args[0];
          int numArguments = callSiteType.parameterCount();
          // simple case: no lambdas
@@ -286,6 +289,7 @@ static MethodHandle lookupMethod(PainlessLookup painlessLookup, MethodHandles.Lo
                      // the implementation is strongly typed, now that we know the interface type,
                      // we have everything.
                      filter = lookupReferenceInternal(painlessLookup,
+                                                      localMethods,
                                                       methodHandlesLookup,
                                                       interfaceType,
                                                       type,
@@ -297,6 +301,7 @@ static MethodHandle lookupMethod(PainlessLookup painlessLookup, MethodHandles.Lo
                      // this cache). It won't blow up since we never nest here (just references)
                      MethodType nestedType = MethodType.methodType(interfaceType, captures);
                      CallSite nested = DefBootstrap.bootstrap(painlessLookup,
+                                                              localMethods,
                                                               methodHandlesLookup,
                                                               call,
                                                               nestedType,
@@ -324,24 +329,23 @@ static MethodHandle lookupMethod(PainlessLookup painlessLookup, MethodHandles.Lo
       * This is just like LambdaMetaFactory, only with a dynamic type. The interface type is known,
       * so we simply need to lookup the matching implementation method based on receiver type.
       */
-    static MethodHandle lookupReference(PainlessLookup painlessLookup, MethodHandles.Lookup methodHandlesLookup, String interfaceClass,
-                                        Class<?> receiverClass, String name) throws Throwable {
+    static MethodHandle lookupReference(PainlessLookup painlessLookup, Map<String, LocalMethod> localMethods,
+            MethodHandles.Lookup methodHandlesLookup, String interfaceClass, Class<?> receiverClass, String name) throws Throwable {
          Class<?> interfaceType = painlessLookup.canonicalTypeNameToType(interfaceClass);
          PainlessMethod interfaceMethod = painlessLookup.lookupPainlessClass(interfaceType).functionalMethod;
          if (interfaceMethod == null) {
              throw new IllegalArgumentException("Class [" + interfaceClass + "] is not a functional interface");
          }
          int arity = interfaceMethod.typeParameters.size();
          PainlessMethod implMethod = lookupMethodInternal(painlessLookup, receiverClass, name, arity);
-        return lookupReferenceInternal(painlessLookup, methodHandlesLookup, interfaceType,
-                PainlessLookupUtility.typeToCanonicalTypeName(implMethod.targetClass),
+        return lookupReferenceInternal(painlessLookup, localMethods, methodHandlesLookup,
+                interfaceType, PainlessLookupUtility.typeToCanonicalTypeName(implMethod.targetClass),
                 implMethod.javaMethod.getName(), receiverClass);
      }
 
      /** Returns a method handle to an implementation of clazz, given method reference signature. */
-    private static MethodHandle lookupReferenceInternal(PainlessLookup painlessLookup, MethodHandles.Lookup methodHandlesLookup,
-                                                        Class<?> clazz, String type, String call, Class<?>... captures)
-            throws Throwable {
+    private static MethodHandle lookupReferenceInternal(PainlessLookup painlessLookup, Map<String, LocalMethod> localMethods,
+            MethodHandles.Lookup methodHandlesLookup, Class<?> clazz, String type, String call, Class<?>... captures) throws Throwable {
          final FunctionRef ref;
          if ("this".equals(type)) {
              // user written method
@@ -351,13 +355,8 @@ private static MethodHandle lookupReferenceInternal(PainlessLookup painlessLooku
                          "to [" + PainlessLookupUtility.typeToCanonicalTypeName(clazz) + "], not a functional interface");
              }
              int arity = interfaceMethod.typeParameters.size() + captures.length;
-             final MethodHandle handle;
-             try {
-                 MethodHandle accessor = methodHandlesLookup.findStaticGetter(methodHandlesLookup.lookupClass(),
-                                                                              getUserFunctionHandleFieldName(call, arity),
-                                                                              MethodHandle.class);
-                 handle = (MethodHandle)accessor.invokeExact();
-             } catch (NoSuchFieldException | IllegalAccessException e) {
+             LocalMethod localMethod = localMethods.get(Locals.buildLocalMethodKey(call, arity));
+             if (localMethod == null) {
                  // is it a synthetic method? If we generated the method ourselves, be more helpful. It can only fail
                  // because the arity does not match the expected interface type.
                  if (call.contains("$")) {
@@ -366,7 +365,7 @@ private static MethodHandle lookupReferenceInternal(PainlessLookup painlessLooku
                  }
                  throw new IllegalArgumentException("Unknown call [" + call + "] with [" + arity + "] arguments.");
              }
-             ref = new FunctionRef(clazz, interfaceMethod, call, handle.type(), captures.length);
+             ref = new FunctionRef(clazz, interfaceMethod, call, localMethod.methodType, captures.length);
          } else {
              // whitelist lookup
              ref = FunctionRef.resolveFromLookup(painlessLookup, clazz, type, call, captures.length);
@@ -385,11 +384,6 @@ private static MethodHandle lookupReferenceInternal(PainlessLookup painlessLooku
          return callSite.dynamicInvoker().asType(MethodType.methodType(clazz, captures));
      }
 
-     /** gets the field name used to lookup up the MethodHandle for a function. */
-     public static String getUserFunctionHandleFieldName(String name, int arity) {
-         return "handle$" + name + "$" + arity;
-     }
-
     /**
      * Looks up handle for a dynamic field getter (field load)
      * <p>

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

@@ -20,6 +20,7 @@
 package org.elasticsearch.painless;
 
 import org.elasticsearch.common.SuppressForbidden;
+import org.elasticsearch.painless.Locals.LocalMethod;
 import org.elasticsearch.painless.lookup.PainlessLookup;
 
 import java.lang.invoke.CallSite;
@@ -28,6 +29,7 @@
 import java.lang.invoke.MethodType;
 import java.lang.invoke.MutableCallSite;
 import java.lang.invoke.WrongMethodTypeException;
+import java.util.Map;
 
 /**
  * Painless invokedynamic bootstrap for the call site.
@@ -105,19 +107,21 @@ static final class PIC extends MutableCallSite {
         static final int MAX_DEPTH = 5;
 
         private final PainlessLookup painlessLookup;
+        private final Map<String, LocalMethod> localMethods;
         private final MethodHandles.Lookup methodHandlesLookup;
         private final String name;
         private final int flavor;
         private final Object[] args;
         int depth; // pkg-protected for testing
 
-        PIC(PainlessLookup painlessLookup, MethodHandles.Lookup methodHandlesLookup,
-                String name, MethodType type, int initialDepth, int flavor, Object[] args) {
+        PIC(PainlessLookup painlessLookup, Map<String, LocalMethod> localMethods,
+                MethodHandles.Lookup methodHandlesLookup, String name, MethodType type, int initialDepth, int flavor, Object[] args) {
             super(type);
             if (type.parameterType(0) != Object.class) {
                 throw new BootstrapMethodError("The receiver type (1st arg) of invokedynamic descriptor must be Object.");
             }
             this.painlessLookup = painlessLookup;
+            this.localMethods = localMethods;
             this.methodHandlesLookup = methodHandlesLookup;
             this.name = name;
             this.flavor = flavor;
@@ -145,7 +149,7 @@ static boolean checkClass(Class<?> clazz, Object receiver) {
         private MethodHandle lookup(int flavor, String name, Class<?> receiver) throws Throwable {
             switch(flavor) {
                 case METHOD_CALL:
-                    return Def.lookupMethod(painlessLookup, methodHandlesLookup, type(), receiver, name, args);
+                    return Def.lookupMethod(painlessLookup, localMethods, methodHandlesLookup, type(), receiver, name, args);
                 case LOAD:
                     return Def.lookupGetter(painlessLookup, receiver, name);
                 case STORE:
@@ -157,7 +161,7 @@ private MethodHandle lookup(int flavor, String name, Class<?> receiver) throws T
                 case ITERATOR:
                     return Def.lookupIterator(receiver);
                 case REFERENCE:
-                    return Def.lookupReference(painlessLookup, methodHandlesLookup, (String) args[0], receiver, name);
+                    return Def.lookupReference(painlessLookup, localMethods, methodHandlesLookup, (String) args[0], receiver, name);
                 case INDEX_NORMALIZE:
                     return Def.lookupIndexNormalize(receiver);
                 default: throw new AssertionError();
@@ -432,8 +436,9 @@ static boolean checkBoth(Class<?> left, Class<?> right, Object leftObject, Objec
      * <p>
      * see https://docs.oracle.com/javase/specs/jvms/se7/html/jvms-6.html#jvms-6.5.invokedynamic
      */
-    public static CallSite bootstrap(PainlessLookup painlessLookup, MethodHandles.Lookup methodHandlesLookup, String name,
-                                     MethodType type, int initialDepth, int flavor, Object... args) {
+    @SuppressWarnings("unchecked")
+    public static CallSite bootstrap(PainlessLookup painlessLookup, Map<String, LocalMethod> localMethods,
+            MethodHandles.Lookup methodHandlesLookup, String name, MethodType type, int initialDepth, int flavor, Object... args) {
         // validate arguments
         switch(flavor) {
             // "function-call" like things get a polymorphic cache
@@ -452,7 +457,7 @@ public static CallSite bootstrap(PainlessLookup painlessLookup, MethodHandles.Lo
                 if (args.length != numLambdas + 1) {
                     throw new BootstrapMethodError("Illegal number of parameters: expected " + numLambdas + " references");
                 }
-                return new PIC(painlessLookup, methodHandlesLookup, name, type, initialDepth, flavor, args);
+                return new PIC(painlessLookup, localMethods, methodHandlesLookup, name, type, initialDepth, flavor, args);
             case LOAD:
             case STORE:
             case ARRAY_LOAD:
@@ -462,15 +467,15 @@ public static CallSite bootstrap(PainlessLookup painlessLookup, MethodHandles.Lo
                 if (args.length > 0) {
                     throw new BootstrapMethodError("Illegal static bootstrap parameters for flavor: " + flavor);
                 }
-                return new PIC(painlessLookup, methodHandlesLookup, name, type, initialDepth, flavor, args);
+                return new PIC(painlessLookup, localMethods, methodHandlesLookup, name, type, initialDepth, flavor, args);
             case REFERENCE:
                 if (args.length != 1) {
                     throw new BootstrapMethodError("Invalid number of parameters for reference call");
                 }
                 if (args[0] instanceof String == false) {
                     throw new BootstrapMethodError("Illegal parameter for reference call: " + args[0]);
                 }
-                return new PIC(painlessLookup, methodHandlesLookup, name, type, initialDepth, flavor, args);
+                return new PIC(painlessLookup, localMethods, methodHandlesLookup, name, type, initialDepth, flavor, args);
 
             // operators get monomorphic cache, with a generic impl for a fallback
             case UNARY_OPERATOR:

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

@@ -32,6 +32,9 @@
 import java.util.List;
 import java.util.Map;
 import java.util.Set;
+import java.util.stream.Collectors;
+
+import static org.elasticsearch.painless.lookup.PainlessLookupUtility.typeToJavaType;
 
 /**
  * Tracks user defined methods and variables across compilation phases.
@@ -74,17 +77,24 @@ public LocalMethod(String name, Class<?> returnType, List<Class<?>> typeParamete
 
     /** Creates a new local variable scope (e.g. loop) inside the current scope */
     public static Locals newLocalScope(Locals currentScope) {
-        return new Locals(currentScope);
+        Locals locals = new Locals(currentScope);
+        locals.methods = currentScope.methods;
+
+        return locals;
     }
 
     /**
      * Creates a new lambda scope inside the current scope
      * <p>
      * This is just like {@link #newFunctionScope}, except the captured parameters are made read-only.
      */
-    public static Locals newLambdaScope(Locals programScope, Class<?> returnType, List<Parameter> parameters,
+    public static Locals newLambdaScope(Locals programScope, String name, Class<?> returnType, List<Parameter> parameters,
                                         int captureCount, int maxLoopCounter) {
         Locals locals = new Locals(programScope, programScope.painlessLookup, returnType, KEYWORDS);
+        locals.methods = programScope.methods;
+        List<Class<?>> typeParameters = parameters.stream().map(parameter -> typeToJavaType(parameter.clazz)).collect(Collectors.toList());
+        locals.methods.put(buildLocalMethodKey(name, parameters.size()), new LocalMethod(name, returnType, typeParameters,
+                MethodType.methodType(typeToJavaType(returnType), typeParameters)));
         for (int i = 0; i < parameters.size(); i++) {
             Parameter parameter = parameters.get(i);
             // TODO: allow non-captures to be r/w:
@@ -104,6 +114,7 @@ public static Locals newLambdaScope(Locals programScope, Class<?> returnType, Li
     /** Creates a new function scope inside the current scope */
     public static Locals newFunctionScope(Locals programScope, Class<?> returnType, List<Parameter> parameters, int maxLoopCounter) {
         Locals locals = new Locals(programScope, programScope.painlessLookup, returnType, KEYWORDS);
+        locals.methods = programScope.methods;
         for (Parameter parameter : parameters) {
             locals.addVariable(parameter.location, parameter.clazz, parameter.name, false);
         }
@@ -118,6 +129,7 @@ public static Locals newFunctionScope(Locals programScope, Class<?> returnType,
     public static Locals newMainMethodScope(ScriptClassInfo scriptClassInfo, Locals programScope, int maxLoopCounter) {
         Locals locals = new Locals(
             programScope, programScope.painlessLookup, scriptClassInfo.getExecuteMethodReturnType(), KEYWORDS);
+        locals.methods = programScope.methods;
         // This reference. Internal use only.
         locals.defineVariable(null, Object.class, THIS, true);
 
@@ -136,6 +148,7 @@ public static Locals newMainMethodScope(ScriptClassInfo scriptClassInfo, Locals
     /** Creates a new program scope: the list of methods. It is the parent for all methods */
     public static Locals newProgramScope(PainlessLookup painlessLookup, Collection<LocalMethod> methods) {
         Locals locals = new Locals(null, painlessLookup, null, null);
+        locals.methods = new HashMap<>();
         for (LocalMethod method : methods) {
             locals.addMethod(method);
         }
@@ -167,15 +180,8 @@ public Variable getVariable(Location location, String name) {
     }
 
     /** Looks up a method. Returns null if the method does not exist. */
-    public LocalMethod getMethod(String key) {
-        LocalMethod method = lookupMethod(key);
-        if (method != null) {
-            return method;
-        }
-        if (parent != null) {
-            return parent.getMethod(key);
-        }
-        return null;
+    public LocalMethod getMethod(String methodName, int methodArity) {
+        return methods.get(buildLocalMethodKey(methodName, methodArity));
     }
 
     /** Creates a new variable. Throws IAE if the variable has already been defined (even in a parent) or reserved. */
@@ -260,15 +266,10 @@ private Variable lookupVariable(Location location, String name) {
         return variables.get(name);
     }
 
-    /** Looks up a method at this scope only. Returns null if the method does not exist. */
-    private LocalMethod lookupMethod(String key) {
-        if (methods == null) {
-            return null;
-        }
-        return methods.get(key);
+    public Map<String, LocalMethod> getMethods() {
+        return Collections.unmodifiableMap(methods);
     }
 
-
     /** Defines a variable at this scope internally. */
     private Variable defineVariable(Location location, Class<?> type, String name, boolean readonly) {
         if (variables == null) {
@@ -281,14 +282,9 @@ private Variable defineVariable(Location location, Class<?> type, String name, b
     }
 
     private void addMethod(LocalMethod method) {
-        if (methods == null) {
-            methods = new HashMap<>();
-        }
         methods.put(buildLocalMethodKey(method.name, method.typeParameters.size()), method);
-        // TODO: check result
     }
 
-
     private int getNextSlot() {
         return nextSlotNumber;
     }