Add needs methods for specific variables to Painless script context factories. (#25267)

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

@@ -33,6 +33,7 @@ import java.security.cert.Certificate;
 import java.util.concurrent.atomic.AtomicInteger;
 
 import static org.elasticsearch.painless.WriterConstants.CLASS_NAME;
+import static org.elasticsearch.painless.node.SSource.MainMethodReserved;
 
 /**
  * The Compiler is the entry point for generating a Painless script.  The compiler will receive a Painless
@@ -143,7 +144,7 @@ final class Compiler {
      * @param settings The CompilerSettings to be used during the compilation.
      * @return An executable script that implements both a specified interface and is a subclass of {@link PainlessScript}
      */
-    Constructor<?> compile(Loader loader, String name, String source, CompilerSettings settings) {
+    Constructor<?> compile(Loader loader, MainMethodReserved reserved, String name, String source, CompilerSettings settings) {
         if (source.length() > MAXIMUM_SOURCE_LENGTH) {
             throw new IllegalArgumentException("Scripts may be no longer than " + MAXIMUM_SOURCE_LENGTH +
                 " characters.  The passed in script is " + source.length() + " characters.  Consider using a" +
@@ -151,7 +152,7 @@ final class Compiler {
         }
 
         ScriptClassInfo scriptClassInfo = new ScriptClassInfo(definition, base);
-        SSource root = Walker.buildPainlessTree(scriptClassInfo, name, source, settings, definition,
+        SSource root = Walker.buildPainlessTree(scriptClassInfo, reserved, name, source, settings, definition,
                 null);
         root.analyze(definition);
         root.write();
@@ -183,7 +184,7 @@ final class Compiler {
         }
 
         ScriptClassInfo scriptClassInfo = new ScriptClassInfo(definition, base);
-        SSource root = Walker.buildPainlessTree(scriptClassInfo, name, source, settings, definition,
+        SSource root = Walker.buildPainlessTree(scriptClassInfo, new MainMethodReserved(), name, source, settings, definition,
                 debugStream);
         root.analyze(definition);
         root.write();

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

@@ -19,6 +19,7 @@
 
 package org.elasticsearch.painless;
 
+import org.apache.logging.log4j.core.tools.Generate;
 import org.apache.lucene.index.LeafReaderContext;
 import org.elasticsearch.SpecialPermission;
 import org.elasticsearch.common.component.AbstractComponent;
@@ -51,6 +52,7 @@ import java.util.List;
 import java.util.Map;
 
 import static org.elasticsearch.painless.WriterConstants.OBJECT_TYPE;
+import static org.elasticsearch.painless.node.SSource.MainMethodReserved;
 
 /**
  * Implementation of a ScriptEngine for the Painless language.
@@ -157,12 +159,13 @@ public final class PainlessScriptEngine extends AbstractComponent implements Scr
                 }
             });
 
-            compile(contextsToCompilers.get(context), loader, scriptName, scriptSource, params);
+            MainMethodReserved reserved = new MainMethodReserved();
+            compile(contextsToCompilers.get(context), loader, reserved, scriptName, scriptSource, params);
 
             if (context.statefulFactoryClazz != null) {
-                return generateFactory(loader, context, generateStatefulFactory(loader, context));
+                return generateFactory(loader, context, reserved, generateStatefulFactory(loader, context, reserved));
             } else {
-                return generateFactory(loader, context, WriterConstants.CLASS_TYPE);
+                return generateFactory(loader, context, reserved, WriterConstants.CLASS_TYPE);
             }
         }
     }
@@ -178,7 +181,7 @@ public final class PainlessScriptEngine extends AbstractComponent implements Scr
      * @param <T> The factory class.
      * @return A factory class that will return script instances.
      */
-    private <T> Type generateStatefulFactory(Loader loader, ScriptContext<T> context) {
+    private <T> Type generateStatefulFactory(Loader loader, ScriptContext<T> context, MainMethodReserved reserved) {
         int classFrames = ClassWriter.COMPUTE_FRAMES | ClassWriter.COMPUTE_MAXS;
         int classAccess = Opcodes.ACC_PUBLIC | Opcodes.ACC_SUPER | Opcodes.ACC_FINAL;
         String interfaceBase = Type.getType(context.statefulFactoryClazz).getInternalName();
@@ -259,6 +262,7 @@ public final class PainlessScriptEngine extends AbstractComponent implements Scr
         adapter.returnValue();
         adapter.endMethod();
 
+        writeNeedsMethods(context.statefulFactoryClazz, writer, reserved);
         writer.visitEnd();
 
         loader.defineFactory(className.replace('/', '.'), writer.toByteArray());
@@ -278,7 +282,7 @@ public final class PainlessScriptEngine extends AbstractComponent implements Scr
      * @param <T> The factory class.
      * @return A factory class that will return script instances.
      */
-    private <T> T generateFactory(Loader loader, ScriptContext<T> context, Type classType) {
+    private <T> T generateFactory(Loader loader, ScriptContext<T> context, MainMethodReserved reserved, Type classType) {
         int classFrames = ClassWriter.COMPUTE_FRAMES | ClassWriter.COMPUTE_MAXS;
         int classAccess = Opcodes.ACC_PUBLIC | Opcodes.ACC_SUPER| Opcodes.ACC_FINAL;
         String interfaceBase = Type.getType(context.factoryClazz).getInternalName();
@@ -329,6 +333,7 @@ public final class PainlessScriptEngine extends AbstractComponent implements Scr
         adapter.returnValue();
         adapter.endMethod();
 
+        writeNeedsMethods(context.factoryClazz, writer, reserved);
         writer.visitEnd();
 
         Class<?> factory = loader.defineFactory(className.replace('/', '.'), writer.toByteArray());
@@ -341,6 +346,27 @@ public final class PainlessScriptEngine extends AbstractComponent implements Scr
         }
     }
 
+    private void writeNeedsMethods(Class<?> clazz, ClassWriter writer, MainMethodReserved reserved) {
+        for (Method method : clazz.getMethods()) {
+            if (method.getName().startsWith("needs") &&
+                method.getReturnType().equals(boolean.class) && method.getParameterTypes().length == 0) {
+                String name = method.getName();
+                name = name.substring(5);
+                name = Character.toLowerCase(name.charAt(0)) + name.substring(1);
+
+                org.objectweb.asm.commons.Method needs = new org.objectweb.asm.commons.Method(method.getName(),
+                    MethodType.methodType(boolean.class).toMethodDescriptorString());
+
+                GeneratorAdapter adapter = new GeneratorAdapter(Opcodes.ASM5, needs,
+                    writer.visitMethod(Opcodes.ACC_PUBLIC, needs.getName(), needs.getDescriptor(), null, null));
+                adapter.visitCode();
+                adapter.push(reserved.getUsedVariables().contains(name));
+                adapter.returnValue();
+                adapter.endMethod();
+            }
+        }
+    }
+
     Object compile(Compiler compiler, String scriptName, String source, Map<String, String> params, Object... args) {
         final CompilerSettings compilerSettings;
 
@@ -398,7 +424,7 @@ public final class PainlessScriptEngine extends AbstractComponent implements Scr
                 @Override
                 public Object run() {
                     String name = scriptName == null ? INLINE_NAME : scriptName;
-                    Constructor<?> constructor = compiler.compile(loader, name, source, compilerSettings);
+                    Constructor<?> constructor = compiler.compile(loader, new MainMethodReserved(), name, source, compilerSettings);
 
                     try {
                         return constructor.newInstance(args);
@@ -414,7 +440,8 @@ public final class PainlessScriptEngine extends AbstractComponent implements Scr
         }
     }
 
-    void compile(Compiler compiler, Loader loader, String scriptName, String source, Map<String, String> params) {
+    void compile(Compiler compiler, Loader loader, MainMethodReserved reserved,
+                 String scriptName, String source, Map<String, String> params) {
         final CompilerSettings compilerSettings;
 
         if (params.isEmpty()) {
@@ -460,7 +487,7 @@ public final class PainlessScriptEngine extends AbstractComponent implements Scr
                 @Override
                 public Void run() {
                     String name = scriptName == null ? INLINE_NAME : scriptName;
-                    compiler.compile(loader, name, source, compilerSettings);
+                    compiler.compile(loader, reserved, name, source, compilerSettings);
 
                     return null;
                 }

modules/lang-painless/src/main/java/org/elasticsearch/painless/antlr/Walker.java

@@ -174,11 +174,10 @@ import java.util.List;
  */
 public final class Walker extends PainlessParserBaseVisitor<ANode> {
 
-    public static SSource buildPainlessTree(ScriptClassInfo mainMethod, String sourceName,
+    public static SSource buildPainlessTree(ScriptClassInfo mainMethod, MainMethodReserved reserved, String sourceName,
                                             String sourceText, CompilerSettings settings, Definition definition,
                                             Printer debugStream) {
-        return new Walker(mainMethod, sourceName, sourceText, settings, definition,
-                debugStream).source;
+        return new Walker(mainMethod, reserved, sourceName, sourceText, settings, definition, debugStream).source;
     }
 
     private final ScriptClassInfo scriptClassInfo;
@@ -193,9 +192,10 @@ public final class Walker extends PainlessParserBaseVisitor<ANode> {
     private final Globals globals;
     private int syntheticCounter = 0;
 
-    private Walker(ScriptClassInfo scriptClassInfo, String sourceName, String sourceText,
+    private Walker(ScriptClassInfo scriptClassInfo, MainMethodReserved reserved, String sourceName, String sourceText,
                    CompilerSettings settings, Definition definition, Printer debugStream) {
         this.scriptClassInfo = scriptClassInfo;
+        this.reserved.push(reserved);
         this.debugStream = debugStream;
         this.settings = settings;
         this.sourceName = Location.computeSourceName(sourceName, sourceText);
@@ -252,8 +252,6 @@ public final class Walker extends PainlessParserBaseVisitor<ANode> {
 
     @Override
     public ANode visitSource(SourceContext ctx) {
-        reserved.push(new MainMethodReserved());
-
         List<SFunction> functions = new ArrayList<>();
 
         for (FunctionContext function : ctx.function()) {

modules/lang-painless/src/test/java/org/elasticsearch/painless/FactoryTests.java

@@ -55,15 +55,41 @@ public class FactoryTests extends ScriptTestCase {
             return y*2;
         }
 
+        public int getC() {
+            return -1;
+        }
+
+        public int getD() {
+            return 2;
+        }
+
         public static final String[] PARAMETERS = new String[] {"test"};
         public abstract Object execute(int test);
 
+        public abstract boolean needsTest();
+        public abstract boolean needsNothing();
+        public abstract boolean needsX();
+        public abstract boolean needsC();
+        public abstract boolean needsD();
+
         public interface StatefulFactory {
             StatefulFactoryTestScript newInstance(int a, int b);
+
+            boolean needsTest();
+            boolean needsNothing();
+            boolean needsX();
+            boolean needsC();
+            boolean needsD();
         }
 
         public interface Factory {
             StatefulFactory newFactory(int x, int y);
+
+            boolean needsTest();
+            boolean needsNothing();
+            boolean needsX();
+            boolean needsC();
+            boolean needsD();
         }
 
         public static final ScriptContext<StatefulFactoryTestScript.Factory> CONTEXT =
@@ -72,12 +98,27 @@ public class FactoryTests extends ScriptTestCase {
 
     public void testStatefulFactory() {
         StatefulFactoryTestScript.Factory factory = scriptEngine.compile(
-            "stateful_factory_test", "test + x + y", StatefulFactoryTestScript.CONTEXT, Collections.emptyMap());
+            "stateful_factory_test", "test + x + y + d", StatefulFactoryTestScript.CONTEXT, Collections.emptyMap());
         StatefulFactoryTestScript.StatefulFactory statefulFactory = factory.newFactory(1, 2);
         StatefulFactoryTestScript script = statefulFactory.newInstance(3, 4);
-        assertEquals(22, script.execute(3));
+        assertEquals(24, script.execute(3));
         statefulFactory.newInstance(5, 6);
-        assertEquals(26, script.execute(7));
+        assertEquals(28, script.execute(7));
+        assertEquals(true, script.needsTest());
+        assertEquals(false, script.needsNothing());
+        assertEquals(true, script.needsX());
+        assertEquals(false, script.needsC());
+        assertEquals(true, script.needsD());
+        assertEquals(true, statefulFactory.needsTest());
+        assertEquals(false, statefulFactory.needsNothing());
+        assertEquals(true, statefulFactory.needsX());
+        assertEquals(false, statefulFactory.needsC());
+        assertEquals(true, statefulFactory.needsD());
+        assertEquals(true, factory.needsTest());
+        assertEquals(false, factory.needsNothing());
+        assertEquals(true, factory.needsX());
+        assertEquals(false, factory.needsC());
+        assertEquals(true, factory.needsD());
     }
 
     public abstract static class FactoryTestScript {
@@ -96,6 +137,9 @@ public class FactoryTests extends ScriptTestCase {
 
         public interface Factory {
             FactoryTestScript newInstance(Map<String, Object> params);
+
+            boolean needsTest();
+            boolean needsNothing();
         }
 
         public static final ScriptContext<FactoryTestScript.Factory> CONTEXT =
@@ -111,6 +155,8 @@ public class FactoryTests extends ScriptTestCase {
         script = factory.newInstance(Collections.singletonMap("test", 3));
         assertEquals(5, script.execute(2));
         assertEquals(2, script.execute(-1));
+        assertEquals(true, factory.needsTest());
+        assertEquals(false, factory.needsNothing());
     }
 
     public abstract static class EmptyTestScript {

modules/lang-painless/src/test/java/org/elasticsearch/painless/ScriptTestCase.java

@@ -20,7 +20,6 @@
 package org.elasticsearch.painless;
 
 import junit.framework.AssertionFailedError;
-
 import org.apache.lucene.search.Scorer;
 import org.elasticsearch.common.lucene.ScorerAware;
 import org.elasticsearch.common.settings.Settings;
@@ -37,6 +36,7 @@ import java.util.Collection;
 import java.util.HashMap;
 import java.util.Map;
 
+import static org.elasticsearch.painless.node.SSource.MainMethodReserved;
 import static org.hamcrest.Matchers.hasSize;
 
 /**
@@ -96,8 +96,7 @@ public abstract class ScriptTestCase extends ESTestCase {
             CompilerSettings pickySettings = new CompilerSettings();
             pickySettings.setPicky(true);
             pickySettings.setRegexesEnabled(CompilerSettings.REGEX_ENABLED.get(scriptEngineSettings()));
-            Walker.buildPainlessTree(scriptClassInfo, getTestName(), script, pickySettings,
-                    definition, null);
+            Walker.buildPainlessTree(scriptClassInfo, new MainMethodReserved(), getTestName(), script, pickySettings, definition, null);
         }
         // test actual script execution
         ExecutableScript.Factory factory = scriptEngine.compile(null, script, ExecutableScript.CONTEXT, compileParams);

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

@@ -31,8 +31,8 @@ import org.elasticsearch.painless.FeatureTest;
 import org.elasticsearch.painless.GenericElasticsearchScript;
 import org.elasticsearch.painless.Locals.Variable;
 import org.elasticsearch.painless.Location;
-import org.elasticsearch.painless.ScriptClassInfo;
 import org.elasticsearch.painless.Operation;
+import org.elasticsearch.painless.ScriptClassInfo;
 import org.elasticsearch.painless.antlr.Walker;
 import org.elasticsearch.test.ESTestCase;
 
@@ -42,6 +42,7 @@ import java.util.Map;
 
 import static java.util.Collections.emptyList;
 import static java.util.Collections.singletonList;
+import static org.elasticsearch.painless.node.SSource.MainMethodReserved;
 
 /**
  * Tests {@link Object#toString} implementations on all extensions of {@link ANode}.
@@ -902,8 +903,8 @@ public class NodeToStringTests extends ESTestCase {
         CompilerSettings compilerSettings = new CompilerSettings();
         compilerSettings.setRegexesEnabled(true);
         try {
-            return Walker.buildPainlessTree(scriptClassInfo, getTestName(), code, compilerSettings,
-                    definition, null);
+            return Walker.buildPainlessTree(
+                scriptClassInfo, new MainMethodReserved(), getTestName(), code, compilerSettings, definition, null);
         } catch (Exception e) {
             throw new AssertionError("Failed to compile: " + code, e);
         }