Remove "extractVariables" phase from Painless user tree (#51690)

Removes the "extractVariables" phase from the Painless user tree which is 
no longer necessary. The information to retrieve used variables is now 
collected during the semantic ("analysis") phase and passed back through 
ScriptRoot to generate the appropriate needs methods in the factories.

Author: jdconrad

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

@@ -36,9 +36,7 @@
 import java.security.cert.Certificate;
 import java.util.Collections;
 import java.util.HashMap;
-import java.util.HashSet;
 import java.util.Map;
-import java.util.Set;
 import java.util.concurrent.atomic.AtomicInteger;
 
 import static org.elasticsearch.painless.WriterConstants.CLASS_NAME;
@@ -208,11 +206,9 @@ private static void addFactoryMethod(Map<String, Class<?>> additionalClasses, Cl
      * @param settings The CompilerSettings to be used during the compilation.
      * @return The ScriptRoot used to compile
      */
-    ScriptRoot compile(Loader loader, Set<String> extractedVariables, String name, String source,
-            CompilerSettings settings) {
+    ScriptRoot compile(Loader loader, String name, String source, CompilerSettings settings) {
         ScriptClassInfo scriptClassInfo = new ScriptClassInfo(painlessLookup, scriptClass);
         SClass root = Walker.buildPainlessTree(scriptClassInfo, name, source, settings, painlessLookup, null);
-        root.extractVariables(extractedVariables);
         ScriptRoot scriptRoot = root.analyze(painlessLookup, settings);
         ClassNode classNode = root.writeClass();
         Map<String, Object> statics = classNode.write();
@@ -243,9 +239,7 @@ ScriptRoot compile(Loader loader, Set<String> extractedVariables, String name, S
      */
     byte[] compile(String name, String source, CompilerSettings settings, Printer debugStream) {
         ScriptClassInfo scriptClassInfo = new ScriptClassInfo(painlessLookup, scriptClass);
-        SClass root = Walker.buildPainlessTree(scriptClassInfo, name, source, settings, painlessLookup,
-                debugStream);
-        root.extractVariables(new HashSet<>());
+        SClass root = Walker.buildPainlessTree(scriptClassInfo, name, source, settings, painlessLookup, debugStream);
         root.analyze(painlessLookup, settings);
         ClassNode classNode = root.writeClass();
         classNode.write();

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

@@ -45,7 +45,6 @@
 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.Set;
@@ -142,14 +141,12 @@ public Loader run() {
             }
         });
 
-        Set<String> extractedVariables = new HashSet<>();
-        ScriptRoot scriptRoot = compile(contextsToCompilers.get(context), loader, extractedVariables, scriptName, scriptSource, params);
+        ScriptRoot scriptRoot = compile(contextsToCompilers.get(context), loader, scriptName, scriptSource, params);
 
         if (context.statefulFactoryClazz != null) {
-            return generateFactory(loader, context, extractedVariables, generateStatefulFactory(loader, context, extractedVariables),
-                scriptRoot);
+            return generateFactory(loader, context, generateStatefulFactory(loader, context, scriptRoot), scriptRoot);
         } else {
-            return generateFactory(loader, context, extractedVariables, WriterConstants.CLASS_TYPE, scriptRoot);
+            return generateFactory(loader, context, WriterConstants.CLASS_TYPE, scriptRoot);
         }
     }
 
@@ -172,7 +169,7 @@ public Set<ScriptContext<?>> getSupportedContexts() {
     private <T> Type generateStatefulFactory(
         Loader loader,
         ScriptContext<T> context,
-        Set<String> extractedVariables
+        ScriptRoot scriptRoot
     ) {
         int classFrames = ClassWriter.COMPUTE_FRAMES | ClassWriter.COMPUTE_MAXS;
         int classAccess = Opcodes.ACC_PUBLIC | Opcodes.ACC_SUPER | Opcodes.ACC_FINAL;
@@ -254,7 +251,7 @@ private <T> Type generateStatefulFactory(
         adapter.returnValue();
         adapter.endMethod();
 
-        writeNeedsMethods(context.statefulFactoryClazz, writer, extractedVariables);
+        writeNeedsMethods(context.statefulFactoryClazz, writer, scriptRoot.getUsedVariables());
         writer.visitEnd();
 
         loader.defineFactory(className.replace('/', '.'), writer.toByteArray());
@@ -278,7 +275,6 @@ private <T> Type generateStatefulFactory(
     private <T> T generateFactory(
         Loader loader,
         ScriptContext<T> context,
-        Set<String> extractedVariables,
         Type classType,
         ScriptRoot scriptRoot
     ) {
@@ -332,7 +328,7 @@ private <T> T generateFactory(
         adapter.returnValue();
         adapter.endMethod();
 
-        writeNeedsMethods(context.factoryClazz, writer, extractedVariables);
+        writeNeedsMethods(context.factoryClazz, writer, scriptRoot.getUsedVariables());
 
         String methodName = "isResultDeterministic";
         org.objectweb.asm.commons.Method isResultDeterministic = new org.objectweb.asm.commons.Method(methodName,
@@ -378,8 +374,7 @@ private void writeNeedsMethods(Class<?> clazz, ClassWriter writer, Set<String> e
         }
     }
 
-    ScriptRoot compile(Compiler compiler, Loader loader, Set<String> extractedVariables,
-                 String scriptName, String source, Map<String, String> params) {
+    ScriptRoot compile(Compiler compiler, Loader loader, String scriptName, String source, Map<String, String> params) {
         final CompilerSettings compilerSettings = buildCompilerSettings(params);
 
         try {
@@ -388,7 +383,7 @@ ScriptRoot compile(Compiler compiler, Loader loader, Set<String> extractedVariab
                 @Override
                 public ScriptRoot run() {
                     String name = scriptName == null ? source : scriptName;
-                    return compiler.compile(loader, extractedVariables, name, source, compilerSettings);
+                    return compiler.compile(loader, name, source, compilerSettings);
                 }
             }, COMPILATION_CONTEXT);
             // Note that it is safe to catch any of the following errors since Painless is stateless.

modules/lang-painless/src/main/java/org/elasticsearch/painless/ir/ClassNode.java

@@ -41,10 +41,8 @@
 import java.util.BitSet;
 import java.util.Collection;
 import java.util.HashMap;
-import java.util.HashSet;
 import java.util.List;
 import java.util.Map;
-import java.util.Set;
 
 import static org.elasticsearch.painless.WriterConstants.BASE_INTERFACE_TYPE;
 import static org.elasticsearch.painless.WriterConstants.BITSET_TYPE;
@@ -109,7 +107,6 @@ public List<StatementNode> getStatementsNodes() {
     private Printer debugStream;
     private ScriptRoot scriptRoot;
     private boolean doesMethodEscape;
-    private final Set<String> extractedVariables = new HashSet<>();
 
     public void setScriptClassInfo(ScriptClassInfo scriptClassInfo) {
         this.scriptClassInfo = scriptClassInfo;
@@ -159,18 +156,6 @@ public boolean doesMethodEscape() {
         return doesMethodEscape;
     }
 
-    public void addExtractedVariable(String extractedVariable) {
-        extractedVariables.add(extractedVariable);
-    }
-
-    public boolean containsExtractedVariable(String extractedVariable) {
-        return extractedVariables.contains(extractedVariable);
-    }
-
-    public Set<String> getExtractedVariables() {
-        return extractedVariables;
-    }
-
     /* ---- end node data ---- */
 
     protected Globals globals;
@@ -299,7 +284,7 @@ public Map<String, Object> write() {
             name = Character.toLowerCase(name.charAt(0)) + name.substring(1);
             MethodWriter ifaceMethod = classWriter.newMethodWriter(Opcodes.ACC_PUBLIC, needsMethod);
             ifaceMethod.visitCode();
-            ifaceMethod.push(extractedVariables.contains(name));
+            ifaceMethod.push(scriptRoot.getUsedVariables().contains(name));
             ifaceMethod.returnValue();
             ifaceMethod.endMethod();
         }
@@ -355,7 +340,7 @@ protected void write(ClassWriter classWriter, MethodWriter methodWriter, Globals
             String name = method.getName().substring(3);
             name = Character.toLowerCase(name.charAt(0)) + name.substring(1);
 
-            if (extractedVariables.contains(name)) {
+            if (scriptRoot.getUsedVariables().contains(name)) {
                 Variable variable = scopeTable.defineVariable(returnType, name);
 
                 methodWriter.loadThis();

modules/lang-painless/src/main/java/org/elasticsearch/painless/node/AExpression.java

@@ -20,8 +20,8 @@
 package org.elasticsearch.painless.node;
 
 import org.elasticsearch.painless.AnalyzerCaster;
-import org.elasticsearch.painless.Scope;
 import org.elasticsearch.painless.Location;
+import org.elasticsearch.painless.Scope;
 import org.elasticsearch.painless.ir.ClassNode;
 import org.elasticsearch.painless.ir.ExpressionNode;
 import org.elasticsearch.painless.lookup.PainlessCast;

modules/lang-painless/src/main/java/org/elasticsearch/painless/node/ANode.java

@@ -30,7 +30,6 @@
 import java.util.Iterator;
 import java.util.List;
 import java.util.Objects;
-import java.util.Set;
 
 import static java.lang.Math.max;
 import static java.util.Collections.emptyList;
@@ -52,14 +51,6 @@ public abstract class ANode {
         this.location = Objects.requireNonNull(location);
     }
 
-    /**
-     * Adds all variable names referenced to the variable set.
-     * <p>
-     * This can be called at any time, e.g. to support lambda capture.
-     * @param variables set of variables referenced (any scope)
-     */
-    abstract void extractVariables(Set<String> variables);
-
     /**
      * Writes ASM based on the data collected during the analysis phase.
      * @param classNode the root {@link ClassNode}

modules/lang-painless/src/main/java/org/elasticsearch/painless/node/EAssignment.java

@@ -21,9 +21,9 @@
 
 
 import org.elasticsearch.painless.AnalyzerCaster;
-import org.elasticsearch.painless.Scope;
 import org.elasticsearch.painless.Location;
 import org.elasticsearch.painless.Operation;
+import org.elasticsearch.painless.Scope;
 import org.elasticsearch.painless.ir.AssignmentNode;
 import org.elasticsearch.painless.ir.ClassNode;
 import org.elasticsearch.painless.lookup.PainlessCast;
@@ -33,7 +33,6 @@
 import java.util.ArrayList;
 import java.util.List;
 import java.util.Objects;
-import java.util.Set;
 
 /**
  * Represents an assignment with the lhs and rhs as child nodes.
@@ -62,15 +61,6 @@ public EAssignment(Location location, AExpression lhs, AExpression rhs, boolean
         this.operation = operation;
     }
 
-    @Override
-    void extractVariables(Set<String> variables) {
-        lhs.extractVariables(variables);
-
-        if (rhs != null) {
-            rhs.extractVariables(variables);
-        }
-    }
-
     @Override
     void analyze(ScriptRoot scriptRoot, Scope scope) {
         analyzeLHS(scriptRoot, scope);

modules/lang-painless/src/main/java/org/elasticsearch/painless/node/EBinary.java

@@ -20,17 +20,16 @@
 package org.elasticsearch.painless.node;
 
 import org.elasticsearch.painless.AnalyzerCaster;
-import org.elasticsearch.painless.Scope;
 import org.elasticsearch.painless.Location;
 import org.elasticsearch.painless.Operation;
+import org.elasticsearch.painless.Scope;
 import org.elasticsearch.painless.ir.BinaryMathNode;
 import org.elasticsearch.painless.ir.ClassNode;
 import org.elasticsearch.painless.lookup.PainlessLookupUtility;
 import org.elasticsearch.painless.lookup.def;
 import org.elasticsearch.painless.symbol.ScriptRoot;
 
 import java.util.Objects;
-import java.util.Set;
 import java.util.regex.Pattern;
 
 /**
@@ -55,12 +54,6 @@ public EBinary(Location location, Operation operation, AExpression left, AExpres
         this.right = Objects.requireNonNull(right);
     }
 
-    @Override
-    void extractVariables(Set<String> variables) {
-        left.extractVariables(variables);
-        right.extractVariables(variables);
-    }
-
     @Override
     void analyze(ScriptRoot scriptRoot, Scope scope) {
         originallyExplicit = explicit;

modules/lang-painless/src/main/java/org/elasticsearch/painless/node/EBool.java

@@ -19,15 +19,14 @@
 
 package org.elasticsearch.painless.node;
 
-import org.elasticsearch.painless.Scope;
 import org.elasticsearch.painless.Location;
 import org.elasticsearch.painless.Operation;
+import org.elasticsearch.painless.Scope;
 import org.elasticsearch.painless.ir.BooleanNode;
 import org.elasticsearch.painless.ir.ClassNode;
 import org.elasticsearch.painless.symbol.ScriptRoot;
 
 import java.util.Objects;
-import java.util.Set;
 
 /**
  * Represents a boolean expression.
@@ -46,12 +45,6 @@ public EBool(Location location, Operation operation, AExpression left, AExpressi
         this.right = Objects.requireNonNull(right);
     }
 
-    @Override
-    void extractVariables(Set<String> variables) {
-        left.extractVariables(variables);
-        right.extractVariables(variables);
-    }
-
     @Override
     void analyze(ScriptRoot scriptRoot, Scope scope) {
         left.expected = boolean.class;

modules/lang-painless/src/main/java/org/elasticsearch/painless/node/EBoolean.java

@@ -19,14 +19,12 @@
 
 package org.elasticsearch.painless.node;
 
-import org.elasticsearch.painless.Scope;
 import org.elasticsearch.painless.Location;
+import org.elasticsearch.painless.Scope;
 import org.elasticsearch.painless.ir.ClassNode;
 import org.elasticsearch.painless.ir.ExpressionNode;
 import org.elasticsearch.painless.symbol.ScriptRoot;
 
-import java.util.Set;
-
 /**
  * Represents a boolean constant.
  */
@@ -38,11 +36,6 @@ public EBoolean(Location location, boolean constant) {
         this.constant = constant;
     }
 
-    @Override
-    void extractVariables(Set<String> variables) {
-        // Do nothing.
-    }
-
     @Override
     void analyze(ScriptRoot scriptRoot, Scope scope) {
         if (!read) {

modules/lang-painless/src/main/java/org/elasticsearch/painless/node/ECallLocal.java

@@ -19,8 +19,8 @@
 
 package org.elasticsearch.painless.node;
 
-import org.elasticsearch.painless.Scope;
 import org.elasticsearch.painless.Location;
+import org.elasticsearch.painless.Scope;
 import org.elasticsearch.painless.ir.ClassNode;
 import org.elasticsearch.painless.ir.UnboundCallNode;
 import org.elasticsearch.painless.lookup.PainlessClassBinding;
@@ -34,7 +34,6 @@
 import java.util.ArrayList;
 import java.util.List;
 import java.util.Objects;
-import java.util.Set;
 
 /**
  * Represents a user-defined call.
@@ -58,13 +57,6 @@ public ECallLocal(Location location, String name, List<AExpression> arguments) {
         this.arguments = Objects.requireNonNull(arguments);
     }
 
-    @Override
-    void extractVariables(Set<String> variables) {
-        for (AExpression argument : arguments) {
-            argument.extractVariables(variables);
-        }
-    }
-
     @Override
     void analyze(ScriptRoot scriptRoot, Scope scope) {
         localFunction = scriptRoot.getFunctionTable().getFunction(name, arguments.size());