Merge pull request #18831 from rmuir/moreRefs

painless: array constructor references

modules/lang-painless/src/main/antlr/PainlessParser.g4

@@ -195,6 +195,30 @@ lamtype
     ;
 
 funcref
-    : TYPE REF ( ID | NEW )
-    | ( ID | THIS ) REF ID
+    : classFuncref
+    | constructorFuncref
+    | capturingFuncref
+    | localFuncref
+    ;
+
+// reference to a static or instance method, e.g. ArrayList::size or Integer::compare
+classFuncref
+    : TYPE REF ID
+    ;
+
+// reference to a constructor, e.g. ArrayList::new
+// currently limited to simple non-array types
+constructorFuncref
+    : decltype REF NEW
+    ;
+
+// reference to an instance method, e.g. object::toString
+// currently limited to capture of a simple variable (id).
+capturingFuncref
+    : ID REF ID
+    ;
+
+// reference to a local function, e.g. this::myfunc
+localFuncref
+    : THIS REF ID
     ;

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

@@ -417,4 +417,32 @@ class PainlessParserBaseVisitor<T> extends AbstractParseTreeVisitor<T> implement
    * {@link #visitChildren} on {@code ctx}.</p>
    */
   @Override public T visitFuncref(PainlessParser.FuncrefContext ctx) { return visitChildren(ctx); }
+  /**
+   * {@inheritDoc}
+   *
+   * <p>The default implementation returns the result of calling
+   * {@link #visitChildren} on {@code ctx}.</p>
+   */
+  @Override public T visitClassFuncref(PainlessParser.ClassFuncrefContext ctx) { return visitChildren(ctx); }
+  /**
+   * {@inheritDoc}
+   *
+   * <p>The default implementation returns the result of calling
+   * {@link #visitChildren} on {@code ctx}.</p>
+   */
+  @Override public T visitConstructorFuncref(PainlessParser.ConstructorFuncrefContext ctx) { return visitChildren(ctx); }
+  /**
+   * {@inheritDoc}
+   *
+   * <p>The default implementation returns the result of calling
+   * {@link #visitChildren} on {@code ctx}.</p>
+   */
+  @Override public T visitCapturingFuncref(PainlessParser.CapturingFuncrefContext ctx) { return visitChildren(ctx); }
+  /**
+   * {@inheritDoc}
+   *
+   * <p>The default implementation returns the result of calling
+   * {@link #visitChildren} on {@code ctx}.</p>
+   */
+  @Override public T visitLocalFuncref(PainlessParser.LocalFuncrefContext ctx) { return visitChildren(ctx); }
 }

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

@@ -397,4 +397,28 @@ interface PainlessParserVisitor<T> extends ParseTreeVisitor<T> {
    * @return the visitor result
    */
   T visitFuncref(PainlessParser.FuncrefContext ctx);
+  /**
+   * Visit a parse tree produced by {@link PainlessParser#classFuncref}.
+   * @param ctx the parse tree
+   * @return the visitor result
+   */
+  T visitClassFuncref(PainlessParser.ClassFuncrefContext ctx);
+  /**
+   * Visit a parse tree produced by {@link PainlessParser#constructorFuncref}.
+   * @param ctx the parse tree
+   * @return the visitor result
+   */
+  T visitConstructorFuncref(PainlessParser.ConstructorFuncrefContext ctx);
+  /**
+   * Visit a parse tree produced by {@link PainlessParser#capturingFuncref}.
+   * @param ctx the parse tree
+   * @return the visitor result
+   */
+  T visitCapturingFuncref(PainlessParser.CapturingFuncrefContext ctx);
+  /**
+   * Visit a parse tree produced by {@link PainlessParser#localFuncref}.
+   * @param ctx the parse tree
+   * @return the visitor result
+   */
+  T visitLocalFuncref(PainlessParser.LocalFuncrefContext ctx);
 }

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

@@ -45,10 +45,13 @@ import org.elasticsearch.painless.antlr.PainlessParser.BraceaccessContext;
 import org.elasticsearch.painless.antlr.PainlessParser.BreakContext;
 import org.elasticsearch.painless.antlr.PainlessParser.CallinvokeContext;
 import org.elasticsearch.painless.antlr.PainlessParser.CalllocalContext;
+import org.elasticsearch.painless.antlr.PainlessParser.CapturingFuncrefContext;
 import org.elasticsearch.painless.antlr.PainlessParser.CastContext;
 import org.elasticsearch.painless.antlr.PainlessParser.ChainprecContext;
