[Painless] Add instance bindings (#34410)

This change adds instance bindings to Painless. This binding allows a whitelisted
 method to be called on an instance instantiated prior to script compilation.
Whitelisting must be done in code as there is no practical way to instantiate a
useful instance from a text file (see the tests for an example). Since an
instance can be shared by multiple scripts, each method called must be
thread-safe.

modules/lang-painless/spi/src/main/java/org/elasticsearch/painless/spi/Whitelist.java

@@ -66,13 +66,17 @@ public final class Whitelist {
     /** The {@link List} of all the whitelisted Painless class bindings. */
     public final List<WhitelistClassBinding> whitelistClassBindings;
 
+    /** The {@link List} of all the whitelisted Painless instance bindings. */
+    public final List<WhitelistInstanceBinding> whitelistInstanceBindings;
+
     /** Standard constructor. All values must be not {@code null}. */
-    public Whitelist(ClassLoader classLoader, List<WhitelistClass> whitelistClasses,
-            List<WhitelistMethod> whitelistImportedMethods, List<WhitelistClassBinding> whitelistClassBindings) {
+    public Whitelist(ClassLoader classLoader, List<WhitelistClass> whitelistClasses, List<WhitelistMethod> whitelistImportedMethods,
+            List<WhitelistClassBinding> whitelistClassBindings, List<WhitelistInstanceBinding> whitelistInstanceBindings) {
 
         this.classLoader = Objects.requireNonNull(classLoader);
         this.whitelistClasses = Collections.unmodifiableList(Objects.requireNonNull(whitelistClasses));
         this.whitelistImportedMethods = Collections.unmodifiableList(Objects.requireNonNull(whitelistImportedMethods));
         this.whitelistClassBindings = Collections.unmodifiableList(Objects.requireNonNull(whitelistClassBindings));
+        this.whitelistInstanceBindings = Collections.unmodifiableList(Objects.requireNonNull(whitelistInstanceBindings));
     }
 }

modules/lang-painless/spi/src/main/java/org/elasticsearch/painless/spi/WhitelistClassBinding.java

@@ -42,9 +42,7 @@ public class WhitelistClassBinding {
     /** The method name for this class binding. */
     public final String methodName;
 
-    /**
-     * The canonical type name for the return type.
-     */
+    /** The canonical type name for the return type. */
     public final String returnCanonicalTypeName;
 
     /**

modules/lang-painless/spi/src/main/java/org/elasticsearch/painless/spi/WhitelistInstanceBinding.java

@@ -0,0 +1,61 @@
+/*
+ * 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.spi;
+
+import java.util.List;
+import java.util.Objects;
+
+/**
+ * An instance binding represents a method call that stores state. Each instance binding must provide
+ * exactly one public method name. The canonical type name parameters provided must match those of the
+ * method. The method for an instance binding will target the specified Java instance.
+ */
+public class WhitelistInstanceBinding {
+
+    /** Information about where this constructor was whitelisted from. */
+    public final String origin;
+
+    /** The Java instance this instance binding targets. */
+    public final Object targetInstance;
+
+    /** The method name for this class binding. */
+    public final String methodName;
+
+    /** The canonical type name for the return type. */
+    public final String returnCanonicalTypeName;
+
+    /**
+     * A {@link List} of {@link String}s that are the Painless type names for the parameters of the
+     * constructor which can be used to look up the Java constructor through reflection.
+     */
+    public final List<String> canonicalTypeNameParameters;
+
+    /** Standard constructor. All values must be not {@code null}. */
+    public WhitelistInstanceBinding(String origin, Object targetInstance,
+            String methodName, String returnCanonicalTypeName, List<String> canonicalTypeNameParameters) {
+
+        this.origin = Objects.requireNonNull(origin);
+        this.targetInstance = Objects.requireNonNull(targetInstance);
+
+        this.methodName = Objects.requireNonNull(methodName);
+        this.returnCanonicalTypeName = Objects.requireNonNull(returnCanonicalTypeName);
+        this.canonicalTypeNameParameters = Objects.requireNonNull(canonicalTypeNameParameters);
+    }
+}

modules/lang-painless/spi/src/main/java/org/elasticsearch/painless/spi/WhitelistLoader.java

@@ -29,6 +29,7 @@
 import java.security.PrivilegedAction;
 import java.util.ArrayList;
 import java.util.Arrays;
+import java.util.Collections;
 import java.util.List;
 
 /** Loads and creates a {@link Whitelist} from one to many text files. */
@@ -392,7 +393,7 @@ public static Whitelist loadFromResourceFiles(Class<?> resource, String... filep
 
         ClassLoader loader = AccessController.doPrivileged((PrivilegedAction<ClassLoader>)resource::getClassLoader);
 
-        return new Whitelist(loader, whitelistClasses, whitelistStatics, whitelistClassBindings);
+        return new Whitelist(loader, whitelistClasses, whitelistStatics, whitelistClassBindings, Collections.emptyList());
     }
 
     private WhitelistLoader() {}

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

@@ -20,7 +20,6 @@
 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;
@@ -222,16 +221,19 @@ 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);
-        Map<String, LocalMethod> localMethods = root.analyze(painlessLookup);
-        root.write();
+        root.analyze(painlessLookup);
+        Map<String, Object> statics = root.write();
 
         try {
             Class<? extends PainlessScript> clazz = loader.defineScript(CLASS_NAME, root.getBytes());
             clazz.getField("$NAME").set(null, 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);
+
+            for (Map.Entry<String, Object> statik : statics.entrySet()) {
+                clazz.getField(statik.getKey()).set(null, statik.getValue());
+            }
 
             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/Globals.java

@@ -31,7 +31,8 @@
 public class Globals {
     private final Map<String,SFunction> syntheticMethods = new HashMap<>();
     private final Map<String,Constant> constantInitializers = new HashMap<>();
-    private final Map<String,Class<?>> bindings = new HashMap<>();
+    private final Map<String,Class<?>> classBindings = new HashMap<>();
+    private final Map<Object,String> instanceBindings = new HashMap<>();
     private final BitSet statements;
 
     /** Create a new Globals from the set of statement boundaries */
@@ -56,14 +57,19 @@ public void addConstantInitializer(Constant constant) {
         }
     }
 
-    /** Adds a new binding to be written as a local variable */
-    public String addBinding(Class<?> type) {
-        String name = "$binding$" + bindings.size();
-        bindings.put(name, type);
+    /** Adds a new class binding to be written as a local variable */
+    public String addClassBinding(Class<?> type) {
+        String name = "$class_binding$" + classBindings.size();
+        classBindings.put(name, type);
 
         return name;
     }
 
+    /** Adds a new binding to be written as a local variable */
+    public String addInstanceBinding(Object instance) {
+        return instanceBindings.computeIfAbsent(instance, key -> "$instance_binding$" + instanceBindings.size());
+    }
+
     /** Returns the current synthetic methods */
     public Map<String,SFunction> getSyntheticMethods() {
         return syntheticMethods;
@@ -75,8 +81,13 @@ public Map<String,Constant> getConstantInitializers() {
     }
 
     /** Returns the current bindings */
-    public Map<String,Class<?>> getBindings() {
-        return bindings;
+    public Map<String,Class<?>> getClassBindings() {
+        return classBindings;
+    }
+
+    /** Returns the current bindings */
+    public Map<Object,String> getInstanceBindings() {
+        return instanceBindings;
     }
 
     /** Returns the set of statement boundaries */

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

@@ -60,7 +60,6 @@ public boolean equals(Object object) {
 
     @Override
     public int hashCode() {
-
         return Objects.hash(javaConstructor, javaMethod, returnType, typeParameters);
     }
 }

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

@@ -0,0 +1,64 @@
+/*
+ * 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.lang.reflect.Method;
+import java.util.List;
+import java.util.Objects;
+
+public class PainlessInstanceBinding {
+
+    public final Object targetInstance;
+    public final Method javaMethod;
+
+    public final Class<?> returnType;
+    public final List<Class<?>> typeParameters;
+
+    PainlessInstanceBinding(Object targetInstance, Method javaMethod, Class<?> returnType, List<Class<?>> typeParameters) {
+        this.targetInstance = targetInstance;
+        this.javaMethod = javaMethod;
+
+        this.returnType = returnType;
+        this.typeParameters = typeParameters;
+    }
+
+    @Override
+    public boolean equals(Object object) {
+        if (this == object) {
+            return true;
+        }
+
+        if (object == null || getClass() != object.getClass()) {
+            return false;
+        }
+
+        PainlessInstanceBinding that = (PainlessInstanceBinding)object;
+
+        return targetInstance == that.targetInstance &&
+                Objects.equals(javaMethod, that.javaMethod) &&
+                Objects.equals(returnType, that.returnType) &&
+                Objects.equals(typeParameters, that.typeParameters);
+    }
+
+    @Override
+    public int hashCode() {
+        return Objects.hash(targetInstance, javaMethod, returnType, typeParameters);
+    }
+}

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

@@ -40,27 +40,31 @@ public final class PainlessLookup {
 
     private final Map<String, PainlessMethod> painlessMethodKeysToImportedPainlessMethods;
     private final Map<String, PainlessClassBinding> painlessMethodKeysToPainlessClassBindings;
+    private final Map<String, PainlessInstanceBinding> painlessMethodKeysToPainlessInstanceBindings;
 
     PainlessLookup(
             Map<String, Class<?>> javaClassNamesToClasses,
             Map<String, Class<?>> canonicalClassNamesToClasses,
             Map<Class<?>, PainlessClass> classesToPainlessClasses,
             Map<String, PainlessMethod> painlessMethodKeysToImportedPainlessMethods,
-            Map<String, PainlessClassBinding> painlessMethodKeysToPainlessClassBindings) {
+            Map<String, PainlessClassBinding> painlessMethodKeysToPainlessClassBindings,
+            Map<String, PainlessInstanceBinding> painlessMethodKeysToPainlessInstanceBindings) {
 
         Objects.requireNonNull(javaClassNamesToClasses);
         Objects.requireNonNull(canonicalClassNamesToClasses);
         Objects.requireNonNull(classesToPainlessClasses);
 
         Objects.requireNonNull(painlessMethodKeysToImportedPainlessMethods);
         Objects.requireNonNull(painlessMethodKeysToPainlessClassBindings);
+        Objects.requireNonNull(painlessMethodKeysToPainlessInstanceBindings);
 
         this.javaClassNamesToClasses = javaClassNamesToClasses;
         this.canonicalClassNamesToClasses = Collections.unmodifiableMap(canonicalClassNamesToClasses);
         this.classesToPainlessClasses = Collections.unmodifiableMap(classesToPainlessClasses);
 
         this.painlessMethodKeysToImportedPainlessMethods = Collections.unmodifiableMap(painlessMethodKeysToImportedPainlessMethods);
         this.painlessMethodKeysToPainlessClassBindings = Collections.unmodifiableMap(painlessMethodKeysToPainlessClassBindings);
+        this.painlessMethodKeysToPainlessInstanceBindings = Collections.unmodifiableMap(painlessMethodKeysToPainlessInstanceBindings);
     }
 
     public Class<?> javaClassNameToClass(String javaClassName) {
@@ -200,6 +204,14 @@ public PainlessClassBinding lookupPainlessClassBinding(String methodName, int ar
         return painlessMethodKeysToPainlessClassBindings.get(painlessMethodKey);
     }
 
+    public PainlessInstanceBinding lookupPainlessInstanceBinding(String methodName, int arity) {
+        Objects.requireNonNull(methodName);
+
+        String painlessMethodKey = buildPainlessMethodKey(methodName, arity);
+
+        return painlessMethodKeysToPainlessInstanceBindings.get(painlessMethodKey);
+    }
+
     public PainlessMethod lookupFunctionalInterfacePainlessMethod(Class<?> targetClass) {
         PainlessClass targetPainlessClass = classesToPainlessClasses.get(targetClass);