Decouple Painless AST Lambda generation from the grammar (#45111)

This is the first step in decoupling the Painless AST from the grammar. The 
Painless AST should be able to generate classes independently of how the 
AST is generated from a grammar. (If I were to build a Painless AST by hand  
in code this should be all that's necessary.) This change removes Lambda 
name generation from the ANTLR grammar tree walker. It also removes 
unnecessary node generation of new array function references from the 
tree walker as well.

Author: jdconrad

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

@@ -38,6 +38,19 @@
  * Tracks user defined methods and variables across compilation phases.
  */
 public final class Locals {
+    private int syntheticCounter = 0;
+
+    /**
+     * Returns a unique identifier for generating the name of a synthetic method.
+     */
+    public String getNextSyntheticName() {
+        Locals locals = this;
+        while (locals.getParent() != null) {
+            locals = locals.getParent();
+        }
+
+        return "lambda$" + locals.syntheticCounter++;
+    }
 
     /**
      * Constructs a local method key used to lookup local methods from a painless class.

modules/lang-painless/src/main/java/org/elasticsearch/painless/action/PainlessContextAction.java

@@ -19,11 +19,11 @@
 
 package org.elasticsearch.painless.action;
 
-import org.elasticsearch.action.ActionType;
 import org.elasticsearch.action.ActionListener;
 import org.elasticsearch.action.ActionRequest;
 import org.elasticsearch.action.ActionRequestValidationException;
 import org.elasticsearch.action.ActionResponse;
+import org.elasticsearch.action.ActionType;
 import org.elasticsearch.action.support.ActionFilters;
 import org.elasticsearch.action.support.HandledTransportAction;
 import org.elasticsearch.client.node.NodeClient;

modules/lang-painless/src/main/java/org/elasticsearch/painless/action/PainlessExecuteAction.java

@@ -30,9 +30,9 @@
 import org.apache.lucene.search.Scorer;
 import org.apache.lucene.search.Weight;
 import org.apache.lucene.store.RAMDirectory;
-import org.elasticsearch.action.ActionType;
 import org.elasticsearch.action.ActionRequestValidationException;
 import org.elasticsearch.action.ActionResponse;
+import org.elasticsearch.action.ActionType;
 import org.elasticsearch.action.support.ActionFilters;
 import org.elasticsearch.action.support.IndicesOptions;
 import org.elasticsearch.action.support.single.shard.SingleShardRequest;

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

@@ -29,7 +29,6 @@
 import org.antlr.v4.runtime.atn.PredictionMode;
 import org.antlr.v4.runtime.tree.TerminalNode;
 import org.elasticsearch.painless.CompilerSettings;
-import org.elasticsearch.painless.Globals;
 import org.elasticsearch.painless.Location;
 import org.elasticsearch.painless.Operation;
 import org.elasticsearch.painless.ScriptClassInfo;
@@ -127,6 +126,7 @@
 import org.elasticsearch.painless.node.EListInit;
 import org.elasticsearch.painless.node.EMapInit;
 import org.elasticsearch.painless.node.ENewArray;
+import org.elasticsearch.painless.node.ENewArrayFunctionRef;
 import org.elasticsearch.painless.node.ENewObj;
 import org.elasticsearch.painless.node.ENull;
 import org.elasticsearch.painless.node.ENumeric;
@@ -163,8 +163,6 @@
 
 import java.util.ArrayDeque;
 import java.util.ArrayList;
-import java.util.Arrays;
-import java.util.BitSet;
 import java.util.Deque;
 import java.util.List;
 
@@ -184,11 +182,10 @@ public static SSource buildPainlessTree(ScriptClassInfo mainMethod, MainMethodRe
     private final CompilerSettings settings;
     private final Printer debugStream;
     private final String sourceName;
+    private final String sourceText;
     private final PainlessLookup painlessLookup;
 
     private final Deque<Reserved> reserved = new ArrayDeque<>();
-    private final Globals globals;
-    private int syntheticCounter = 0;
 
     private Walker(ScriptClassInfo scriptClassInfo, MainMethodReserved reserved, String sourceName, String sourceText,
                    CompilerSettings settings, PainlessLookup painlessLookup, Printer debugStream) {
@@ -197,7 +194,7 @@ private Walker(ScriptClassInfo scriptClassInfo, MainMethodReserved reserved, Str
         this.debugStream = debugStream;
         this.settings = settings;
         this.sourceName = Location.computeSourceName(sourceName);
-        this.globals = new Globals(new BitSet(sourceText.length()));
+        this.sourceText = sourceText;
         this.painlessLookup = painlessLookup;
         this.source = (SSource)visit(buildAntlrTree(sourceText));
     }
@@ -242,11 +239,6 @@ private Location location(ParserRuleContext ctx) {
         return new Location(sourceName, ctx.getStart().getStartIndex());
     }
 
-    /** Returns name of next lambda */
-    private String nextLambda() {
-        return "lambda$" + syntheticCounter++;
-    }
-
     @Override
     public ANode visitSource(SourceContext ctx) {
         List<SFunction> functions = new ArrayList<>();
@@ -261,8 +253,8 @@ public ANode visitSource(SourceContext ctx) {
             statements.add((AStatement)visit(statement));
         }
 
-        return new SSource(scriptClassInfo, settings, sourceName, debugStream, (MainMethodReserved)reserved.pop(),
-                           location(ctx), functions, globals, statements);
+        return new SSource(scriptClassInfo, settings, sourceName, sourceText, debugStream,
+                           (MainMethodReserved)reserved.pop(), location(ctx), functions, statements);
     }
 
     @Override
@@ -1099,8 +1091,7 @@ public ANode visitLambda(LambdaContext ctx) {
         FunctionReserved lambdaReserved = (FunctionReserved)reserved.pop();
         reserved.peek().addUsedVariables(lambdaReserved);
 
-        String name = nextLambda();
-        return new ELambda(name, lambdaReserved, location(ctx), paramTypes, paramNames, statements);
+        return new ELambda(lambdaReserved, location(ctx), paramTypes, paramNames, statements);
     }
 
     @Override
@@ -1115,22 +1106,9 @@ public ANode visitClassfuncref(ClassfuncrefContext ctx) {
 
     @Override
     public ANode visitConstructorfuncref(ConstructorfuncrefContext ctx) {
-        if (!ctx.decltype().LBRACE().isEmpty()) {
-            // array constructors are special: we need to make a synthetic method
-            // taking integer as argument and returning a new instance, and return a ref to that.
-            Location location = location(ctx);
-            String arrayType = ctx.decltype().getText();
-            SReturn code = new SReturn(location,
-                new ENewArray(location, arrayType, Arrays.asList(
-                    new EVariable(location, "size")), false));
-            String name = nextLambda();
-            globals.addSyntheticMethod(new SFunction(new FunctionReserved(), location, arrayType, name,
-                Arrays.asList("int"), Arrays.asList("size"), Arrays.asList(code), true));
-
-            return new EFunctionRef(location(ctx), "this", name);
-        }
-
-        return new EFunctionRef(location(ctx), ctx.decltype().getText(), ctx.NEW().getText());
+        return ctx.decltype().LBRACE().isEmpty() ?
+                new EFunctionRef(location(ctx), ctx.decltype().getText(), ctx.NEW().getText()) :
+                new ENewArrayFunctionRef(location(ctx), ctx.decltype().getText());
     }
 
     @Override

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

@@ -63,7 +63,6 @@
  */
 public final class ELambda extends AExpression implements ILambda {
 
-    private final String name;
     private final FunctionReserved reserved;
     private final List<String> paramTypeStrs;
     private final List<String> paramNameStrs;
@@ -78,11 +77,10 @@ public final class ELambda extends AExpression implements ILambda {
     // dynamic parent, deferred until link time
     private String defPointer;
 
-    public ELambda(String name, FunctionReserved reserved,
-                   Location location, List<String> paramTypes, List<String> paramNames,
+    public ELambda(FunctionReserved reserved, Location location,
+                   List<String> paramTypes, List<String> paramNames,
                    List<AStatement> statements) {
         super(location);
-        this.name = Objects.requireNonNull(name);
         this.reserved = Objects.requireNonNull(reserved);
         this.paramTypeStrs = Collections.unmodifiableList(paramTypes);
         this.paramNameStrs = Collections.unmodifiableList(paramNames);
@@ -167,6 +165,7 @@ void analyze(Locals locals) {
         paramNames.addAll(paramNameStrs);
 
         // desugar lambda body into a synthetic method
+        String name = locals.getNextSyntheticName();
         desugared = new SFunction(reserved, location, PainlessLookupUtility.typeToCanonicalTypeName(returnType), name,
                                   paramTypes, paramNames, statements, true);
         desugared.generateSignature(locals.getPainlessLookup());

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

@@ -0,0 +1,99 @@
+/*
+ * Licensed to Elasticsearch under one or more contributor
+ * license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright
+ * ownership. Elasticsearch licenses this file to you under
+ * the Apache License, Version 2.0 (the "License"); you may
+ * not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.elasticsearch.painless.node;
+
+import org.elasticsearch.painless.FunctionRef;
+import org.elasticsearch.painless.Globals;
+import org.elasticsearch.painless.Locals;
+import org.elasticsearch.painless.Location;
+import org.elasticsearch.painless.MethodWriter;
+import org.objectweb.asm.Type;
+
+import java.util.Arrays;
+import java.util.Objects;
+import java.util.Set;
+
+/**
+ * Represents a function reference.
+ */
+public final class ENewArrayFunctionRef extends AExpression implements ILambda {
+    private final String type;
+
+    private SFunction function;
+    private FunctionRef ref;
+    private String defPointer;
+
+    public ENewArrayFunctionRef(Location location, String type) {
+        super(location);
+
+        this.type = Objects.requireNonNull(type);
+    }
+
+    @Override
+    void extractVariables(Set<String> variables) {}
+
+    @Override
+    void analyze(Locals locals) {
+        SReturn code = new SReturn(location, new ENewArray(location, type, Arrays.asList(new EVariable(location, "size")), false));
+        function = new SFunction(new SFunction.FunctionReserved(), location, type, locals.getNextSyntheticName(),
+                Arrays.asList("int"), Arrays.asList("size"), Arrays.asList(code), true);
+        function.generateSignature(locals.getPainlessLookup());
+        function.analyze(Locals.newLambdaScope(locals.getProgramScope(), function.name, function.returnType,
+                function.parameters, 0, 0));
+
+        if (expected == null) {
+            ref = null;
+            actual = String.class;
+            defPointer = "Sthis." + function.name + ",0";
+        } else {
+            defPointer = null;
+            ref = FunctionRef.create(locals.getPainlessLookup(), locals.getMethods(), location, expected, "this", function.name, 0);
+            actual = expected;
+        }
+    }
+
+    @Override
+    void write(MethodWriter writer, Globals globals) {
+        if (ref != null) {
+            writer.writeDebugInfo(location);
+            writer.invokeLambdaCall(ref);
+        } else {
+            // push a null instruction as a placeholder for future lambda instructions
+            writer.push((String)null);
+        }
+
+        globals.addSyntheticMethod(function);
+    }
+
+    @Override
+    public String getPointer() {
+        return defPointer;
+    }
+
+    @Override
+    public Type[] getCaptures() {
+        return new Type[0]; // no captures
+    }
+
+    @Override
+    public String toString() {
+        return singleLineToString(type + "[]", "new");
+    }
+}

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

@@ -140,20 +140,17 @@ public int getMaxLoopCounter() {
     private final List<org.objectweb.asm.commons.Method> getMethods;
     private byte[] bytes;
 
-    public SSource(ScriptClassInfo scriptClassInfo, CompilerSettings settings, String name, Printer debugStream,
-            MainMethodReserved reserved, Location location, List<SFunction> functions, Globals globals, List<AStatement> statements) {
+    public SSource(ScriptClassInfo scriptClassInfo, CompilerSettings settings, String name, String sourceText, Printer debugStream,
+            MainMethodReserved reserved, Location location, List<SFunction> functions, List<AStatement> statements) {
         super(location);
         this.scriptClassInfo = Objects.requireNonNull(scriptClassInfo);
         this.settings = Objects.requireNonNull(settings);
         this.name = Objects.requireNonNull(name);
         this.debugStream = debugStream;
         this.reserved = Objects.requireNonNull(reserved);
-        // process any synthetic functions generated by walker (because right now, thats still easy)
-        functions.addAll(globals.getSyntheticMethods().values());
-        globals.getSyntheticMethods().clear();
         this.functions = Collections.unmodifiableList(functions);
         this.statements = Collections.unmodifiableList(statements);
-        this.globals = globals;
+        this.globals = new Globals(new BitSet(sourceText.length()));
 
         this.getMethods = new ArrayList<>();
     }