+import org.elasticsearch.painless.antlr.PainlessParser.ClassFuncrefContext;
 import org.elasticsearch.painless.antlr.PainlessParser.CompContext;
 import org.elasticsearch.painless.antlr.PainlessParser.ConditionalContext;
+import org.elasticsearch.painless.antlr.PainlessParser.ConstructorFuncrefContext;
 import org.elasticsearch.painless.antlr.PainlessParser.ContinueContext;
 import org.elasticsearch.painless.antlr.PainlessParser.DeclContext;
 import org.elasticsearch.painless.antlr.PainlessParser.DeclarationContext;
@@ -71,6 +74,7 @@ import org.elasticsearch.painless.antlr.PainlessParser.IfContext;
 import org.elasticsearch.painless.antlr.PainlessParser.InitializerContext;
 import org.elasticsearch.painless.antlr.PainlessParser.LambdaContext;
 import org.elasticsearch.painless.antlr.PainlessParser.LamtypeContext;
+import org.elasticsearch.painless.antlr.PainlessParser.LocalFuncrefContext;
 import org.elasticsearch.painless.antlr.PainlessParser.NewarrayContext;
 import org.elasticsearch.painless.antlr.PainlessParser.NewobjectContext;
 import org.elasticsearch.painless.antlr.PainlessParser.NullContext;
@@ -144,6 +148,7 @@ import org.elasticsearch.painless.node.SWhile;
 
 import java.util.ArrayDeque;
 import java.util.ArrayList;
+import java.util.Arrays;
 import java.util.Deque;
 import java.util.List;
 
@@ -162,6 +167,7 @@ public final class Walker extends PainlessParserBaseVisitor<Object> {
     private final String sourceText;
 
     private final Deque<Reserved> reserved = new ArrayDeque<>();
+    private final List<SFunction> synthetic = new ArrayList<>();
 
     private Walker(String sourceName, String sourceText, CompilerSettings settings) {
         this.settings = settings;
@@ -225,6 +231,8 @@ public final class Walker extends PainlessParserBaseVisitor<Object> {
         for (StatementContext statement : ctx.statement()) {
             statements.add((AStatement)visit(statement));
         }
+        
+        functions.addAll(synthetic);
 
         return new SSource(sourceName, sourceText, (ExecuteReserved)reserved.pop(), location(ctx), functions, statements);
     }
@@ -251,7 +259,7 @@ public final class Walker extends PainlessParserBaseVisitor<Object> {
             statements.add((AStatement)visit(statement));
         }
 
-        return new SFunction((FunctionReserved)reserved.pop(), location(ctx), rtnType, name, paramTypes, paramNames, statements);
+        return new SFunction((FunctionReserved)reserved.pop(), location(ctx), rtnType, name, paramTypes, paramNames, statements, false);
     }
 
     @Override
@@ -950,20 +958,51 @@ public final class Walker extends PainlessParserBaseVisitor<Object> {
 
     @Override
     public Object visitFuncref(FuncrefContext ctx) {
-        if (ctx.TYPE() != null) {
-            // non-capturing Type::method or Type::new
-            final String methodText;
-            if (ctx.NEW() != null) {
-                methodText = ctx.NEW().getText();
-            } else {
-                methodText = ctx.ID(0).getText();
-            }
-            return new EFunctionRef(location(ctx), ctx.TYPE().getText(), methodText);
-        } else if (ctx.THIS() != null) {
-            return new EFunctionRef(location(ctx), ctx.THIS().getText(), ctx.ID(0).getText());
+        if (ctx.classFuncref() != null) {
+            return visit(ctx.classFuncref());
+        } else if (ctx.constructorFuncref() != null) {
+            return visit(ctx.constructorFuncref());
+        } else if (ctx.capturingFuncref() != null) {
+            return visit(ctx.capturingFuncref());
+        } else if (ctx.localFuncref() != null) {
+            return visit(ctx.localFuncref());
         } else {
-            // capturing object::method
-            return new ECapturingFunctionRef(location(ctx), ctx.ID(0).getText(), ctx.ID(1).getText());
+            throw location(ctx).createError(new IllegalStateException("Illegal tree structure."));
+        }
+    }
+
+    @Override
+    public Object visitClassFuncref(ClassFuncrefContext ctx) {
+        return new EFunctionRef(location(ctx), ctx.TYPE().getText(), ctx.ID().getText());
+    }
+
+    @Override
+    public Object 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 EChain(location,
+                           new LNewArray(location, arrayType, Arrays.asList(
+                           new EChain(location, 
+                           new LVariable(location, "size"))))));
+            String name = "lambda$" + synthetic.size();
+            synthetic.add(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());
+    }
+
+    @Override
+    public Object visitCapturingFuncref(CapturingFuncrefContext ctx) {
+        return new ECapturingFunctionRef(location(ctx), ctx.ID(0).getText(), ctx.ID(1).getText());
+    }
+
+    @Override
+    public Object visitLocalFuncref(LocalFuncrefContext ctx) {
+        return new EFunctionRef(location(ctx), ctx.THIS().getText(), ctx.ID().getText());
     }
 }

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

