Remove "extractVariables" phase from Painless user tree

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());

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

@@ -20,9 +20,9 @@
 package org.elasticsearch.painless.node;
 
 import org.elasticsearch.painless.FunctionRef;
+import org.elasticsearch.painless.Location;
 import org.elasticsearch.painless.Scope;
 import org.elasticsearch.painless.Scope.Variable;
-import org.elasticsearch.painless.Location;
 import org.elasticsearch.painless.ir.CapturingFuncRefNode;
 import org.elasticsearch.painless.ir.ClassNode;
 import org.elasticsearch.painless.lookup.def;
@@ -31,7 +31,6 @@
 import java.util.Collections;
 import java.util.List;
 import java.util.Objects;
-import java.util.Set;
 
 /**
  * Represents a capturing function reference.  For member functions that require a this reference, ie not static.
@@ -51,11 +50,6 @@ public ECapturingFunctionRef(Location location, String variable, String call) {
         this.call = Objects.requireNonNull(call);
     }
 
-    @Override
-    void extractVariables(Set<String> variables) {
-        variables.add(variable);
-    }
-
     @Override
     void analyze(ScriptRoot scriptRoot, Scope scope) {
         captured = scope.getVariable(location, variable);