Painless: Add PainlessClassBuilder (#32141)

Several pieces of data in PainlessClass cannot be passed in at the time the 
PainlessClass is created so it must be "frozen" after all the data is collected. This means 
PainlessClass is currently serving two functions as both a builder and a set of data. This 
separates the two pieces into clearly distinct values.

This change also removes the PainlessMethodKey in favor of a simple String. The goal is 
to have the painless method key be completely internal to the PainlessLookup eventually 
and this simplifies the way there. Note that this was added since PainlessClass and 
PainlessClassBuilder were already being changed instead of a follow up PR.

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

@@ -23,7 +23,6 @@ import org.elasticsearch.painless.lookup.PainlessClass;
 import org.elasticsearch.painless.lookup.PainlessLookup;
 import org.elasticsearch.painless.lookup.PainlessLookupUtility;
 import org.elasticsearch.painless.lookup.PainlessMethod;
-import org.elasticsearch.painless.lookup.PainlessMethodKey;
 
 import java.lang.invoke.CallSite;
 import java.lang.invoke.MethodHandle;
@@ -185,7 +184,7 @@ public final class Def {
      * @throws IllegalArgumentException if no matching whitelisted method was found.
      */
     static PainlessMethod lookupMethodInternal(PainlessLookup painlessLookup, Class<?> receiverClass, String name, int arity) {
-        PainlessMethodKey key = new PainlessMethodKey(name, arity);
+        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);

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

@@ -23,7 +23,6 @@ import org.elasticsearch.painless.lookup.PainlessClass;
 import org.elasticsearch.painless.lookup.PainlessLookup;
 import org.elasticsearch.painless.lookup.PainlessLookupUtility;
 import org.elasticsearch.painless.lookup.PainlessMethod;
-import org.elasticsearch.painless.lookup.PainlessMethodKey;
 import org.objectweb.asm.Type;
 
 import java.lang.invoke.MethodType;
@@ -177,10 +176,11 @@ public class FunctionRef {
         final PainlessMethod impl;
         // ctor ref
         if ("new".equals(call)) {
-            impl = struct.constructors.get(new PainlessMethodKey("<init>", method.arguments.size()));
+            impl = struct.constructors.get(PainlessLookupUtility.buildPainlessMethodKey("<init>", method.arguments.size()));
         } else {
             // look for a static impl first
-            PainlessMethod staticImpl = struct.staticMethods.get(new PainlessMethodKey(call, method.arguments.size()));
+            PainlessMethod staticImpl =
+                    struct.staticMethods.get(PainlessLookupUtility.buildPainlessMethodKey(call, method.arguments.size()));
             if (staticImpl == null) {
                 // otherwise a virtual impl
                 final int arity;
@@ -191,7 +191,7 @@ public class FunctionRef {
                     // receiver passed
                     arity = method.arguments.size() - 1;
                 }
-                impl = struct.methods.get(new PainlessMethodKey(call, arity));
+                impl = struct.methods.get(PainlessLookupUtility.buildPainlessMethodKey(call, arity));
             } else {
                 impl = staticImpl;
             }

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

@@ -23,7 +23,6 @@ import org.elasticsearch.painless.ScriptClassInfo.MethodArgument;
 import org.elasticsearch.painless.lookup.PainlessLookup;
 import org.elasticsearch.painless.lookup.PainlessLookupUtility;
 import org.elasticsearch.painless.lookup.PainlessMethod;
-import org.elasticsearch.painless.lookup.PainlessMethodKey;
 
 import java.util.Arrays;
 import java.util.Collection;
@@ -144,7 +143,7 @@ public final class Locals {
     }
 
     /** Looks up a method. Returns null if the method does not exist. */
-    public PainlessMethod getMethod(PainlessMethodKey key) {
+    public PainlessMethod getMethod(String key) {
         PainlessMethod method = lookupMethod(key);
         if (method != null) {
             return method;
@@ -200,7 +199,7 @@ public final class Locals {
     // variable name -> variable
     private Map<String,Variable> variables;
     // method name+arity -> methods
-    private Map<PainlessMethodKey,PainlessMethod> methods;
+    private Map<String,PainlessMethod> methods;
 
     /**
      * Create a new Locals
@@ -238,7 +237,7 @@ public final class Locals {
     }
 
     /** Looks up a method at this scope only. Returns null if the method does not exist. */
-    private PainlessMethod lookupMethod(PainlessMethodKey key) {
+    private PainlessMethod lookupMethod(String key) {
         if (methods == null) {
             return null;
         }
@@ -261,7 +260,7 @@ public final class Locals {
         if (methods == null) {
             methods = new HashMap<>();
         }
-        methods.put(new PainlessMethodKey(method.name, method.arguments.size()), method);
+        methods.put(PainlessLookupUtility.buildPainlessMethodKey(method.name, method.arguments.size()), method);
         // TODO: check result
     }
 

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

@@ -19,19 +19,20 @@
 
 package org.elasticsearch.painless.lookup;
 
+import org.objectweb.asm.Type;
+
 import java.lang.invoke.MethodHandle;
 import java.util.Collections;
-import java.util.HashMap;
 import java.util.Map;
 
 public final class PainlessClass {
     public final String name;
     public final Class<?> clazz;
-    public final org.objectweb.asm.Type type;
+    public final Type type;
 
-    public final Map<PainlessMethodKey, PainlessMethod> constructors;
+    public final Map<String, PainlessMethod> constructors;
-    public final Map<PainlessMethodKey, PainlessMethod> staticMethods;
+    public final Map<String, PainlessMethod> staticMethods;
-    public final Map<PainlessMethodKey, PainlessMethod> methods;
+    public final Map<String, PainlessMethod> methods;
 
     public final Map<String, PainlessField> staticMembers;
     public final Map<String, PainlessField> members;
@@ -41,63 +42,25 @@ public final class PainlessClass {
 
     public final PainlessMethod functionalMethod;
 
-    PainlessClass(String name, Class<?> clazz, org.objectweb.asm.Type type) {
+    PainlessClass(String name, Class<?> clazz, Type type,
+                  Map<String, PainlessMethod> constructors, Map<String, PainlessMethod> staticMethods, Map<String, PainlessMethod> methods,
+                  Map<String, PainlessField> staticMembers, Map<String, PainlessField> members,
+                  Map<String, MethodHandle> getters, Map<String, MethodHandle> setters,
+                  PainlessMethod functionalMethod) {
         this.name = name;
         this.clazz = clazz;
         this.type = type;
 
-        constructors = new HashMap<>();
+        this.constructors = Collections.unmodifiableMap(constructors);
-        staticMethods = new HashMap<>();
+        this.staticMethods = Collections.unmodifiableMap(staticMethods);
-        methods = new HashMap<>();
+        this.methods = Collections.unmodifiableMap(methods);
 
-        staticMembers = new HashMap<>();
+        this.staticMembers = Collections.unmodifiableMap(staticMembers);
-        members = new HashMap<>();
+        this.members = Collections.unmodifiableMap(members);
 
-        getters = new HashMap<>();
+        this.getters = Collections.unmodifiableMap(getters);
-        setters = new HashMap<>();
+        this.setters = Collections.unmodifiableMap(setters);
-
-        functionalMethod = null;
-    }
-
-    private PainlessClass(PainlessClass struct, PainlessMethod functionalMethod) {
-        name = struct.name;
-        clazz = struct.clazz;
-        type = struct.type;
-
-        constructors = Collections.unmodifiableMap(struct.constructors);
-        staticMethods = Collections.unmodifiableMap(struct.staticMethods);
-        methods = Collections.unmodifiableMap(struct.methods);
-
-        staticMembers = Collections.unmodifiableMap(struct.staticMembers);
-        members = Collections.unmodifiableMap(struct.members);
-
-        getters = Collections.unmodifiableMap(struct.getters);
-        setters = Collections.unmodifiableMap(struct.setters);
 
         this.functionalMethod = functionalMethod;
     }
-
-    public PainlessClass freeze(PainlessMethod functionalMethod) {
-        return new PainlessClass(this, functionalMethod);
-    }
-
-    @Override
-    public boolean equals(Object object) {
-        if (this == object) {
-            return true;
-        }
-
-        if (object == null || getClass() != object.getClass()) {
-            return false;
-        }
-
-        PainlessClass struct = (PainlessClass)object;
-
-        return name.equals(struct.name);
-    }
-
-    @Override
-    public int hashCode() {
-        return name.hashCode();
-    }
 }

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

@@ -0,0 +1,70 @@
+/*
+ * Licensed to Elasticsearch under one or more contributor
+ * license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright
+ * ownership. Elasticsearch licenses this file to you under
+ * the Apache License, Version 2.0 (the "License"); you may
+ * not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.elasticsearch.painless.lookup;
+
+import org.objectweb.asm.Type;
+
+import java.lang.invoke.MethodHandle;
+import java.util.HashMap;
+import java.util.Map;
+
+final class PainlessClassBuilder {
+    final String name;
+    final Class<?> clazz;
+    final Type type;
+
+    final Map<String, PainlessMethod> constructors;
+    final Map<String, PainlessMethod> staticMethods;
+    final Map<String, PainlessMethod> methods;
+
+    final Map<String, PainlessField> staticMembers;
+    final Map<String, PainlessField> members;
+
+    final Map<String, MethodHandle> getters;
+    final Map<String, MethodHandle> setters;
+
+    PainlessMethod functionalMethod;
+
+    PainlessClassBuilder(String name, Class<?> clazz, Type type) {
+        this.name = name;
+        this.clazz = clazz;
+        this.type = type;
+
+        constructors = new HashMap<>();
+        staticMethods = new HashMap<>();
+        methods = new HashMap<>();
+
+        staticMembers = new HashMap<>();
+        members = new HashMap<>();
+
+        getters = new HashMap<>();
+        setters = new HashMap<>();
+
+        functionalMethod = null;
+    }
+
+    PainlessClass build() {
+        return new PainlessClass(name, clazz, type,
+            constructors, staticMethods, methods,
+            staticMembers, members,
+            getters, setters,
+            functionalMethod);
+    }
+}

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

@@ -37,6 +37,8 @@ import java.util.Map;
 import java.util.Stack;
 import java.util.regex.Pattern;
 
+import static org.elasticsearch.painless.lookup.PainlessLookupUtility.buildPainlessMethodKey;
+
 public class PainlessLookupBuilder {
     private static final Pattern TYPE_NAME_PATTERN = Pattern.compile("^[_a-zA-Z][._a-zA-Z0-9]*$");
 
@@ -60,16 +62,16 @@ public class PainlessLookupBuilder {
     }
 
     private final Map<String, Class<?>> painlessTypesToJavaClasses;
-    private final Map<Class<?>, PainlessClass> javaClassesToPainlessStructs;
+    private final Map<Class<?>, PainlessClassBuilder> javaClassesToPainlessClassBuilders;
 
     public PainlessLookupBuilder(List<Whitelist> whitelists) {
         painlessTypesToJavaClasses = new HashMap<>();
-        javaClassesToPainlessStructs = new HashMap<>();
+        javaClassesToPainlessClassBuilders = new HashMap<>();
 
         String origin = null;
 
         painlessTypesToJavaClasses.put("def", def.class);
-        javaClassesToPainlessStructs.put(def.class, new PainlessClass("def", Object.class, Type.getType(Object.class)));
+        javaClassesToPainlessClassBuilders.put(def.class, new PainlessClassBuilder("def", Object.class, Type.getType(Object.class)));
 
         try {
             // first iteration collects all the Painless type names that
@@ -77,7 +79,8 @@ public class PainlessLookupBuilder {
             for (Whitelist whitelist : whitelists) {
                 for (WhitelistClass whitelistStruct : whitelist.whitelistStructs) {
                     String painlessTypeName = whitelistStruct.javaClassName.replace('$', '.');
-                    PainlessClass painlessStruct = javaClassesToPainlessStructs.get(painlessTypesToJavaClasses.get(painlessTypeName));
+                    PainlessClassBuilder painlessStruct =
+                            javaClassesToPainlessClassBuilders.get(painlessTypesToJavaClasses.get(painlessTypeName));
 
                     if (painlessStruct != null && painlessStruct.clazz.getName().equals(whitelistStruct.javaClassName) == false) {
                         throw new IllegalArgumentException("struct [" + painlessStruct.name + "] cannot represent multiple classes " +
@@ -87,8 +90,8 @@ public class PainlessLookupBuilder {
                     origin = whitelistStruct.origin;
                     addStruct(whitelist.javaClassLoader, whitelistStruct);
 
-                    painlessStruct = javaClassesToPainlessStructs.get(painlessTypesToJavaClasses.get(painlessTypeName));
+                    painlessStruct = javaClassesToPainlessClassBuilders.get(painlessTypesToJavaClasses.get(painlessTypeName));
-                    javaClassesToPainlessStructs.put(painlessStruct.clazz, painlessStruct);
+                    javaClassesToPainlessClassBuilders.put(painlessStruct.clazz, painlessStruct);
                 }
             }
 
@@ -121,8 +124,8 @@ public class PainlessLookupBuilder {
 
         // goes through each Painless struct and determines the inheritance list,
         // and then adds all inherited types to the Painless struct's whitelist
-        for (Class<?> javaClass : javaClassesToPainlessStructs.keySet()) {
+        for (Class<?> javaClass : javaClassesToPainlessClassBuilders.keySet()) {
-            PainlessClass painlessStruct = javaClassesToPainlessStructs.get(javaClass);
+            PainlessClassBuilder painlessStruct = javaClassesToPainlessClassBuilders.get(javaClass);
 
             List<String> painlessSuperStructs = new ArrayList<>();
             Class<?> javaSuperClass = painlessStruct.clazz.getSuperclass();
@@ -133,7 +136,7 @@ public class PainlessLookupBuilder {
             // adds super classes to the inheritance list
             if (javaSuperClass != null && javaSuperClass.isInterface() == false) {
                 while (javaSuperClass != null) {
-                    PainlessClass painlessSuperStruct = javaClassesToPainlessStructs.get(javaSuperClass);
+                    PainlessClassBuilder painlessSuperStruct = javaClassesToPainlessClassBuilders.get(javaSuperClass);
 
                     if (painlessSuperStruct != null) {
                         painlessSuperStructs.add(painlessSuperStruct.name);
@@ -149,7 +152,7 @@ public class PainlessLookupBuilder {
                 Class<?> javaInterfaceLookup = javaInteraceLookups.pop();
 
                 for (Class<?> javaSuperInterface : javaInterfaceLookup.getInterfaces()) {
-                    PainlessClass painlessInterfaceStruct = javaClassesToPainlessStructs.get(javaSuperInterface);
+                    PainlessClassBuilder painlessInterfaceStruct = javaClassesToPainlessClassBuilders.get(javaSuperInterface);
 
                     if (painlessInterfaceStruct != null) {
                         String painlessInterfaceStructName = painlessInterfaceStruct.name;
@@ -170,7 +173,7 @@ public class PainlessLookupBuilder {
 
             // copies methods and fields from Object into interface types
             if (painlessStruct.clazz.isInterface() || (def.class.getSimpleName()).equals(painlessStruct.name)) {
-                PainlessClass painlessObjectStruct = javaClassesToPainlessStructs.get(Object.class);
+                PainlessClassBuilder painlessObjectStruct = javaClassesToPainlessClassBuilders.get(Object.class);
 
                 if (painlessObjectStruct != null) {
                     copyStruct(painlessStruct.name, Collections.singletonList(painlessObjectStruct.name));
@@ -179,14 +182,9 @@ public class PainlessLookupBuilder {
         }
 
         // precompute runtime classes
-        for (PainlessClass painlessStruct : javaClassesToPainlessStructs.values()) {
+        for (PainlessClassBuilder painlessStruct : javaClassesToPainlessClassBuilders.values()) {
             addRuntimeClass(painlessStruct);
         }
-
-        // copy all structs to make them unmodifiable for outside users:
-        for (Map.Entry<Class<?>,PainlessClass> entry : javaClassesToPainlessStructs.entrySet()) {
-            entry.setValue(entry.getValue().freeze(computeFunctionalInterfaceMethod(entry.getValue())));
-        }
     }
 
     private void addStruct(ClassLoader whitelistClassLoader, WhitelistClass whitelistStruct) {
@@ -223,12 +221,12 @@ public class PainlessLookupBuilder {
             }
         }
 
-        PainlessClass existingStruct = javaClassesToPainlessStructs.get(javaClass);
+        PainlessClassBuilder existingStruct = javaClassesToPainlessClassBuilders.get(javaClass);
 
         if (existingStruct == null) {
-            PainlessClass struct = new PainlessClass(painlessTypeName, javaClass, org.objectweb.asm.Type.getType(javaClass));
+            PainlessClassBuilder struct = new PainlessClassBuilder(painlessTypeName, javaClass, org.objectweb.asm.Type.getType(javaClass));
             painlessTypesToJavaClasses.put(painlessTypeName, javaClass);
-            javaClassesToPainlessStructs.put(javaClass, struct);
+            javaClassesToPainlessClassBuilders.put(javaClass, struct);
         } else if (existingStruct.clazz.equals(javaClass) == false) {
             throw new IllegalArgumentException("struct [" + painlessTypeName + "] is used to " +
                 "illegally represent multiple java classes [" + whitelistStruct.javaClassName + "] and " +
@@ -261,7 +259,7 @@ public class PainlessLookupBuilder {
     }
 
     private void addConstructor(String ownerStructName, WhitelistConstructor whitelistConstructor) {
-        PainlessClass ownerStruct = javaClassesToPainlessStructs.get(painlessTypesToJavaClasses.get(ownerStructName));
+        PainlessClassBuilder ownerStruct = javaClassesToPainlessClassBuilders.get(painlessTypesToJavaClasses.get(ownerStructName));
 
         if (ownerStruct == null) {
             throw new IllegalArgumentException("owner struct [" + ownerStructName + "] not defined for constructor with " +
@@ -295,7 +293,7 @@ public class PainlessLookupBuilder {
                 " with constructor parameters " + whitelistConstructor.painlessParameterTypeNames, exception);
         }
 
-        PainlessMethodKey painlessMethodKey = new PainlessMethodKey("<init>", whitelistConstructor.painlessParameterTypeNames.size());
+        String painlessMethodKey = buildPainlessMethodKey("<init>", whitelistConstructor.painlessParameterTypeNames.size());
         PainlessMethod painlessConstructor = ownerStruct.constructors.get(painlessMethodKey);
 
         if (painlessConstructor == null) {
@@ -321,7 +319,7 @@ public class PainlessLookupBuilder {
     }
 
     private void addMethod(ClassLoader whitelistClassLoader, String ownerStructName, WhitelistMethod whitelistMethod) {
-        PainlessClass ownerStruct = javaClassesToPainlessStructs.get(painlessTypesToJavaClasses.get(ownerStructName));
+        PainlessClassBuilder ownerStruct = javaClassesToPainlessClassBuilders.get(painlessTypesToJavaClasses.get(ownerStructName));
 
         if (ownerStruct == null) {
             throw new IllegalArgumentException("owner struct [" + ownerStructName + "] not defined for method with " +
@@ -400,8 +398,8 @@ public class PainlessLookupBuilder {
                 "and parameters " + whitelistMethod.painlessParameterTypeNames);
         }
 
-        PainlessMethodKey painlessMethodKey =
+        String painlessMethodKey =
-            new PainlessMethodKey(whitelistMethod.javaMethodName, whitelistMethod.painlessParameterTypeNames.size());
+            buildPainlessMethodKey(whitelistMethod.javaMethodName, whitelistMethod.painlessParameterTypeNames.size());
 
         if (javaAugmentedClass == null && Modifier.isStatic(javaMethod.getModifiers())) {
             PainlessMethod painlessMethod = ownerStruct.staticMethods.get(painlessMethodKey);
@@ -459,7 +457,7 @@ public class PainlessLookupBuilder {
     }
 
     private void addField(String ownerStructName, WhitelistField whitelistField) {
-        PainlessClass ownerStruct = javaClassesToPainlessStructs.get(painlessTypesToJavaClasses.get(ownerStructName));
+        PainlessClassBuilder ownerStruct = javaClassesToPainlessClassBuilders.get(painlessTypesToJavaClasses.get(ownerStructName));
 
         if (ownerStruct == null) {
             throw new IllegalArgumentException("owner struct [" + ownerStructName + "] not defined for method with " +
@@ -540,14 +538,14 @@ public class PainlessLookupBuilder {
     }
 
     private void copyStruct(String struct, List<String> children) {
-        final PainlessClass owner = javaClassesToPainlessStructs.get(painlessTypesToJavaClasses.get(struct));
+        final PainlessClassBuilder owner = javaClassesToPainlessClassBuilders.get(painlessTypesToJavaClasses.get(struct));
 
         if (owner == null) {
             throw new IllegalArgumentException("Owner struct [" + struct + "] not defined for copy.");
         }
 
         for (int count = 0; count < children.size(); ++count) {
-            final PainlessClass child = javaClassesToPainlessStructs.get(painlessTypesToJavaClasses.get(children.get(count)));
+            final PainlessClassBuilder child = javaClassesToPainlessClassBuilders.get(painlessTypesToJavaClasses.get(children.get(count)));
 
             if (child == null) {
                 throw new IllegalArgumentException("Child struct [" + children.get(count) + "]" +
@@ -559,8 +557,8 @@ public class PainlessLookupBuilder {
                     " is not a super type of owner struct [" + owner.name + "] in copy.");
             }
 
-            for (Map.Entry<PainlessMethodKey,PainlessMethod> kvPair : child.methods.entrySet()) {
+            for (Map.Entry<String,PainlessMethod> kvPair : child.methods.entrySet()) {
-                PainlessMethodKey methodKey = kvPair.getKey();
+                String methodKey = kvPair.getKey();
                 PainlessMethod method = kvPair.getValue();
                 if (owner.methods.get(methodKey) == null) {
                     // TODO: some of these are no longer valid or outright don't work
@@ -625,10 +623,10 @@ public class PainlessLookupBuilder {
     /**
      * Precomputes a more efficient structure for dynamic method/field access.
      */
-    private void addRuntimeClass(final PainlessClass struct) {
+    private void addRuntimeClass(final PainlessClassBuilder struct) {
         // add all getters/setters
-        for (Map.Entry<PainlessMethodKey, PainlessMethod> method : struct.methods.entrySet()) {
+        for (Map.Entry<String, PainlessMethod> method : struct.methods.entrySet()) {
-            String name = method.getKey().name;
+            String name = method.getValue().name;
             PainlessMethod m = method.getValue();
 
             if (m.arguments.size() == 0 &&
@@ -668,7 +666,7 @@ public class PainlessLookupBuilder {
     }
 
     /** computes the functional interface method for a class, or returns null */
-    private PainlessMethod computeFunctionalInterfaceMethod(PainlessClass clazz) {
+    private PainlessMethod computeFunctionalInterfaceMethod(PainlessClassBuilder clazz) {
         if (!clazz.clazz.isInterface()) {
             return null;
         }
@@ -703,7 +701,7 @@ public class PainlessLookupBuilder {
         }
         // inspect the one method found from the reflection API, it should match the whitelist!
         java.lang.reflect.Method oneMethod = methods.get(0);
-        PainlessMethod painless = clazz.methods.get(new PainlessMethodKey(oneMethod.getName(), oneMethod.getParameterCount()));
+        PainlessMethod painless = clazz.methods.get(buildPainlessMethodKey(oneMethod.getName(), oneMethod.getParameterCount()));
         if (painless == null || painless.method.equals(org.objectweb.asm.commons.Method.getMethod(oneMethod)) == false) {
             throw new IllegalArgumentException("Class: " + clazz.name + " is functional but the functional " +
                 "method is not whitelisted!");
@@ -712,7 +710,15 @@ public class PainlessLookupBuilder {
     }
 
     public PainlessLookup build() {
-        return new PainlessLookup(painlessTypesToJavaClasses, javaClassesToPainlessStructs);
+        Map<Class<?>, PainlessClass> javaClassesToPainlessClasses = new HashMap<>();
+
+        // copy all structs to make them unmodifiable for outside users:
+        for (Map.Entry<Class<?>,PainlessClassBuilder> entry : javaClassesToPainlessClassBuilders.entrySet()) {
+            entry.getValue().functionalMethod = computeFunctionalInterfaceMethod(entry.getValue());
+            javaClassesToPainlessClasses.put(entry.getKey(), entry.getValue().build());
+        }
+
+        return new PainlessLookup(painlessTypesToJavaClasses, javaClassesToPainlessClasses);
     }
 
     public Class<?> getJavaClassFromPainlessType(String painlessType) {

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

@@ -1,75 +0,0 @@
-/*
- * Licensed to Elasticsearch under one or more contributor
- * license agreements. See the NOTICE file distributed with
- * this work for additional information regarding copyright
- * ownership. Elasticsearch licenses this file to you under
- * the Apache License, Version 2.0 (the "License"); you may
- * not use this file except in compliance with the License.
- * You may obtain a copy of the License at
- *
- *    http://www.apache.org/licenses/LICENSE-2.0
- *
- * Unless required by applicable law or agreed to in writing,
- * software distributed under the License is distributed on an
- * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
- * KIND, either express or implied.  See the License for the
- * specific language governing permissions and limitations
- * under the License.
- */
-
-package org.elasticsearch.painless.lookup;
-
-import java.util.Objects;
-
-/**
- * Key for looking up a method.
- * <p>
- * Methods are keyed on both name and arity, and can be overloaded once per arity.
- * This allows signatures such as {@code String.indexOf(String) vs String.indexOf(String, int)}.
- * <p>
- * It is less flexible than full signature overloading where types can differ too, but
- * better than just the name, and overloading types adds complexity to users, too.
- */
-public final class PainlessMethodKey {
-    public final String name;
-    public final int arity;
-
-    /**
-     * Create a new lookup key
-     * @param name name of the method
-     * @param arity number of parameters
-     */
-    public PainlessMethodKey(String name, int arity) {
-        this.name = Objects.requireNonNull(name);
-        this.arity = arity;
-    }
-
-    @Override
-    public int hashCode() {
-        final int prime = 31;
-        int result = 1;
-        result = prime * result + arity;
-        result = prime * result + name.hashCode();
-        return result;
-    }
-
-    @Override
-    public boolean equals(Object obj) {
-        if (this == obj) return true;
-        if (obj == null) return false;
-        if (getClass() != obj.getClass()) return false;
-        PainlessMethodKey other = (PainlessMethodKey) obj;
-        if (arity != other.arity) return false;
-        if (!name.equals(other.name)) return false;
-        return true;
-    }
-
-    @Override
-    public String toString() {
-        StringBuilder sb = new StringBuilder();
-        sb.append(name);
-        sb.append('/');
-        sb.append(arity);
-        return sb.toString();
-    }
-}

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

@@ -23,8 +23,8 @@ import org.elasticsearch.painless.Globals;
 import org.elasticsearch.painless.Locals;
 import org.elasticsearch.painless.Location;
 import org.elasticsearch.painless.MethodWriter;
+import org.elasticsearch.painless.lookup.PainlessLookupUtility;
 import org.elasticsearch.painless.lookup.PainlessMethod;
-import org.elasticsearch.painless.lookup.PainlessMethodKey;
 
 import java.util.List;
 import java.util.Objects;
@@ -58,7 +58,7 @@ public final class ECallLocal extends AExpression {
 
     @Override
     void analyze(Locals locals) {
-        PainlessMethodKey methodKey = new PainlessMethodKey(name, arguments.size());
+        String methodKey = PainlessLookupUtility.buildPainlessMethodKey(name, arguments.size());
         method = locals.getMethod(methodKey);
 
         if (method == null) {

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

@@ -27,7 +27,6 @@ import org.elasticsearch.painless.Location;
 import org.elasticsearch.painless.MethodWriter;
 import org.elasticsearch.painless.lookup.PainlessLookupUtility;
 import org.elasticsearch.painless.lookup.PainlessMethod;
-import org.elasticsearch.painless.lookup.PainlessMethodKey;
 import org.objectweb.asm.Type;
 
 import java.util.Objects;
@@ -71,7 +70,8 @@ public final class EFunctionRef extends AExpression implements ILambda {
                         throw new IllegalArgumentException("Cannot convert function reference [" + type + "::" + call + "] " +
                                 "to [" + PainlessLookupUtility.anyTypeToPainlessTypeName(expected) + "], not a functional interface");
                     }
-                    PainlessMethod delegateMethod = locals.getMethod(new PainlessMethodKey(call, interfaceMethod.arguments.size()));
+                    PainlessMethod delegateMethod =
+                            locals.getMethod(PainlessLookupUtility.buildPainlessMethodKey(call, interfaceMethod.arguments.size()));
                     if (delegateMethod == null) {
                         throw new IllegalArgumentException("Cannot convert function reference [" + type + "::" + call + "] " +
                                 "to [" + PainlessLookupUtility.anyTypeToPainlessTypeName(expected) + "], function not found");

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

@@ -23,8 +23,8 @@ import org.elasticsearch.painless.Globals;
 import org.elasticsearch.painless.Locals;
 import org.elasticsearch.painless.Location;
 import org.elasticsearch.painless.MethodWriter;
+import org.elasticsearch.painless.lookup.PainlessLookupUtility;
 import org.elasticsearch.painless.lookup.PainlessMethod;
-import org.elasticsearch.painless.lookup.PainlessMethodKey;
 import org.elasticsearch.painless.lookup.def;
 import org.objectweb.asm.Type;
 
@@ -62,14 +62,15 @@ public final class EListInit extends AExpression {
 
         actual = ArrayList.class;
 
-        constructor =
+        constructor = locals.getPainlessLookup().getPainlessStructFromJavaClass(actual).constructors
-                locals.getPainlessLookup().getPainlessStructFromJavaClass(actual).constructors.get(new PainlessMethodKey("<init>", 0));
+                .get(PainlessLookupUtility.buildPainlessMethodKey("<init>", 0));
 
         if (constructor == null) {
             throw createError(new IllegalStateException("Illegal tree structure."));
         }
 
-        method = locals.getPainlessLookup().getPainlessStructFromJavaClass(actual).methods.get(new PainlessMethodKey("add", 1));
+        method = locals.getPainlessLookup().getPainlessStructFromJavaClass(actual).methods
+                .get(PainlessLookupUtility.buildPainlessMethodKey("add", 1));
 
         if (method == null) {
             throw createError(new IllegalStateException("Illegal tree structure."));

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

@@ -23,8 +23,8 @@ import org.elasticsearch.painless.Globals;
 import org.elasticsearch.painless.Locals;
 import org.elasticsearch.painless.Location;
 import org.elasticsearch.painless.MethodWriter;
+import org.elasticsearch.painless.lookup.PainlessLookupUtility;
 import org.elasticsearch.painless.lookup.PainlessMethod;
-import org.elasticsearch.painless.lookup.PainlessMethodKey;
 import org.elasticsearch.painless.lookup.def;
 import org.objectweb.asm.Type;
 
@@ -68,14 +68,15 @@ public final class EMapInit extends AExpression {
 
         actual = HashMap.class;
 
-        constructor =
+        constructor = locals.getPainlessLookup().getPainlessStructFromJavaClass(actual).constructors
-                locals.getPainlessLookup().getPainlessStructFromJavaClass(actual).constructors.get(new PainlessMethodKey("<init>", 0));
+                .get(PainlessLookupUtility.buildPainlessMethodKey("<init>", 0));
 
         if (constructor == null) {
             throw createError(new IllegalStateException("Illegal tree structure."));
         }
 
-        method = locals.getPainlessLookup().getPainlessStructFromJavaClass(actual).methods.get(new PainlessMethodKey("put", 2));
+        method = locals.getPainlessLookup().getPainlessStructFromJavaClass(actual).methods
+                .get(PainlessLookupUtility.buildPainlessMethodKey("put", 2));
 
         if (method == null) {
             throw createError(new IllegalStateException("Illegal tree structure."));

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

@@ -24,8 +24,8 @@ import org.elasticsearch.painless.Locals;
 import org.elasticsearch.painless.Location;
 import org.elasticsearch.painless.MethodWriter;
 import org.elasticsearch.painless.lookup.PainlessClass;
+import org.elasticsearch.painless.lookup.PainlessLookupUtility;
 import org.elasticsearch.painless.lookup.PainlessMethod;
-import org.elasticsearch.painless.lookup.PainlessMethodKey;
 import org.objectweb.asm.Type;
 
 import java.util.List;
@@ -65,7 +65,7 @@ public final class ENewObj extends AExpression {
         }
 
         PainlessClass struct = locals.getPainlessLookup().getPainlessStructFromJavaClass(actual);
-        constructor = struct.constructors.get(new PainlessMethodKey("<init>", arguments.size()));
+        constructor = struct.constructors.get(PainlessLookupUtility.buildPainlessMethodKey("<init>", arguments.size()));
 
         if (constructor != null) {
             Class<?>[] types = new Class<?>[constructor.arguments.size()];

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

@@ -26,7 +26,6 @@ import org.elasticsearch.painless.MethodWriter;
 import org.elasticsearch.painless.lookup.PainlessClass;
 import org.elasticsearch.painless.lookup.PainlessLookupUtility;
 import org.elasticsearch.painless.lookup.PainlessMethod;
-import org.elasticsearch.painless.lookup.PainlessMethodKey;
 import org.elasticsearch.painless.lookup.def;
 
 import java.util.List;
@@ -77,7 +76,7 @@ public final class PCallInvoke extends AExpression {
             struct = locals.getPainlessLookup().getPainlessStructFromJavaClass(PainlessLookupUtility.getBoxedAnyType(prefix.actual));
         }
 
-        PainlessMethodKey methodKey = new PainlessMethodKey(name, arguments.size());
+        String methodKey = PainlessLookupUtility.buildPainlessMethodKey(name, arguments.size());
         PainlessMethod method = prefix instanceof EStatic ? struct.staticMethods.get(methodKey) : struct.methods.get(methodKey);
 
         if (method != null) {

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

@@ -27,7 +27,6 @@ import org.elasticsearch.painless.lookup.PainlessClass;
 import org.elasticsearch.painless.lookup.PainlessField;
 import org.elasticsearch.painless.lookup.PainlessLookupUtility;
 import org.elasticsearch.painless.lookup.PainlessMethod;
-import org.elasticsearch.painless.lookup.PainlessMethodKey;
 import org.elasticsearch.painless.lookup.def;
 
 import java.util.List;
@@ -74,16 +73,16 @@ public final class PField extends AStoreable {
             if (field != null) {
                 sub = new PSubField(location, field);
             } else {
-                PainlessMethod getter = struct.methods.get(
+                PainlessMethod getter = struct.methods.get(PainlessLookupUtility.buildPainlessMethodKey(
-                    new PainlessMethodKey("get" + Character.toUpperCase(value.charAt(0)) + value.substring(1), 0));
+                        "get" + Character.toUpperCase(value.charAt(0)) + value.substring(1), 0));
 
                 if (getter == null) {
-                    getter = struct.methods.get(
+                    getter = struct.methods.get(PainlessLookupUtility.buildPainlessMethodKey(
-                        new PainlessMethodKey("is" + Character.toUpperCase(value.charAt(0)) + value.substring(1), 0));
+                            "is" + Character.toUpperCase(value.charAt(0)) + value.substring(1), 0));
                 }
 
-                PainlessMethod setter = struct.methods.get(
+                PainlessMethod setter = struct.methods.get(PainlessLookupUtility.buildPainlessMethodKey(
-                    new PainlessMethodKey("set" + Character.toUpperCase(value.charAt(0)) + value.substring(1), 1));
+                        "set" + Character.toUpperCase(value.charAt(0)) + value.substring(1), 1));
 
                 if (getter != null || setter != null) {
                     sub = new PSubShortcut(location, value, PainlessLookupUtility.anyTypeToPainlessTypeName(prefix.actual), getter, setter);

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

@@ -25,8 +25,8 @@ import org.elasticsearch.painless.Location;
 import org.elasticsearch.painless.MethodWriter;
 import org.elasticsearch.painless.WriterConstants;
 import org.elasticsearch.painless.lookup.PainlessClass;
+import org.elasticsearch.painless.lookup.PainlessLookupUtility;
 import org.elasticsearch.painless.lookup.PainlessMethod;
-import org.elasticsearch.painless.lookup.PainlessMethodKey;
 
 import java.util.Objects;
 import java.util.Set;
@@ -56,8 +56,8 @@ final class PSubListShortcut extends AStoreable {
 
     @Override
     void analyze(Locals locals) {
-        getter = struct.methods.get(new PainlessMethodKey("get", 1));
+        getter = struct.methods.get(PainlessLookupUtility.buildPainlessMethodKey("get", 1));
-        setter = struct.methods.get(new PainlessMethodKey("set", 2));
+        setter = struct.methods.get(PainlessLookupUtility.buildPainlessMethodKey("set", 2));
 
         if (getter != null && (getter.rtn == void.class || getter.arguments.size() != 1 ||
             getter.arguments.get(0) != int.class)) {

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

@@ -24,8 +24,8 @@ import org.elasticsearch.painless.Locals;
 import org.elasticsearch.painless.Location;
 import org.elasticsearch.painless.MethodWriter;
 import org.elasticsearch.painless.lookup.PainlessClass;
+import org.elasticsearch.painless.lookup.PainlessLookupUtility;
 import org.elasticsearch.painless.lookup.PainlessMethod;
-import org.elasticsearch.painless.lookup.PainlessMethodKey;
 
 import java.util.Objects;
 import java.util.Set;
@@ -55,8 +55,8 @@ final class PSubMapShortcut extends AStoreable {
 
     @Override
     void analyze(Locals locals) {
-        getter = struct.methods.get(new PainlessMethodKey("get", 1));
+        getter = struct.methods.get(PainlessLookupUtility.buildPainlessMethodKey("get", 1));
-        setter = struct.methods.get(new PainlessMethodKey("put", 2));
+        setter = struct.methods.get(PainlessLookupUtility.buildPainlessMethodKey("put", 2));
 
         if (getter != null && (getter.rtn == void.class || getter.arguments.size() != 1)) {
             throw createError(new IllegalArgumentException("Illegal map get shortcut for type [" + struct.name + "]."));

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

@@ -30,8 +30,8 @@ import org.elasticsearch.painless.ScriptClassInfo;
 import org.elasticsearch.painless.SimpleChecksAdapter;
 import org.elasticsearch.painless.WriterConstants;
 import org.elasticsearch.painless.lookup.PainlessLookup;
+import org.elasticsearch.painless.lookup.PainlessLookupUtility;
 import org.elasticsearch.painless.lookup.PainlessMethod;
-import org.elasticsearch.painless.lookup.PainlessMethodKey;
 import org.elasticsearch.painless.node.SFunction.FunctionReserved;
 import org.objectweb.asm.ClassVisitor;
 import org.objectweb.asm.ClassWriter;
@@ -165,12 +165,12 @@ public final class SSource extends AStatement {
     }
 
     public void analyze(PainlessLookup painlessLookup) {
-        Map<PainlessMethodKey, PainlessMethod> methods = new HashMap<>();
+        Map<String, PainlessMethod> methods = new HashMap<>();
 
         for (SFunction function : functions) {
             function.generateSignature(painlessLookup);
 
-            PainlessMethodKey key = new PainlessMethodKey(function.name, function.parameters.size());
+            String key = PainlessLookupUtility.buildPainlessMethodKey(function.name, function.parameters.size());
 
             if (methods.put(key, function.method) != null) {
                 throw createError(new IllegalArgumentException("Duplicate functions with name [" + function.name + "]."));

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

@@ -29,7 +29,6 @@ import org.elasticsearch.painless.MethodWriter;
 import org.elasticsearch.painless.lookup.PainlessCast;
 import org.elasticsearch.painless.lookup.PainlessLookupUtility;
 import org.elasticsearch.painless.lookup.PainlessMethod;
-import org.elasticsearch.painless.lookup.PainlessMethodKey;
 import org.elasticsearch.painless.lookup.def;
 import org.objectweb.asm.Label;
 import org.objectweb.asm.Opcodes;
@@ -77,8 +76,8 @@ final class SSubEachIterable extends AStatement {
         if (expression.actual == def.class) {
             method = null;
         } else {
-            method = locals.getPainlessLookup().
+            method = locals.getPainlessLookup().getPainlessStructFromJavaClass(expression.actual).methods
-                    getPainlessStructFromJavaClass(expression.actual).methods.get(new PainlessMethodKey("iterator", 0));
+                    .get(PainlessLookupUtility.buildPainlessMethodKey("iterator", 0));
 
             if (method == null) {
                 throw createError(new IllegalArgumentException("Unable to create iterator for the type " +

modules/lang-painless/src/test/java/org/elasticsearch/painless/node/NodeToStringTests.java

@@ -24,8 +24,8 @@ import org.elasticsearch.painless.lookup.PainlessLookup;
 import org.elasticsearch.painless.lookup.PainlessCast;
 import org.elasticsearch.painless.lookup.PainlessField;
 import org.elasticsearch.painless.lookup.PainlessLookupBuilder;
+import org.elasticsearch.painless.lookup.PainlessLookupUtility;
 import org.elasticsearch.painless.lookup.PainlessMethod;
-import org.elasticsearch.painless.lookup.PainlessMethodKey;
 import org.elasticsearch.painless.lookup.PainlessClass;
 import org.elasticsearch.painless.FeatureTest;
 import org.elasticsearch.painless.GenericElasticsearchScript;
@@ -405,14 +405,14 @@ public class NodeToStringTests extends ESTestCase {
     public void testPSubCallInvoke() {
         Location l = new Location(getTestName(), 0);
         PainlessClass c = painlessLookup.getPainlessStructFromJavaClass(Integer.class);
-        PainlessMethod m = c.methods.get(new PainlessMethodKey("toString", 0));
+        PainlessMethod m = c.methods.get(PainlessLookupUtility.buildPainlessMethodKey("toString", 0));
         PSubCallInvoke node = new PSubCallInvoke(l, m, null, emptyList());
         node.prefix = new EVariable(l, "a");
         assertEquals("(PSubCallInvoke (EVariable a) toString)", node.toString());
         assertEquals("(PSubNullSafeCallInvoke (PSubCallInvoke (EVariable a) toString))", new PSubNullSafeCallInvoke(l, node).toString());
 
         l = new Location(getTestName(), 1);
-        m = c.methods.get(new PainlessMethodKey("equals", 1));
+        m = c.methods.get(PainlessLookupUtility.buildPainlessMethodKey("equals", 1));
         node = new PSubCallInvoke(l, m, null, singletonList(new EVariable(l, "b")));
         node.prefix = new EVariable(l, "a");
         assertEquals("(PSubCallInvoke (EVariable a) equals (Args (EVariable b)))", node.toString());
@@ -502,8 +502,8 @@ public class NodeToStringTests extends ESTestCase {
     public void testPSubShortcut() {
         Location l = new Location(getTestName(), 0);
         PainlessClass s = painlessLookup.getPainlessStructFromJavaClass(FeatureTest.class);
-        PainlessMethod getter = s.methods.get(new PainlessMethodKey("getX", 0));
+        PainlessMethod getter = s.methods.get(PainlessLookupUtility.buildPainlessMethodKey("getX", 0));
-        PainlessMethod setter = s.methods.get(new PainlessMethodKey("setX", 1));
+        PainlessMethod setter = s.methods.get(PainlessLookupUtility.buildPainlessMethodKey("setX", 1));
         PSubShortcut node = new PSubShortcut(l, "x", FeatureTest.class.getName(), getter, setter);
         node.prefix = new EVariable(l, "a");
         assertEquals("(PSubShortcut (EVariable a) x)", node.toString());