@@ -29,6 +29,7 @@ import org.elasticsearch.painless.Operation;
 import org.elasticsearch.painless.Locals;
 import org.elasticsearch.painless.MethodWriter;
 
+import java.util.Arrays;
 import java.util.List;
 
 /**
@@ -46,6 +47,11 @@ public final class EChain extends AExpression {
     Type promote = null;
     Cast there = null;
     Cast back = null;
+    
+    /** Creates a new RHS-only EChain */
+    public EChain(Location location, ALink link) {
+        this(location, Arrays.asList(link), false, false, null, null);
+    }
 
     public EChain(Location location, List<ALink> links,
                   boolean pre, boolean post, Operation operation, AExpression expression) {

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

@@ -50,6 +50,7 @@ public class SFunction extends AStatement {
     final List<String> paramTypeStrs;
     final List<String> paramNameStrs;
     final List<AStatement> statements;
+    final boolean synthetic;
 
     Type rtnType = null;
     List<Parameter> parameters = new ArrayList<>();
@@ -58,7 +59,8 @@ public class SFunction extends AStatement {
     Locals locals = null;
 
     public SFunction(FunctionReserved reserved, Location location,
-                     String rtnType, String name, List<String> paramTypes, List<String> paramNames, List<AStatement> statements) {
+                     String rtnType, String name, List<String> paramTypes, 
+                     List<String> paramNames, List<AStatement> statements, boolean synthetic) {
         super(location);
 
         this.reserved = reserved;
@@ -67,6 +69,7 @@ public class SFunction extends AStatement {
         this.paramTypeStrs = Collections.unmodifiableList(paramTypes);
         this.paramNameStrs = Collections.unmodifiableList(paramNames);
         this.statements = Collections.unmodifiableList(statements);
+        this.synthetic = synthetic;
     }
 
     void generate() {
@@ -138,7 +141,11 @@ public class SFunction extends AStatement {
     
     /** Writes the function to given ClassWriter. */
     void write (ClassWriter writer, BitSet statements) {
-        final MethodWriter function = new MethodWriter(Opcodes.ACC_PRIVATE | Opcodes.ACC_STATIC, method.method, writer, statements);
+        int access = Opcodes.ACC_PRIVATE | Opcodes.ACC_STATIC;
+        if (synthetic) {
+            access |= Opcodes.ACC_SYNTHETIC;
+        }
+        final MethodWriter function = new MethodWriter(access, method.method, writer, statements);
         write(function);
         function.endMethod();
     }

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

@@ -70,6 +70,10 @@ public class ArrayTests extends ScriptTestCase {
         assertEquals(1, exec("int x; def y = new def[1]; x = y[0] = 1; return x;"));
     }
 
+    public void testArrayVariable() {
+        assertEquals(1, exec("int x = 1; int[] y = new int[x]; return y.length"));
+    }
+
     public void testForLoop() {
         assertEquals(999*1000/2, exec("def a = new int[1000]; for (int x = 0; x < a.length; x++) { a[x] = x; } "+
             "int total = 0; for (int x = 0; x < a.length; x++) { total += a[x]; } return total;"));

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

@@ -56,7 +56,21 @@ public class FunctionRefTests extends ScriptTestCase {
                                                   "DoubleSummaryStatistics::combine); " + 
                  "return stats.getSum()"));
     }
-    
+
+    public void testArrayCtorMethodRef() {
+        assertEquals(1.0D, 
+                exec("List l = new ArrayList(); l.add(1.0); l.add(2.0); " + 
+                     "def[] array = l.stream().toArray(Double[]::new);" + 
+                     "return array[0];"));
+    }
+
+    public void testArrayCtorMethodRefDef() {
+        assertEquals(1.0D, 
+                exec("def l = new ArrayList(); l.add(1.0); l.add(2.0); " + 
+                     "def[] array = l.stream().toArray(Double[]::new);" + 
+                     "return array[0];"));
+    }
+
     public void testCapturingMethodReference() {
         assertEquals("5", exec("Integer x = Integer.valueOf(5); return Optional.empty().orElseGet(x::toString);"));
         assertEquals("[]", exec("List l = new ArrayList(); return Optional.empty().orElseGet(l::toString);"));