Painless: Move Some Lookup Logic to PainlessLookup

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

@@ -96,7 +96,7 @@ public Class<?> findClass(String name) throws ClassNotFoundException {
             if (statefulFactoryClass != null && statefulFactoryClass.getName().equals(name)) {
                 return statefulFactoryClass;
             }
-            Class<?> found = painlessLookup.getClassFromBinaryName(name);
+            Class<?> found = painlessLookup.canonicalTypeNameToType(name.replace('$', '.'));
 
             return found != null ? found : super.findClass(name);
         }

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

@@ -187,7 +187,7 @@ static PainlessMethod lookupMethodInternal(PainlessLookup painlessLookup, Class<
         String key = PainlessLookupUtility.buildPainlessMethodKey(name, arity);
         // check whitelist for matching method
         for (Class<?> clazz = receiverClass; clazz != null; clazz = clazz.getSuperclass()) {
-            PainlessClass struct = painlessLookup.getPainlessStructFromJavaClass(clazz);
+            PainlessClass struct = painlessLookup.lookupPainlessClass(clazz);
 
             if (struct != null) {
                 PainlessMethod method = struct.methods.get(key);
@@ -197,7 +197,7 @@ static PainlessMethod lookupMethodInternal(PainlessLookup painlessLookup, Class<
             }
 
             for (Class<?> iface : clazz.getInterfaces()) {
-                struct = painlessLookup.getPainlessStructFromJavaClass(iface);
+                struct = painlessLookup.lookupPainlessClass(iface);
 
                 if (struct != null) {
                     PainlessMethod method = struct.methods.get(key);
@@ -326,8 +326,8 @@ static MethodHandle lookupMethod(PainlessLookup painlessLookup, MethodHandles.Lo
       */
     static MethodHandle lookupReference(PainlessLookup painlessLookup, MethodHandles.Lookup methodHandlesLookup, String interfaceClass,
                                         Class<?> receiverClass, String name) throws Throwable {
-         Class<?> interfaceType = painlessLookup.getJavaClassFromPainlessType(interfaceClass);
-         PainlessMethod interfaceMethod = painlessLookup.getPainlessStructFromJavaClass(interfaceType).functionalMethod;
+         Class<?> interfaceType = painlessLookup.canonicalTypeNameToType(interfaceClass);
+         PainlessMethod interfaceMethod = painlessLookup.lookupPainlessClass(interfaceType).functionalMethod;
          if (interfaceMethod == null) {
              throw new IllegalArgumentException("Class [" + interfaceClass + "] is not a functional interface");
          }
@@ -345,7 +345,7 @@ private static MethodHandle lookupReferenceInternal(PainlessLookup painlessLooku
          final FunctionRef ref;
          if ("this".equals(type)) {
              // user written method
-             PainlessMethod interfaceMethod = painlessLookup.getPainlessStructFromJavaClass(clazz).functionalMethod;
+             PainlessMethod interfaceMethod = painlessLookup.lookupPainlessClass(clazz).functionalMethod;
              if (interfaceMethod == null) {
                  throw new IllegalArgumentException("Cannot convert function reference [" + type + "::" + call + "] " +
                          "to [" + PainlessLookupUtility.typeToCanonicalTypeName(clazz) + "], not a functional interface");
@@ -419,7 +419,7 @@ public static String getUserFunctionHandleFieldName(String name, int arity) {
     static MethodHandle lookupGetter(PainlessLookup painlessLookup, Class<?> receiverClass, String name) {
         // first try whitelist
         for (Class<?> clazz = receiverClass; clazz != null; clazz = clazz.getSuperclass()) {
-            PainlessClass struct = painlessLookup.getPainlessStructFromJavaClass(clazz);
+            PainlessClass struct = painlessLookup.lookupPainlessClass(clazz);
 
             if (struct != null) {
                 MethodHandle handle = struct.getterMethodHandles.get(name);
@@ -429,7 +429,7 @@ static MethodHandle lookupGetter(PainlessLookup painlessLookup, Class<?> receive
             }
 
             for (final Class<?> iface : clazz.getInterfaces()) {
-                struct = painlessLookup.getPainlessStructFromJavaClass(iface);
+                struct = painlessLookup.lookupPainlessClass(iface);
 
                 if (struct != null) {
                     MethodHandle handle = struct.getterMethodHandles.get(name);
@@ -490,7 +490,7 @@ static MethodHandle lookupGetter(PainlessLookup painlessLookup, Class<?> receive
     static MethodHandle lookupSetter(PainlessLookup painlessLookup, Class<?> receiverClass, String name) {
         // first try whitelist
         for (Class<?> clazz = receiverClass; clazz != null; clazz = clazz.getSuperclass()) {
-            PainlessClass struct = painlessLookup.getPainlessStructFromJavaClass(clazz);
+            PainlessClass struct = painlessLookup.lookupPainlessClass(clazz);
 
             if (struct != null) {
                 MethodHandle handle = struct.setterMethodHandles.get(name);
@@ -500,7 +500,7 @@ static MethodHandle lookupSetter(PainlessLookup painlessLookup, Class<?> receive
             }
 
             for (final Class<?> iface : clazz.getInterfaces()) {
-                struct = painlessLookup.getPainlessStructFromJavaClass(iface);
+                struct = painlessLookup.lookupPainlessClass(iface);
 
                 if (struct != null) {
                     MethodHandle handle = struct.setterMethodHandles.get(name);

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

@@ -90,10 +90,10 @@ public static FunctionRef resolveFromLookup(
             PainlessLookup painlessLookup, Class<?> expected, String type, String call, int numCaptures) {
 
         if ("new".equals(call)) {
-            return new FunctionRef(expected, painlessLookup.getPainlessStructFromJavaClass(expected).functionalMethod,
+            return new FunctionRef(expected, painlessLookup.lookupPainlessClass(expected).functionalMethod,
                     lookup(painlessLookup, expected, type), numCaptures);
         } else {
-            return new FunctionRef(expected, painlessLookup.getPainlessStructFromJavaClass(expected).functionalMethod,
+            return new FunctionRef(expected, painlessLookup.lookupPainlessClass(expected).functionalMethod,
                     lookup(painlessLookup, expected, type, call, numCaptures > 0), numCaptures);
         }
     }
@@ -230,14 +230,14 @@ public FunctionRef(Class<?> expected,
     private static PainlessConstructor lookup(PainlessLookup painlessLookup, Class<?> expected, String type) {
         // check its really a functional interface
         // for e.g. Comparable
-        PainlessMethod method = painlessLookup.getPainlessStructFromJavaClass(expected).functionalMethod;
+        PainlessMethod method = painlessLookup.lookupPainlessClass(expected).functionalMethod;
         if (method == null) {
             throw new IllegalArgumentException("Cannot convert function reference [" + type + "::new] " +
                     "to [" + PainlessLookupUtility.typeToCanonicalTypeName(expected) + "], not a functional interface");
         }
 
         // lookup requested constructor
-        PainlessClass struct = painlessLookup.getPainlessStructFromJavaClass(painlessLookup.getJavaClassFromPainlessType(type));
+        PainlessClass struct = painlessLookup.lookupPainlessClass(painlessLookup.canonicalTypeNameToType(type));
         PainlessConstructor impl = struct.constructors.get(PainlessLookupUtility.buildPainlessConstructorKey(method.typeParameters.size()));
 
         if (impl == null) {
@@ -254,14 +254,14 @@ private static PainlessMethod lookup(PainlessLookup painlessLookup, Class<?> exp
                                          String type, String call, boolean receiverCaptured) {
         // check its really a functional interface
         // for e.g. Comparable
-        PainlessMethod method = painlessLookup.getPainlessStructFromJavaClass(expected).functionalMethod;
+        PainlessMethod method = painlessLookup.lookupPainlessClass(expected).functionalMethod;
         if (method == null) {
             throw new IllegalArgumentException("Cannot convert function reference [" + type + "::" + call + "] " +
                     "to [" + PainlessLookupUtility.typeToCanonicalTypeName(expected) + "], not a functional interface");
         }
 
         // lookup requested method
-        PainlessClass struct = painlessLookup.getPainlessStructFromJavaClass(painlessLookup.getJavaClassFromPainlessType(type));
+        PainlessClass struct = painlessLookup.lookupPainlessClass(painlessLookup.canonicalTypeNameToType(type));
         final PainlessMethod impl;
         // look for a static impl first
         PainlessMethod staticImpl =

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

@@ -57,7 +57,7 @@ public Map<String, List<String>> getHeaders(PainlessLookup painlessLookup) {
         if (objectToExplain != null) {
             toString = objectToExplain.toString();
             javaClassName = objectToExplain.getClass().getName();
-            PainlessClass struct = painlessLookup.getPainlessStructFromJavaClass(objectToExplain.getClass());
+            PainlessClass struct = painlessLookup.lookupPainlessClass(objectToExplain.getClass());
             if (struct != null) {
                 painlessClassName = PainlessLookupUtility.typeToCanonicalTypeName(objectToExplain.getClass());
             }

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

@@ -190,7 +190,7 @@ private static Class<?> definitionTypeForClass(PainlessLookup painlessLookup, Cl
             componentType = componentType.getComponentType();
         }
 
-        if (painlessLookup.getPainlessStructFromJavaClass(componentType) == null) {
+        if (painlessLookup.lookupPainlessClass(componentType) == null) {
             throw new IllegalArgumentException(unknownErrorMessageSource.apply(componentType));
         }
 

modules/lang-painless/src/main/java/org/elasticsearch/painless/antlr/EnhancedPainlessLexer.java

@@ -75,7 +75,7 @@ public void recover(final LexerNoViableAltException lnvae) {
 
     @Override
     protected boolean isType(String name) {
-        return painlessLookup.isSimplePainlessType(name);
+        return painlessLookup.isValidCanonicalClassName(name);
     }
 
     @Override

modules/lang-painless/src/main/java/org/elasticsearch/painless/lookup/PainlessLookup.java

@@ -19,41 +19,129 @@
 
 package org.elasticsearch.painless.lookup;
 
-import java.util.Collection;
 import java.util.Collections;
+import java.util.Comparator;
+import java.util.List;
 import java.util.Map;
+import java.util.Objects;
+import java.util.stream.Collectors;
 
-/**
- * The entire API for Painless.  Also used as a whitelist for checking for legal
- * methods and fields during at both compile-time and runtime.
- */
-public final class PainlessLookup {
+import static org.elasticsearch.painless.lookup.PainlessLookupUtility.buildPainlessConstructorKey;
+import static org.elasticsearch.painless.lookup.PainlessLookupUtility.buildPainlessFieldKey;
+import static org.elasticsearch.painless.lookup.PainlessLookupUtility.buildPainlessMethodKey;
+import static org.elasticsearch.painless.lookup.PainlessLookupUtility.typeToCanonicalTypeName;
 
-    public Collection<Class<?>> getStructs() {
-        return classesToPainlessClasses.keySet();
-    }
+public final class PainlessLookup {
 
     private final Map<String, Class<?>> canonicalClassNamesToClasses;
     private final Map<Class<?>, PainlessClass> classesToPainlessClasses;
 
+    private final List<Class<?>> sortedClassesByCanonicalClassName;
+
     PainlessLookup(Map<String, Class<?>> canonicalClassNamesToClasses, Map<Class<?>, PainlessClass> classesToPainlessClasses) {
+        Objects.requireNonNull(canonicalClassNamesToClasses);
+        Objects.requireNonNull(classesToPainlessClasses);
+
         this.canonicalClassNamesToClasses = Collections.unmodifiableMap(canonicalClassNamesToClasses);
         this.classesToPainlessClasses = Collections.unmodifiableMap(classesToPainlessClasses);
+
+        this.sortedClassesByCanonicalClassName = Collections.unmodifiableList(
+                classesToPainlessClasses.keySet().stream().sorted(
+                        Comparator.comparing(Class::getCanonicalName)
+                ).collect(Collectors.toList())
+        );
     }
 
-    public Class<?> getClassFromBinaryName(String painlessType) {
-        return canonicalClassNamesToClasses.get(painlessType.replace('$', '.'));
+    public boolean isValidCanonicalClassName(String canonicalClassName) {
+        Objects.requireNonNull(canonicalClassName);
+
+        return canonicalClassNamesToClasses.containsKey(canonicalClassName);
     }
 
-    public boolean isSimplePainlessType(String painlessType) {
-        return canonicalClassNamesToClasses.containsKey(painlessType);
+    public Class<?> canonicalTypeNameToType(String painlessType) {
+        Objects.requireNonNull(painlessType);
+
+        return PainlessLookupUtility.canonicalTypeNameToType(painlessType, canonicalClassNamesToClasses);
     }
 
-    public PainlessClass getPainlessStructFromJavaClass(Class<?> clazz) {
-        return classesToPainlessClasses.get(clazz);
+    public List<Class<?>> getSortedClassesByCanonicalClassName() {
+        return sortedClassesByCanonicalClassName;
     }
 
-    public Class<?> getJavaClassFromPainlessType(String painlessType) {
-        return PainlessLookupUtility.canonicalTypeNameToType(painlessType, canonicalClassNamesToClasses);
+    public PainlessClass lookupPainlessClass(Class<?> targetClass) {
+        return classesToPainlessClasses.get(targetClass);
+    }
+
+    public PainlessConstructor lookupPainlessConstructor(Class<?> targetClass, int constructorArity) {
+        Objects.requireNonNull(targetClass);
+
+        PainlessClass targetPainlessClass = classesToPainlessClasses.get(targetClass);
+        String painlessConstructorKey = buildPainlessConstructorKey(constructorArity);
+
+        if (targetPainlessClass == null) {
+            throw new IllegalArgumentException("target class [" + typeToCanonicalTypeName(targetClass) + "] " +
+                    "not found for constructor [" + painlessConstructorKey + "]");
+        }
+
+        PainlessConstructor painlessConstructor = targetPainlessClass.constructors.get(painlessConstructorKey);
+
+        if (painlessConstructor == null) {
+            throw new IllegalArgumentException(
+                    "constructor [" + typeToCanonicalTypeName(targetClass) + ", " + painlessConstructorKey + "] not found");
+        }
+
+        return painlessConstructor;
+    }
+
+    public PainlessMethod lookupPainlessMethod(Class<?> targetClass, boolean isStatic, String methodName, int methodArity) {
+        Objects.requireNonNull(targetClass);
+        Objects.requireNonNull(methodName);
+
+        if (targetClass.isPrimitive()) {
+            targetClass = PainlessLookupUtility.typeToBoxedType(targetClass);
+        }
+
+        PainlessClass targetPainlessClass = classesToPainlessClasses.get(targetClass);
+        String painlessMethodKey = buildPainlessMethodKey(methodName, methodArity);
+
+        if (targetPainlessClass == null) {
+            throw new IllegalArgumentException(
+                    "target class [" + typeToCanonicalTypeName(targetClass) + "] not found for method [" + painlessMethodKey + "]");
+        }
+
+        PainlessMethod painlessMethod = isStatic ?
+                targetPainlessClass.staticMethods.get(painlessMethodKey) :
+                targetPainlessClass.methods.get(painlessMethodKey);
+
+        if (painlessMethod == null) {
+            throw new IllegalArgumentException(
+                    "method [" + typeToCanonicalTypeName(targetClass) + ", " + painlessMethodKey + "] not found");
+        }
+
+        return painlessMethod;
+    }
+
+    public PainlessField lookupPainlessField(Class<?> targetClass, boolean isStatic, String fieldName) {
+        Objects.requireNonNull(targetClass);
+        Objects.requireNonNull(fieldName);
+
+        PainlessClass targetPainlessClass = classesToPainlessClasses.get(targetClass);
+        String painlessFieldKey = buildPainlessFieldKey(fieldName);
+
+        if (targetPainlessClass == null) {
+            throw new IllegalArgumentException(
+                    "target class [" + typeToCanonicalTypeName(targetClass) + "] not found for field [" + painlessFieldKey + "]");
+        }
+
+        PainlessField painlessField = isStatic ?
+                targetPainlessClass.staticFields.get(painlessFieldKey) :
+                targetPainlessClass.fields.get(painlessFieldKey);
+
+        if (painlessField == null) {
+            throw new IllegalArgumentException(
+                    "field [" + typeToCanonicalTypeName(targetClass) + ", " + painlessFieldKey + "] not found");
+        }
+
+        return painlessField;
     }
 }

modules/lang-painless/src/main/java/org/elasticsearch/painless/node/EExplicit.java

@@ -50,7 +50,7 @@ void extractVariables(Set<String> variables) {
     @Override
     void analyze(Locals locals) {
         try {
-            actual = locals.getPainlessLookup().getJavaClassFromPainlessType(type);
+            actual = locals.getPainlessLookup().canonicalTypeNameToType(type);
         } catch (IllegalArgumentException exception) {
             throw createError(new IllegalArgumentException("Not a type [" + type + "]."));
         }

modules/lang-painless/src/main/java/org/elasticsearch/painless/node/EFunctionRef.java

@@ -66,7 +66,7 @@ void analyze(Locals locals) {
             try {
                 if ("this".equals(type)) {
                     // user's own function
-                    PainlessMethod interfaceMethod = locals.getPainlessLookup().getPainlessStructFromJavaClass(expected).functionalMethod;
+                    PainlessMethod interfaceMethod = locals.getPainlessLookup().lookupPainlessClass(expected).functionalMethod;
                     if (interfaceMethod == null) {
                         throw new IllegalArgumentException("Cannot convert function reference [" + type + "::" + call + "] " +
                                 "to [" + PainlessLookupUtility.typeToCanonicalTypeName(expected) + "], not a functional interface");

modules/lang-painless/src/main/java/org/elasticsearch/painless/node/EInstanceof.java

@@ -58,7 +58,7 @@ void analyze(Locals locals) {
 
         // ensure the specified type is part of the definition
         try {
-            clazz = locals.getPainlessLookup().getJavaClassFromPainlessType(this.type);
+            clazz = locals.getPainlessLookup().canonicalTypeNameToType(this.type);
         } catch (IllegalArgumentException exception) {
             throw createError(new IllegalArgumentException("Not a type [" + this.type + "]."));
         }

modules/lang-painless/src/main/java/org/elasticsearch/painless/node/ELambda.java

@@ -121,7 +121,7 @@ void analyze(Locals locals) {
             }
         } else {
             // we know the method statically, infer return type and any unknown/def types
-            interfaceMethod = locals.getPainlessLookup().getPainlessStructFromJavaClass(expected).functionalMethod;
+            interfaceMethod = locals.getPainlessLookup().lookupPainlessClass(expected).functionalMethod;
             if (interfaceMethod == null) {
                 throw createError(new IllegalArgumentException("Cannot pass lambda to " +
                         "[" + PainlessLookupUtility.typeToCanonicalTypeName(expected) + "], not a functional interface"));

modules/lang-painless/src/main/java/org/elasticsearch/painless/node/EListInit.java

@@ -24,7 +24,6 @@
 import org.elasticsearch.painless.Location;
 import org.elasticsearch.painless.MethodWriter;
 import org.elasticsearch.painless.lookup.PainlessConstructor;
-import org.elasticsearch.painless.lookup.PainlessLookupUtility;
 import org.elasticsearch.painless.lookup.PainlessMethod;
 import org.elasticsearch.painless.lookup.def;
 import org.objectweb.asm.Type;
@@ -64,18 +63,16 @@ void analyze(Locals locals) {
 
         actual = ArrayList.class;
 
-        constructor = locals.getPainlessLookup().getPainlessStructFromJavaClass(actual).constructors.get(
-                PainlessLookupUtility.buildPainlessConstructorKey(0));
-
-        if (constructor == null) {
-            throw createError(new IllegalStateException("Illegal tree structure."));
+        try {
+            constructor = locals.getPainlessLookup().lookupPainlessConstructor(actual, 0);
+        } catch (IllegalArgumentException iae) {
+            throw createError(iae);
         }
 
-        method = locals.getPainlessLookup().getPainlessStructFromJavaClass(actual).methods
-                .get(PainlessLookupUtility.buildPainlessMethodKey("add", 1));
-
-        if (method == null) {
-            throw createError(new IllegalStateException("Illegal tree structure."));
+        try {
+            method = locals.getPainlessLookup().lookupPainlessMethod(actual, false, "add", 1);
+        } catch (IllegalArgumentException iae) {
+            throw createError(iae);
         }
 
         for (int index = 0; index < values.size(); ++index) {