Skip to content

Commit

Permalink
Adds a lookup method to Painless for finding methods of all sub class…
Browse files Browse the repository at this point in the history 
…es (#77044)

This change adds a method to the PainlessLookup used to find methods of all allow listed sub classes. 
A method is specified for a specific super class, and then a list is built with all matching methods from 
the all the allow listed sub classes. If no matches are found, null is returned which is consistent with 
the behavior of the other look up methods. It is up to the caller to check.
  • Loading branch information
@jdconrad
jdconrad committed Aug 31, 2021
1 parent 5aa0dd7 commit 9e9a6d5
Show file tree
Hide file tree
Showing 3 changed files with 256 additions and 12 deletions.
67 changes: 67 additions & 0 deletions modules/lang-painless/src/main/java/org/elasticsearch/painless/lookup/PainlessLookup.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -150,8 +150,16 @@ public PainlessMethod lookupPainlessMethod(Class<?> targetClass, boolean isStati
Objects.requireNonNull(targetClass);
Objects.requireNonNull(methodName);

if (classesToPainlessClasses.containsKey(targetClass) == false) {
return null;
}

if (targetClass.isPrimitive()) {
targetClass = typeToBoxedType(targetClass);

if (classesToPainlessClasses.containsKey(targetClass) == false) {
return null;
}
}

String painlessMethodKey = buildPainlessMethodKey(methodName, methodArity);
Expand All @@ -162,6 +170,61 @@ public PainlessMethod lookupPainlessMethod(Class<?> targetClass, boolean isStati
return lookupPainlessObject(targetClass, objectLookup);
}

public List<PainlessMethod> lookupPainlessSubClassesMethod(String targetCanonicalClassName, String methodName, int methodArity) {
Objects.requireNonNull(targetCanonicalClassName);

Class<?> targetClass = canonicalTypeNameToType(targetCanonicalClassName);

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

return lookupPainlessSubClassesMethod(targetClass, methodName, methodArity);
}

public List<PainlessMethod> lookupPainlessSubClassesMethod(Class<?> targetClass, String methodName, int methodArity) {
Objects.requireNonNull(targetClass);
Objects.requireNonNull(methodName);

if (classesToPainlessClasses.containsKey(targetClass) == false) {
return null;
}

if (targetClass.isPrimitive()) {
targetClass = typeToBoxedType(targetClass);

if (classesToPainlessClasses.containsKey(targetClass) == false) {
return null;
}
}

String painlessMethodKey = buildPainlessMethodKey(methodName, methodArity);
List<Class<?>> subClasses = new ArrayList<>(classesToDirectSubClasses.get(targetClass));
Set<Class<?>> resolvedSubClasses = new HashSet<>();
List<PainlessMethod> subMethods = null;

while (subClasses.isEmpty() == false) {
Class<?> subClass = subClasses.remove(0);

if (resolvedSubClasses.add(subClass)) {
subClasses.addAll(classesToDirectSubClasses.get(subClass));

PainlessClass painlessClass = classesToPainlessClasses.get(subClass);
PainlessMethod painlessMethod = painlessClass.methods.get(painlessMethodKey);

if (painlessMethod != null) {
if (subMethods == null) {
subMethods = new ArrayList<>();
}

subMethods.add(painlessMethod);
}
}
}

return subMethods;
}

public PainlessField lookupPainlessField(String targetCanonicalClassName, boolean isStatic, String fieldName) {
Objects.requireNonNull(targetCanonicalClassName);

Expand All @@ -178,6 +241,10 @@ public PainlessField lookupPainlessField(Class<?> targetClass, boolean isStatic,
Objects.requireNonNull(targetClass);
Objects.requireNonNull(fieldName);

if (classesToPainlessClasses.containsKey(targetClass) == false) {
return null;
}

String painlessFieldKey = buildPainlessFieldKey(fieldName);
Function<PainlessClass, PainlessField> objectLookup = isStatic ?
targetPainlessClass -> targetPainlessClass.staticFields.get(painlessFieldKey) :
Expand Down
184 changes: 172 additions & 12 deletions modules/lang-painless/src/test/java/org/elasticsearch/painless/LookupTests.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -10,11 +10,13 @@

import org.elasticsearch.painless.lookup.PainlessLookup;
import org.elasticsearch.painless.lookup.PainlessLookupBuilder;
import org.elasticsearch.painless.lookup.PainlessMethod;
import org.elasticsearch.painless.spi.WhitelistLoader;
import org.elasticsearch.test.ESTestCase;
import org.junit.Before;

import java.util.Collections;
import java.util.List;
import java.util.Set;

public class LookupTests extends ESTestCase {
Expand All @@ -28,31 +30,54 @@ public void setup() {
));
}

public static class A { } // in whitelist
public static class B extends A { } // not in whitelist
public static class C extends B { } // in whitelist
public static class D extends B { } // in whitelist
public static class A { } // in whitelist
public static class B extends A { } // not in whitelist
public static class C extends B { // in whitelist
public String getString0() { return "C/0"; } // in whitelist
}
public static class D extends B { // in whitelist
public String getString0() { return "D/0"; } // in whitelist
public String getString1(int param0) { return "D/1 (" + param0 + ")"; } // in whitelist
}

public interface Z { } // in whitelist
public interface Y { } // not in whitelist
public interface X extends Y, Z { } // not in whitelist
public interface V extends Y, Z { } // in whitelist
public interface U extends X { } // in whitelist
public interface T extends V { } // in whitelist
public interface U extends X { // in whitelist
String getString2(int x, int y); // in whitelist
String getString1(int param0); // in whitelist
String getString0(); // not in whitelist
}
public interface T extends V { // in whitelist
String getString1(int param0); // in whitelist
int getInt0(); // in whitelist
}
public interface S extends U, X { } // in whitelist

public static class AA implements X { } // in whitelist
public static class AB extends AA implements S { } // not in whitelist
public static class AC extends AB implements V { } // in whitelist
public static class AD implements X, S, T { } // in whitelist
public static class AA implements X { } // in whitelist
public static class AB extends AA implements S { // not in whitelist
public String getString2(int x, int y) { return "" + x + y; } // not in whitelist
public String getString1(int param0) { return "" + param0; } // not in whitelist
public String getString0() { return ""; } // not in whitelist
}
public static class AC extends AB implements V { // in whitelist
public String getString2(int x, int y) { return "" + x + y; } // in whitelist
}
public static class AD extends AA implements X, S, T { // in whitelist
public String getString2(int x, int y) { return "" + x + y; } // in whitelist
public String getString1(int param0) { return "" + param0; } // in whitelist
public String getString0() { return ""; } // not in whitelist
public int getInt0() { return 0; } // in whitelist
}

public void testDirectSubClasses() {
Set<Class<?>> directSubClasses = painlessLookup.getDirectSubClasses(Object.class);
assertEquals(4, directSubClasses.size());
assertTrue(directSubClasses.contains(String.class));
assertTrue(directSubClasses.contains(A.class));
assertTrue(directSubClasses.contains(Z.class));
assertTrue(directSubClasses.contains(AA.class));
assertTrue(directSubClasses.contains(AD.class));

directSubClasses = painlessLookup.getDirectSubClasses(A.class);
assertEquals(2, directSubClasses.size());
Expand Down Expand Up @@ -101,8 +126,9 @@ public void testDirectSubClasses() {
assertTrue(directSubClasses.contains(AD.class));

directSubClasses = painlessLookup.getDirectSubClasses(AA.class);
assertEquals(1, directSubClasses.size());
assertEquals(2, directSubClasses.size());
assertTrue(directSubClasses.contains(AC.class));
assertTrue(directSubClasses.contains(AD.class));

directSubClasses = painlessLookup.getDirectSubClasses(AB.class);
assertNull(directSubClasses);
Expand All @@ -113,4 +139,138 @@ public void testDirectSubClasses() {
directSubClasses = painlessLookup.getDirectSubClasses(AD.class);
assertTrue(directSubClasses.isEmpty());
}

public void testDirectSubClassMethods() {
PainlessMethod CgetString0 = painlessLookup.lookupPainlessMethod(C.class, false, "getString0", 0);
PainlessMethod DgetString0 = painlessLookup.lookupPainlessMethod(D.class, false, "getString0", 0);
List<PainlessMethod> subMethods = painlessLookup.lookupPainlessSubClassesMethod(A.class, "getString0", 0);
assertNotNull(subMethods);
assertEquals(2, subMethods.size());
assertTrue(subMethods.contains(CgetString0));
assertTrue(subMethods.contains(DgetString0));

PainlessMethod DgetString1 = painlessLookup.lookupPainlessMethod(D.class, false, "getString1", 1);
subMethods = painlessLookup.lookupPainlessSubClassesMethod(A.class, "getString1", 1);
assertNotNull(subMethods);
assertEquals(1, subMethods.size());
assertTrue(subMethods.contains(DgetString1));

subMethods = painlessLookup.lookupPainlessSubClassesMethod(A.class, "getString2", 0);
assertNull(subMethods);

PainlessMethod ACgetString2 = painlessLookup.lookupPainlessMethod(AC.class, false, "getString2", 2);
PainlessMethod ADgetString2 = painlessLookup.lookupPainlessMethod(AD.class, false, "getString2", 2);
subMethods = painlessLookup.lookupPainlessSubClassesMethod(AA.class, "getString2", 2);
assertNotNull(subMethods);
assertEquals(2, subMethods.size());
assertTrue(subMethods.contains(ACgetString2));
assertTrue(subMethods.contains(ADgetString2));

PainlessMethod ADgetString1 = painlessLookup.lookupPainlessMethod(AD.class, false, "getString1", 1);
subMethods = painlessLookup.lookupPainlessSubClassesMethod(AA.class, "getString1", 1);
assertNotNull(subMethods);
assertEquals(1, subMethods.size());
assertTrue(subMethods.contains(ADgetString1));

subMethods = painlessLookup.lookupPainlessSubClassesMethod(AA.class, "getString0", 0);
assertNull(subMethods);

PainlessMethod ADgetInt0 = painlessLookup.lookupPainlessMethod(AD.class, false, "getInt0", 0);
subMethods = painlessLookup.lookupPainlessSubClassesMethod(AA.class, "getInt0", 0);
assertNotNull(subMethods);
assertEquals(1, subMethods.size());
assertTrue(subMethods.contains(ADgetInt0));

PainlessMethod UgetString2 = painlessLookup.lookupPainlessMethod(U.class, false, "getString2", 2);
subMethods = painlessLookup.lookupPainlessSubClassesMethod(Z.class, "getString2", 2);
assertNotNull(subMethods);
assertEquals(3, subMethods.size());
assertTrue(subMethods.contains(UgetString2));
assertTrue(subMethods.contains(ACgetString2));
assertTrue(subMethods.contains(ADgetString2));

PainlessMethod UgetString1 = painlessLookup.lookupPainlessMethod(U.class, false, "getString1", 1);
PainlessMethod TgetString1 = painlessLookup.lookupPainlessMethod(T.class, false, "getString1", 1);
subMethods = painlessLookup.lookupPainlessSubClassesMethod(Z.class, "getString1", 1);
assertNotNull(subMethods);
assertEquals(3, subMethods.size());
assertTrue(subMethods.contains(UgetString1));
assertTrue(subMethods.contains(TgetString1));
assertTrue(subMethods.contains(ADgetString1));

subMethods = painlessLookup.lookupPainlessSubClassesMethod(Z.class, "getString0", 0);
assertNull(subMethods);

PainlessMethod TgetInt0 = painlessLookup.lookupPainlessMethod(T.class, false, "getInt0", 0);
subMethods = painlessLookup.lookupPainlessSubClassesMethod(Z.class, "getInt0", 0);
assertNotNull(subMethods);
assertEquals(2, subMethods.size());
assertTrue(subMethods.contains(TgetInt0));
assertTrue(subMethods.contains(ADgetInt0));

subMethods = painlessLookup.lookupPainlessSubClassesMethod(V.class, "getString2", 2);
assertNotNull(subMethods);
assertEquals(2, subMethods.size());
assertTrue(subMethods.contains(ACgetString2));
assertTrue(subMethods.contains(ADgetString2));

subMethods = painlessLookup.lookupPainlessSubClassesMethod(V.class, "getString1", 1);
assertNotNull(subMethods);
assertEquals(2, subMethods.size());
assertTrue(subMethods.contains(TgetString1));
assertTrue(subMethods.contains(ADgetString1));

subMethods = painlessLookup.lookupPainlessSubClassesMethod(V.class, "getString0", 0);
assertNull(subMethods);

subMethods = painlessLookup.lookupPainlessSubClassesMethod(V.class, "getInt0", 0);
assertNotNull(subMethods);
assertEquals(2, subMethods.size());
assertTrue(subMethods.contains(TgetInt0));
assertTrue(subMethods.contains(ADgetInt0));

subMethods = painlessLookup.lookupPainlessSubClassesMethod(U.class, "getString2", 2);
assertNotNull(subMethods);
assertEquals(2, subMethods.size());
assertTrue(subMethods.contains(ACgetString2));
assertTrue(subMethods.contains(ADgetString2));

subMethods = painlessLookup.lookupPainlessSubClassesMethod(U.class, "getString1", 1);
assertNotNull(subMethods);
assertEquals(1, subMethods.size());
assertTrue(subMethods.contains(ADgetString1));

subMethods = painlessLookup.lookupPainlessSubClassesMethod(U.class, "getString0", 0);
assertNull(subMethods);

subMethods = painlessLookup.lookupPainlessSubClassesMethod(U.class, "getInt0", 0);
assertNotNull(subMethods);
assertEquals(1, subMethods.size());
assertTrue(subMethods.contains(ADgetInt0));

subMethods = painlessLookup.lookupPainlessSubClassesMethod(V.class, "getInt0", 0);
assertNotNull(subMethods);
assertEquals(2, subMethods.size());
assertTrue(subMethods.contains(TgetInt0));
assertTrue(subMethods.contains(ADgetInt0));

subMethods = painlessLookup.lookupPainlessSubClassesMethod(S.class, "getString2", 2);
assertNotNull(subMethods);
assertEquals(2, subMethods.size());
assertTrue(subMethods.contains(ACgetString2));
assertTrue(subMethods.contains(ADgetString2));

subMethods = painlessLookup.lookupPainlessSubClassesMethod(S.class, "getString1", 1);
assertNotNull(subMethods);
assertEquals(1, subMethods.size());
assertTrue(subMethods.contains(ADgetString1));

subMethods = painlessLookup.lookupPainlessSubClassesMethod(S.class, "getString0", 0);
assertNull(subMethods);

subMethods = painlessLookup.lookupPainlessSubClassesMethod(S.class, "getInt0", 0);
assertNotNull(subMethods);
assertEquals(1, subMethods.size());
assertTrue(subMethods.contains(ADgetInt0));
}
}
17 changes: 17 additions & 0 deletions ...-painless/src/test/resources/org/elasticsearch/painless/org.elasticsearch.painless.lookup
View file
Original file line number Diff line number Diff line change
@@ -1,13 +1,22 @@
class int @no_import {
}

class java.lang.Object {
}

class java.lang.String {
}

class org.elasticsearch.painless.LookupTests$A {
}

class org.elasticsearch.painless.LookupTests$C {
String getString0()
}

class org.elasticsearch.painless.LookupTests$D {
String getString0()
String getString1(int)
}

class org.elasticsearch.painless.LookupTests$Z {
Expand All @@ -17,9 +26,13 @@ class org.elasticsearch.painless.LookupTests$V {
}

class org.elasticsearch.painless.LookupTests$U {
String getString2(int, int);
String getString1(int);
}

class org.elasticsearch.painless.LookupTests$T {
String getString1(int);
int getInt0();
}

class org.elasticsearch.painless.LookupTests$S {
Expand All @@ -29,7 +42,11 @@ class org.elasticsearch.painless.LookupTests$AA {
}

class org.elasticsearch.painless.LookupTests$AC {
String getString2(int, int);
}

class org.elasticsearch.painless.LookupTests$AD {
String getString2(int, int)
String getString1(int)
int getInt0()
}

0 comments on commit 9e9a6d5

Please sign in to comment.