Painless: Add Imported Static Method (#33440)

Allow static methods to be imported in Painless and called using just the method name.

Author: jdconrad

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

@@ -61,12 +61,19 @@ public final class Whitelist {
     /** The {@link List} of all the whitelisted Painless classes. */
     public final List<WhitelistClass> whitelistClasses;
 
+    /** The {@link List} of all the whitelisted static Painless methods. */
+    public final List<WhitelistMethod> whitelistImportedMethods;
+
+    /** The {@link List} of all the whitelisted Painless bindings. */
     public final List<WhitelistBinding> whitelistBindings;
 
-    /** Standard constructor.  All values must be not {@code null}. */
-    public Whitelist(ClassLoader classLoader, List<WhitelistClass> whitelistClasses, List<WhitelistBinding> whitelistBindings) {
+    /** Standard constructor. All values must be not {@code null}. */
+    public Whitelist(ClassLoader classLoader, List<WhitelistClass> whitelistClasses,
+            List<WhitelistMethod> whitelistImportedMethods, List<WhitelistBinding> whitelistBindings) {
+
         this.classLoader = Objects.requireNonNull(classLoader);
         this.whitelistClasses = Collections.unmodifiableList(Objects.requireNonNull(whitelistClasses));
+        this.whitelistImportedMethods = Collections.unmodifiableList(Objects.requireNonNull(whitelistImportedMethods));
         this.whitelistBindings = Collections.unmodifiableList(Objects.requireNonNull(whitelistBindings));
     }
 }

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

@@ -133,6 +133,7 @@ public final class WhitelistLoader {
      */
     public static Whitelist loadFromResourceFiles(Class<?> resource, String... filepaths) {
         List<WhitelistClass> whitelistClasses = new ArrayList<>();
+        List<WhitelistMethod> whitelistStatics = new ArrayList<>();
         List<WhitelistBinding> whitelistBindings = new ArrayList<>();
 
         // Execute a single pass through the whitelist text files.  This will gather all the
@@ -192,18 +193,18 @@ public static Whitelist loadFromResourceFiles(Class<?> resource, String... filep
                         whitelistConstructors = new ArrayList<>();
                         whitelistMethods = new ArrayList<>();
                         whitelistFields = new ArrayList<>();
-                    } else if (line.startsWith("static ")) {
+                    } else if (line.startsWith("static_import ")) {
                         // Ensure the final token of the line is '{'.
                         if (line.endsWith("{") == false) {
                             throw new IllegalArgumentException(
-                                    "invalid static definition: failed to parse static opening bracket [" + line + "]");
+                                    "invalid static import definition: failed to parse static import opening bracket [" + line + "]");
                         }
 
                         if (parseType != null) {
-                            throw new IllegalArgumentException("invalid definition: cannot embed static definition [" + line + "]");
+                            throw new IllegalArgumentException("invalid definition: cannot embed static import definition [" + line + "]");
                         }
 
-                        parseType = "static";
+                        parseType = "static_import";
 
                     // Handle the end of a definition and reset all previously gathered values.
                     // Expects the following format: '}' '\n'
@@ -229,9 +230,9 @@ public static Whitelist loadFromResourceFiles(Class<?> resource, String... filep
                         // Reset the parseType.
                         parseType = null;
 
-                    // Handle static definition types.
-                    // Expects the following format: ID ID '(' ( ID ( ',' ID )* )? ')' 'bound_to' ID '\n'
-                    } else if ("static".equals(parseType)) {
+                    // Handle static import definition types.
+                    // Expects the following format: ID ID '(' ( ID ( ',' ID )* )? ')' ( 'from_class' | 'bound_to' ) ID '\n'
+                    } else if ("static_import".equals(parseType)) {
                         // Mark the origin of this parsable object.
                         String origin = "[" + filepath + "]:[" + number + "]";
 
@@ -240,7 +241,7 @@ public static Whitelist loadFromResourceFiles(Class<?> resource, String... filep
 
                         if (parameterStartIndex == -1) {
                             throw new IllegalArgumentException(
-                                    "illegal static definition: start of method parameters not found [" + line + "]");
+                                    "illegal static import definition: start of method parameters not found [" + line + "]");
                         }
 
                         String[] tokens = line.substring(0, parameterStartIndex).trim().split("\\s+");
@@ -261,7 +262,7 @@ public static Whitelist loadFromResourceFiles(Class<?> resource, String... filep
 
                         if (parameterEndIndex == -1) {
                             throw new IllegalArgumentException(
-                                    "illegal static definition: end of method parameters not found [" + line + "]");
+                                    "illegal static import definition: end of method parameters not found [" + line + "]");
                         }
 
                         String[] canonicalTypeNameParameters =
@@ -272,39 +273,37 @@ public static Whitelist loadFromResourceFiles(Class<?> resource, String... filep
                             canonicalTypeNameParameters = new String[0];
                         }
 
-                        // Parse the static type and class.
+                        // Parse the static import type and class.
                         tokens = line.substring(parameterEndIndex + 1).trim().split("\\s+");
 
-                        String staticType;
+                        String staticImportType;
                         String targetJavaClassName;
 
                         // Based on the number of tokens, look up the type and class.
                         if (tokens.length == 2) {
-                            staticType = tokens[0];
+                            staticImportType = tokens[0];
                             targetJavaClassName = tokens[1];
                         } else {
-                            throw new IllegalArgumentException("invalid static definition: unexpected format [" + line + "]");
+                            throw new IllegalArgumentException("invalid static import definition: unexpected format [" + line + "]");
                         }
 
-                        // Check the static type is valid.
-                        if ("bound_to".equals(staticType) == false) {
-                            throw new IllegalArgumentException(
-                                    "invalid static definition: unexpected static type [" + staticType + "] [" + line + "]");
+                        // Add a static import method or binding depending on the static import type.
+                        if ("from_class".equals(staticImportType)) {
+                            whitelistStatics.add(new WhitelistMethod(origin, targetJavaClassName,
+                                    methodName, returnCanonicalTypeName, Arrays.asList(canonicalTypeNameParameters)));
+                        } else if ("bound_to".equals(staticImportType)) {
+                            whitelistBindings.add(new WhitelistBinding(origin, targetJavaClassName,
+                                    methodName, returnCanonicalTypeName, Arrays.asList(canonicalTypeNameParameters)));
+                        } else {
+                            throw new IllegalArgumentException("invalid static import definition: " +
+                                    "unexpected static import type [" + staticImportType + "] [" + line + "]");
                         }
 
-                        whitelistBindings.add(new WhitelistBinding(origin, targetJavaClassName,
-                                methodName, returnCanonicalTypeName, Arrays.asList(canonicalTypeNameParameters)));
-
                     // Handle class definition types.
                     } else if ("class".equals(parseType)) {
                         // Mark the origin of this parsable object.
                         String origin = "[" + filepath + "]:[" + number + "]";
 
-                        // Ensure we have a defined class before adding any constructors, methods, augmented methods, or fields.
-                        if (parseType == null) {
-                            throw new IllegalArgumentException("invalid definition: expected one of ['class', 'static'] [" + line + "]");
-                        }
-
                         // Handle the case for a constructor definition.
                         // Expects the following format: '(' ( ID ( ',' ID )* )? ')' '\n'
                         if (line.startsWith("(")) {
@@ -393,7 +392,7 @@ public static Whitelist loadFromResourceFiles(Class<?> resource, String... filep
 
         ClassLoader loader = AccessController.doPrivileged((PrivilegedAction<ClassLoader>)resource::getClassLoader);
 
-        return new Whitelist(loader, whitelistClasses, whitelistBindings);
+        return new Whitelist(loader, whitelistClasses, whitelistStatics, whitelistBindings);
     }
 
     private WhitelistLoader() {}

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

@@ -24,6 +24,21 @@
 
 /** Currently just a dummy class for testing a few features not yet exposed by whitelist! */
 public class FeatureTest {
+    /** static method that returns true */
+    public static boolean overloadedStatic() {
+        return true;
+    }
+
+    /** static method that returns what you ask it */
+    public static boolean overloadedStatic(boolean whatToReturn) {
+        return whatToReturn;
+    }
+
+    /** static method only whitelisted as a static */
+    public static float staticAddFloatsTest(float x, float y) {
+        return x + y;
+    }
+
     private int x;
     private int y;
     public int z;
@@ -58,21 +73,12 @@ public void setY(int y) {
         this.y = y;
     }
 
-    /** static method that returns true */
-    public static boolean overloadedStatic() {
-        return true;
-    }
-
-    /** static method that returns what you ask it */
-    public static boolean overloadedStatic(boolean whatToReturn) {
-        return whatToReturn;
-    }
-
     /** method taking two functions! */
     public Object twoFunctionsOfX(Function<Object,Object> f, Function<Object,Object> g) {
         return f.apply(g.apply(x));
     }
 
+    /** method to take in a list */
     public void listInput(List<Object> list) {
 
     }

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

@@ -37,16 +37,23 @@ public final class PainlessLookup {
     private final Map<String, Class<?>> canonicalClassNamesToClasses;
     private final Map<Class<?>, PainlessClass> classesToPainlessClasses;
 
+    private final Map<String, PainlessMethod> painlessMethodKeysToImportedPainlessMethods;
     private final Map<String, PainlessBinding> painlessMethodKeysToPainlessBindings;
 
     PainlessLookup(Map<String, Class<?>> canonicalClassNamesToClasses, Map<Class<?>, PainlessClass> classesToPainlessClasses,
+            Map<String, PainlessMethod> painlessMethodKeysToImportedPainlessMethods,
             Map<String, PainlessBinding> painlessMethodKeysToPainlessBindings) {
+
         Objects.requireNonNull(canonicalClassNamesToClasses);
         Objects.requireNonNull(classesToPainlessClasses);
 
+        Objects.requireNonNull(painlessMethodKeysToImportedPainlessMethods);
+        Objects.requireNonNull(painlessMethodKeysToPainlessBindings);
+
         this.canonicalClassNamesToClasses = Collections.unmodifiableMap(canonicalClassNamesToClasses);
         this.classesToPainlessClasses = Collections.unmodifiableMap(classesToPainlessClasses);
 
+        this.painlessMethodKeysToImportedPainlessMethods = Collections.unmodifiableMap(painlessMethodKeysToImportedPainlessMethods);
         this.painlessMethodKeysToPainlessBindings = Collections.unmodifiableMap(painlessMethodKeysToPainlessBindings);
     }
 
@@ -167,6 +174,14 @@ public PainlessField lookupPainlessField(Class<?> targetClass, boolean isStatic,
         return painlessField;
     }
 
+    public PainlessMethod lookupImportedPainlessMethod(String methodName, int arity) {
+        Objects.requireNonNull(methodName);
+
+        String painlessMethodKey = buildPainlessMethodKey(methodName, arity);
+
+        return painlessMethodKeysToImportedPainlessMethods.get(painlessMethodKey);
+    }
+
     public PainlessBinding lookupPainlessBinding(String methodName, int arity) {
         Objects.requireNonNull(methodName);