Skip to content

Add a direct sub classes data structure to the Painless lookup (#76955) #9

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
MitchLewis930 wants to merge 1 commit into
base: pr_019_before
Choose a base branch
from
Open

Add a direct sub classes data structure to the Painless lookup (#76955) #9

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
86 changes: 54 additions & 32 deletions modules/lang-painless/src/main/java/org/elasticsearch/painless/lookup/PainlessLookup.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -9,6 +9,10 @@
package org.elasticsearch.painless.lookup;

import java.lang.invoke.MethodHandle;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.HashSet;
import java.util.List;
import java.util.Map;
import java.util.Objects;
import java.util.Set;
Expand All @@ -25,6 +29,7 @@ public final class PainlessLookup {
private final Map<String, Class<?>> javaClassNamesToClasses;
private final Map<String, Class<?>> canonicalClassNamesToClasses;
private final Map<Class<?>, PainlessClass> classesToPainlessClasses;
private final Map<Class<?>, Set<Class<?>>> classesToDirectSubClasses;

private final Map<String, PainlessMethod> painlessMethodKeysToImportedPainlessMethods;
private final Map<String, PainlessClassBinding> painlessMethodKeysToPainlessClassBindings;
Expand All @@ -34,13 +39,15 @@ public final class PainlessLookup {
Map<String, Class<?>> javaClassNamesToClasses,
Map<String, Class<?>> canonicalClassNamesToClasses,
Map<Class<?>, PainlessClass> classesToPainlessClasses,
Map<Class<?>, Set<Class<?>>> classesToDirectSubClasses,
Map<String, PainlessMethod> painlessMethodKeysToImportedPainlessMethods,
Map<String, PainlessClassBinding> painlessMethodKeysToPainlessClassBindings,
Map<String, PainlessInstanceBinding> painlessMethodKeysToPainlessInstanceBindings) {

Objects.requireNonNull(javaClassNamesToClasses);
Objects.requireNonNull(canonicalClassNamesToClasses);
Objects.requireNonNull(classesToPainlessClasses);
Objects.requireNonNull(classesToDirectSubClasses);

Objects.requireNonNull(painlessMethodKeysToImportedPainlessMethods);
Objects.requireNonNull(painlessMethodKeysToPainlessClassBindings);
Expand All @@ -49,6 +56,7 @@ public final class PainlessLookup {
this.javaClassNamesToClasses = javaClassNamesToClasses;
this.canonicalClassNamesToClasses = Map.copyOf(canonicalClassNamesToClasses);
this.classesToPainlessClasses = Map.copyOf(classesToPainlessClasses);
this.classesToDirectSubClasses = Map.copyOf(classesToDirectSubClasses);

this.painlessMethodKeysToImportedPainlessMethods = Map.copyOf(painlessMethodKeysToImportedPainlessMethods);
this.painlessMethodKeysToPainlessClassBindings = Map.copyOf(painlessMethodKeysToPainlessClassBindings);
Expand All @@ -75,6 +83,10 @@ public Set<Class<?>> getClasses() {
return classesToPainlessClasses.keySet();
}

public Set<Class<?>> getDirectSubClasses(Class<?> superClass) {
return classesToDirectSubClasses.get(superClass);
}

public Set<String> getImportedPainlessMethodsKeys() {
return painlessMethodKeysToImportedPainlessMethods.keySet();
}
Expand Down Expand Up @@ -142,16 +154,12 @@ public PainlessMethod lookupPainlessMethod(Class<?> targetClass, boolean isStati
targetClass = typeToBoxedType(targetClass);
}

PainlessClass targetPainlessClass = classesToPainlessClasses.get(targetClass);
String painlessMethodKey = buildPainlessMethodKey(methodName, methodArity);
Function<PainlessClass, PainlessMethod> objectLookup = isStatic ?
targetPainlessClass -> targetPainlessClass.staticMethods.get(painlessMethodKey) :
targetPainlessClass -> targetPainlessClass.methods.get(painlessMethodKey);

if (targetPainlessClass == null) {
return null;
}

return isStatic ?
targetPainlessClass.staticMethods.get(painlessMethodKey) :
targetPainlessClass.methods.get(painlessMethodKey);
return lookupPainlessObject(targetClass, objectLookup);
}

public PainlessField lookupPainlessField(String targetCanonicalClassName, boolean isStatic, String fieldName) {
Expand All @@ -170,22 +178,12 @@ public PainlessField lookupPainlessField(Class<?> targetClass, boolean isStatic,
Objects.requireNonNull(targetClass);
Objects.requireNonNull(fieldName);

PainlessClass targetPainlessClass = classesToPainlessClasses.get(targetClass);
String painlessFieldKey = buildPainlessFieldKey(fieldName);
Function<PainlessClass, PainlessField> objectLookup = isStatic ?
targetPainlessClass -> targetPainlessClass.staticFields.get(painlessFieldKey) :
targetPainlessClass -> targetPainlessClass.fields.get(painlessFieldKey);

if (targetPainlessClass == null) {
return null;
}

PainlessField painlessField = isStatic ?
targetPainlessClass.staticFields.get(painlessFieldKey) :
targetPainlessClass.fields.get(painlessFieldKey);

if (painlessField == null) {
return null;
}

return painlessField;
return lookupPainlessObject(targetClass, objectLookup);
}

public PainlessMethod lookupImportedPainlessMethod(String methodName, int arity) {
Expand Down Expand Up @@ -230,7 +228,7 @@ public PainlessMethod lookupRuntimePainlessMethod(Class<?> originalTargetClass,
Function<PainlessClass, PainlessMethod> objectLookup =
targetPainlessClass -> targetPainlessClass.runtimeMethods.get(painlessMethodKey);

return lookupRuntimePainlessObject(originalTargetClass, objectLookup);
return lookupPainlessObject(originalTargetClass, objectLookup);
}

public MethodHandle lookupRuntimeGetterMethodHandle(Class<?> originalTargetClass, String getterName) {
Expand All @@ -239,7 +237,7 @@ public MethodHandle lookupRuntimeGetterMethodHandle(Class<?> originalTargetClass

Function<PainlessClass, MethodHandle> objectLookup = targetPainlessClass -> targetPainlessClass.getterMethodHandles.get(getterName);

return lookupRuntimePainlessObject(originalTargetClass, objectLookup);
return lookupPainlessObject(originalTargetClass, objectLookup);
}

public MethodHandle lookupRuntimeSetterMethodHandle(Class<?> originalTargetClass, String setterName) {
Expand All @@ -248,10 +246,13 @@ public MethodHandle lookupRuntimeSetterMethodHandle(Class<?> originalTargetClass

Function<PainlessClass, MethodHandle> objectLookup = targetPainlessClass -> targetPainlessClass.setterMethodHandles.get(setterName);

return lookupRuntimePainlessObject(originalTargetClass, objectLookup);
return lookupPainlessObject(originalTargetClass, objectLookup);
}

private <T> T lookupRuntimePainlessObject(Class<?> originalTargetClass, Function<PainlessClass, T> objectLookup) {
private <T> T lookupPainlessObject(Class<?> originalTargetClass, Function<PainlessClass, T> objectLookup) {
Objects.requireNonNull(originalTargetClass);
Objects.requireNonNull(objectLookup);

Class<?> currentTargetClass = originalTargetClass;

while (currentTargetClass != null) {
Expand All @@ -268,17 +269,38 @@ private <T> T lookupRuntimePainlessObject(Class<?> originalTargetClass, Function
currentTargetClass = currentTargetClass.getSuperclass();
}

if (originalTargetClass.isInterface()) {
PainlessClass targetPainlessClass = classesToPainlessClasses.get(Object.class);

if (targetPainlessClass != null) {
T painlessObject = objectLookup.apply(targetPainlessClass);

if (painlessObject != null) {
return painlessObject;
}
}
}

currentTargetClass = originalTargetClass;
Set<Class<?>> resolvedInterfaces = new HashSet<>();

while (currentTargetClass != null) {
for (Class<?> targetInterface : currentTargetClass.getInterfaces()) {
PainlessClass targetPainlessClass = classesToPainlessClasses.get(targetInterface);
List<Class<?>> targetInterfaces = new ArrayList<>(Arrays.asList(currentTargetClass.getInterfaces()));

while (targetInterfaces.isEmpty() == false) {
Class<?> targetInterface = targetInterfaces.remove(0);

if (resolvedInterfaces.add(targetInterface)) {
PainlessClass targetPainlessClass = classesToPainlessClasses.get(targetInterface);

if (targetPainlessClass != null) {
T painlessObject = objectLookup.apply(targetPainlessClass);

if (targetPainlessClass != null) {
T painlessObject = objectLookup.apply(targetPainlessClass);
if (painlessObject != null) {
return painlessObject;
}

if (painlessObject != null) {
return painlessObject;
targetInterfaces.addAll(Arrays.asList(targetInterface.getInterfaces()));
}
Comment on lines +289 to 304
Copy link

@coderabbitai coderabbitai bot Jan 29, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

⚠️ Potential issue | 🟠 Major

Ensure interface traversal reaches allow-listed ancestors even through non-allow-listed interfaces.

Currently, parent interfaces are only enqueued when the current interface is allow‑listed. If a class implements a non‑allow‑listed interface that extends an allow‑listed one, lookups can miss valid members. This is a correctness regression in hierarchy resolution.

🔧 Proposed fix 
-                if (resolvedInterfaces.add(targetInterface)) {
-                    PainlessClass targetPainlessClass = classesToPainlessClasses.get(targetInterface);
-
-                    if (targetPainlessClass != null) {
-                        T painlessObject = objectLookup.apply(targetPainlessClass);
-
-                        if (painlessObject != null) {
-                            return painlessObject;
-                        }
-
-                        targetInterfaces.addAll(Arrays.asList(targetInterface.getInterfaces()));
-                    }
-                }
+                if (resolvedInterfaces.add(targetInterface)) {
+                    PainlessClass targetPainlessClass = classesToPainlessClasses.get(targetInterface);
+
+                    if (targetPainlessClass != null) {
+                        T painlessObject = objectLookup.apply(targetPainlessClass);
+                        if (painlessObject != null) {
+                            return painlessObject;
+                        }
+                    }
+                    // Always traverse parent interfaces, even when this interface isn't allow-listed.
+                    targetInterfaces.addAll(Arrays.asList(targetInterface.getInterfaces()));
+                }
📝 Committable suggestion

‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.

Suggested change
while (targetInterfaces.isEmpty() == false) {
Class<?> targetInterface = targetInterfaces.remove(0);
if (resolvedInterfaces.add(targetInterface)) {
PainlessClass targetPainlessClass = classesToPainlessClasses.get(targetInterface);
if (targetPainlessClass != null) {
T painlessObject = objectLookup.apply(targetPainlessClass);
if (targetPainlessClass != null) {
T painlessObject = objectLookup.apply(targetPainlessClass);
if (painlessObject != null) {
return painlessObject;
}
if (painlessObject != null) {
return painlessObject;
targetInterfaces.addAll(Arrays.asList(targetInterface.getInterfaces()));
}
while (targetInterfaces.isEmpty() == false) {
Class<?> targetInterface = targetInterfaces.remove(0);
if (resolvedInterfaces.add(targetInterface)) {
PainlessClass targetPainlessClass = classesToPainlessClasses.get(targetInterface);
if (targetPainlessClass != null) {
T painlessObject = objectLookup.apply(targetPainlessClass);
if (painlessObject != null) {
return painlessObject;
}
}
// Always traverse parent interfaces, even when this interface isn't allow-listed.
targetInterfaces.addAll(Arrays.asList(targetInterface.getInterfaces()));
}
🤖 Prompt for AI Agents 
In
`@modules/lang-painless/src/main/java/org/elasticsearch/painless/lookup/PainlessLookup.java`
around lines 289 - 304, The traversal skips enqueuing parent interfaces when a
visited interface isn't allow‑listed because targetInterface.getInterfaces() is
only added inside the "if (targetPainlessClass != null)" block; modify the loop
in PainlessLookup so that after a targetInterface is marked visited
(resolvedInterfaces.add) you always enqueue its super-interfaces (call
targetInterfaces.addAll(Arrays.asList(targetInterface.getInterfaces())))
regardless of whether classesToPainlessClasses.containsKey(targetInterface) or
objectLookup returned null, preserving the existing checks for painlessObject
via objectLookup and avoiding duplicate visits via resolvedInterfaces.

}
}
Expand Down
106 changes: 53 additions & 53 deletions .../lang-painless/src/main/java/org/elasticsearch/painless/lookup/PainlessLookupBuilder.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -42,11 +42,14 @@
import java.security.SecureClassLoader;
import java.security.cert.Certificate;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.Collections;
import java.util.HashMap;
import java.util.HashSet;
import java.util.List;
import java.util.Map;
import java.util.Objects;
import java.util.Set;
import java.util.function.Supplier;
import java.util.regex.Pattern;

Expand Down Expand Up @@ -189,6 +192,7 @@ public static PainlessLookup buildFromWhitelists(List<Whitelist> whitelists) {
// of the values of javaClassNamesToClasses.
private final Map<String, Class<?>> canonicalClassNamesToClasses;
private final Map<Class<?>, PainlessClassBuilder> classesToPainlessClassBuilders;
private final Map<Class<?>, Set<Class<?>>> classesToDirectSubClasses;

private final Map<String, PainlessMethod> painlessMethodKeysToImportedPainlessMethods;
private final Map<String, PainlessClassBinding> painlessMethodKeysToPainlessClassBindings;
Expand All @@ -198,6 +202,7 @@ public PainlessLookupBuilder() {
javaClassNamesToClasses = new HashMap<>();
canonicalClassNamesToClasses = new HashMap<>();
classesToPainlessClassBuilders = new HashMap<>();
classesToDirectSubClasses = new HashMap<>();

painlessMethodKeysToImportedPainlessMethods = new HashMap<>();
painlessMethodKeysToPainlessClassBindings = new HashMap<>();
Expand Down Expand Up @@ -1255,7 +1260,7 @@ public void addPainlessInstanceBinding(
}

public PainlessLookup build() {
copyPainlessClassMembers();
buildPainlessClassHierarchy();
setFunctionalInterfaceMethods();
generateRuntimeMethods();
cacheRuntimeHandles();
Expand Down Expand Up @@ -1286,71 +1291,66 @@ public PainlessLookup build() {
javaClassNamesToClasses,
canonicalClassNamesToClasses,
classesToPainlessClasses,
classesToDirectSubClasses,
painlessMethodKeysToImportedPainlessMethods,
painlessMethodKeysToPainlessClassBindings,
painlessMethodKeysToPainlessInstanceBindings);
}

private void copyPainlessClassMembers() {
for (Class<?> parentClass : classesToPainlessClassBuilders.keySet()) {
copyPainlessInterfaceMembers(parentClass, parentClass);

Class<?> childClass = parentClass.getSuperclass();

while (childClass != null) {
if (classesToPainlessClassBuilders.containsKey(childClass)) {
copyPainlessClassMembers(childClass, parentClass);
}

copyPainlessInterfaceMembers(childClass, parentClass);
childClass = childClass.getSuperclass();
}
}

for (Class<?> javaClass : classesToPainlessClassBuilders.keySet()) {
if (javaClass.isInterface()) {
copyPainlessClassMembers(Object.class, javaClass);
}
}
}

private void copyPainlessInterfaceMembers(Class<?> parentClass, Class<?> targetClass) {
for (Class<?> childClass : parentClass.getInterfaces()) {
if (classesToPainlessClassBuilders.containsKey(childClass)) {
copyPainlessClassMembers(childClass, targetClass);
}

copyPainlessInterfaceMembers(childClass, targetClass);
private void buildPainlessClassHierarchy() {
for (Class<?> targetClass : classesToPainlessClassBuilders.keySet()) {
classesToDirectSubClasses.put(targetClass, new HashSet<>());
}
}

private void copyPainlessClassMembers(Class<?> originalClass, Class<?> targetClass) {
PainlessClassBuilder originalPainlessClassBuilder = classesToPainlessClassBuilders.get(originalClass);
PainlessClassBuilder targetPainlessClassBuilder = classesToPainlessClassBuilders.get(targetClass);
for (Class<?> subClass : classesToPainlessClassBuilders.keySet()) {
List<Class<?>> superInterfaces = new ArrayList<>(Arrays.asList(subClass.getInterfaces()));

Objects.requireNonNull(originalPainlessClassBuilder);
Objects.requireNonNull(targetPainlessClassBuilder);
// we check for Object.class as part of the allow listed classes because
// it is possible for the compiler to work without Object
if (subClass.isInterface() && superInterfaces.isEmpty() && classesToPainlessClassBuilders.containsKey(Object.class)) {
classesToDirectSubClasses.get(Object.class).add(subClass);
} else {
Class<?> superClass = subClass.getSuperclass();

// this finds the nearest super class for a given sub class
// because the allow list may have gaps between classes
// example:
// class A {} // allowed
// class B extends A // not allowed
// class C extends B // allowed
// in this case C is considered a direct sub class of A
while (superClass != null) {
if (classesToPainlessClassBuilders.containsKey(superClass)) {
break;
} else {
// this ensures all interfaces from a sub class that
// is not allow listed are checked if they are
// considered a direct super class of the sub class
// because these interfaces may still be allow listed
// even if their sub class is not
superInterfaces.addAll(Arrays.asList(superClass.getInterfaces()));
}

for (Map.Entry<String, PainlessMethod> painlessMethodEntry : originalPainlessClassBuilder.methods.entrySet()) {
String painlessMethodKey = painlessMethodEntry.getKey();
PainlessMethod newPainlessMethod = painlessMethodEntry.getValue();
PainlessMethod existingPainlessMethod = targetPainlessClassBuilder.methods.get(painlessMethodKey);
superClass = superClass.getSuperclass();
}

if (existingPainlessMethod == null || existingPainlessMethod.targetClass != newPainlessMethod.targetClass &&
existingPainlessMethod.targetClass.isAssignableFrom(newPainlessMethod.targetClass)) {
targetPainlessClassBuilder.methods.put(painlessMethodKey.intern(), newPainlessMethod);
if (superClass != null) {
classesToDirectSubClasses.get(superClass).add(subClass);
}
}
}

for (Map.Entry<String, PainlessField> painlessFieldEntry : originalPainlessClassBuilder.fields.entrySet()) {
String painlessFieldKey = painlessFieldEntry.getKey();
PainlessField newPainlessField = painlessFieldEntry.getValue();
PainlessField existingPainlessField = targetPainlessClassBuilder.fields.get(painlessFieldKey);
Set<Class<?>> resolvedInterfaces = new HashSet<>();

while (superInterfaces.isEmpty() == false) {
Class<?> superInterface = superInterfaces.remove(0);

if (existingPainlessField == null ||
existingPainlessField.javaField.getDeclaringClass() != newPainlessField.javaField.getDeclaringClass() &&
existingPainlessField.javaField.getDeclaringClass().isAssignableFrom(newPainlessField.javaField.getDeclaringClass())) {
targetPainlessClassBuilder.fields.put(painlessFieldKey.intern(), newPainlessField);
if (resolvedInterfaces.add(superInterface)) {
if (classesToPainlessClassBuilders.containsKey(superInterface)) {
classesToDirectSubClasses.get(superInterface).add(subClass);
} else {
superInterfaces.addAll(Arrays.asList(superInterface.getInterfaces()));
}
}
}
}
}
Expand Down
Loading