Merge pull request #18818 from rmuir/capturingReferences

painless: add capturing method references

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

@@ -83,6 +83,7 @@ decltype
 
 funcref
     : TYPE REF ( ID | NEW )
+    | ID REF ID
     ;
 
 declvar

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

@@ -212,6 +212,8 @@ public final class Def {
      * until it finds a matching whitelisted method. If one is not found, it throws an exception.
      * Otherwise it returns a handle to the matching method.
      * <p>
+     * @param lookup caller's lookup
+     * @param callSiteType callsite's type
      * @param receiverClass Class of the object to invoke the method on.
      * @param name Name of the method.
      * @param args args passed to callsite
@@ -220,31 +222,103 @@ public final class Def {
      * @throws LambdaConversionException if a method reference cannot be converted to an functional interface
      * @throws IllegalArgumentException if no matching whitelisted method was found.
      */
-     static MethodHandle lookupMethod(Lookup lookup, Class<?> receiverClass, String name,
+     static MethodHandle lookupMethod(Lookup lookup, MethodType callSiteType, 
-             Object args[], long recipe) throws LambdaConversionException {
+             Class<?> receiverClass, String name, Object args[], long recipe) throws LambdaConversionException {
-         Method method = lookupMethodInternal(receiverClass, name, args.length - 1);
+         // simple case: no lambdas
+         if (recipe == 0) {
+             return lookupMethodInternal(receiverClass, name, args.length - 1).handle;
+         }
+         
+         // otherwise: first we have to compute the "real" arity. This is because we have extra arguments:
+         // e.g. f(a, g(x), b, h(y), i()) looks like f(a, g, x, b, h, y, i). 
+         int arity = args.length - 1;
+         for (int i = 0; i < args.length; i++) {
+             if ((recipe & (1L << (i - 1))) != 0) {
+                 String signature = (String) args[i];
+                 int numCaptures = Integer.parseInt(signature.substring(signature.indexOf(',')+1));
+                 arity -= numCaptures;
+             }
+         }
+         
+         // lookup the method with the proper arity, then we know everything (e.g. interface types of parameters).
+         // based on these we can finally link any remaining lambdas that were deferred.
+         Method method = lookupMethodInternal(receiverClass, name, arity);
          MethodHandle handle = method.handle;
 
-         if (recipe != 0) {
+         int replaced = 0;
-             MethodHandle filters[] = new MethodHandle[args.length];
+         for (int i = 1; i < args.length; i++) {
-             for (int i = 0; i < args.length; i++) {
              // its a functional reference, replace the argument with an impl
              if ((recipe & (1L << (i - 1))) != 0) {
-                     filters[i] = lookupReference(lookup, method.arguments.get(i - 1), (String) args[i]);
+                 // decode signature of form 'type.call,2' 
+                 String signature = (String) args[i];
+                 int separator = signature.indexOf('.');
+                 int separator2 = signature.indexOf(',');
+                 String type = signature.substring(1, separator);
+                 String call = signature.substring(separator+1, separator2);
+                 int numCaptures = Integer.parseInt(signature.substring(separator2+1));
+                 Class<?> captures[] = new Class<?>[numCaptures];
+                 for (int capture = 0; capture < captures.length; capture++) {
+                     captures[capture] = callSiteType.parameterType(i + 1 + capture);
                  }
+                 MethodHandle filter;
+                 Definition.Type interfaceType = method.arguments.get(i - 1 - replaced);
+                 if (signature.charAt(0) == 'S') {
+                     // the implementation is strongly typed, now that we know the interface type,
+                     // we have everything.
+                     filter = lookupReferenceInternal(lookup,
+                                                      interfaceType,
+                                                      type,
+                                                      call,
+                                                      captures);
+                 } else if (signature.charAt(0) == 'D') {
+                     // the interface type is now known, but we need to get the implementation.
+                     // this is dynamically based on the receiver type (and cached separately, underneath
+                     // this cache). It won't blow up since we never nest here (just references)
+                     MethodType nestedType = MethodType.methodType(interfaceType.clazz, captures);
+                     CallSite nested = DefBootstrap.bootstrap(lookup, 
+                                                              call,
+                                                              nestedType, 
+                                                              DefBootstrap.REFERENCE,
+                                                              interfaceType.name);
+                     filter = nested.dynamicInvoker();
+                 } else {
+                     throw new AssertionError();
+                 }
+                 // the filter now ignores the signature (placeholder) on the stack
+                 filter = MethodHandles.dropArguments(filter, 0, String.class);
+                 handle = MethodHandles.collectArguments(handle, i, filter);
+                 i += numCaptures;
+                 replaced += numCaptures;
              }
-             handle = MethodHandles.filterArguments(handle, 0, filters);
          }
          
          return handle;
      }
      
+     /**
+      * Returns an implementation of interfaceClass that calls receiverClass.name
+      * <p>
+      * This is just like LambdaMetaFactory, only with a dynamic type. The interface type is known,
+      * so we simply need to lookup the matching implementation method based on receiver type.
+      */
+     static MethodHandle lookupReference(Lookup lookup, String interfaceClass, 
+                                         Class<?> receiverClass, String name) throws LambdaConversionException {
+         Definition.Type interfaceType = Definition.getType(interfaceClass);
+         Method interfaceMethod = interfaceType.struct.getFunctionalMethod();
+         if (interfaceMethod == null) {
+             throw new IllegalArgumentException("Class [" + interfaceClass + "] is not a functional interface");
+         }
+         int arity = interfaceMethod.arguments.size();
+         Method implMethod = lookupMethodInternal(receiverClass, name, arity);
+         return lookupReferenceInternal(lookup, interfaceType, implMethod.owner.name, implMethod.name, receiverClass);
+     }
+     
      /** Returns a method handle to an implementation of clazz, given method reference signature 
       * @throws LambdaConversionException if a method reference cannot be converted to an functional interface
       */
-     private static MethodHandle lookupReference(Lookup lookup, Definition.Type clazz, String signature) throws LambdaConversionException {
+     private static MethodHandle lookupReferenceInternal(Lookup lookup, Definition.Type clazz, String type,
-         int separator = signature.indexOf('.');
+                                                         String call, Class<?>... captures) throws LambdaConversionException {
-         FunctionRef ref = new FunctionRef(clazz, signature.substring(0, separator), signature.substring(separator+1));
+         FunctionRef ref = new FunctionRef(clazz, type, call, captures);
          final CallSite callSite;
          if (ref.needsBridges()) {
              callSite = LambdaMetafactory.altMetafactory(lookup, 
@@ -265,9 +339,7 @@ public final class Def {
                      ref.samMethodType,
                      0);
          }
-         // we could actually invoke and cache here (in non-capturing cases), but this is not a speedup.
+         return callSite.dynamicInvoker().asType(MethodType.methodType(clazz.clazz, captures));
-         MethodHandle factory = callSite.dynamicInvoker().asType(MethodType.methodType(clazz.clazz));
-         return MethodHandles.dropArguments(factory, 0, Object.class);
      }
      
 

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

@@ -60,6 +60,8 @@ public final class DefBootstrap {
     public static final int ARRAY_STORE = 4;
     /** static bootstrap parameter indicating a dynamic iteration, e.g. for (x : y) */
     public static final int ITERATOR = 5;
+    /** static bootstrap parameter indicating a dynamic method reference, e.g. foo::bar */
+    public static final int REFERENCE = 6;
 
     /**
      * CallSite that implements the polymorphic inlining cache (PIC).
@@ -71,17 +73,15 @@ public final class DefBootstrap {
         private final Lookup lookup;
         private final String name;
         private final int flavor;
-        private final long recipe;
+        private final Object[] args;
         int depth; // pkg-protected for testing
 
-        PIC(Lookup lookup, String name, MethodType type, int flavor, long recipe) {
+        PIC(Lookup lookup, String name, MethodType type, int flavor, Object[] args) {
             super(type);
             this.lookup = lookup;
             this.name = name;
             this.flavor = flavor;
-            this.recipe = recipe;
+            this.args = args;
-            assert recipe == 0 || flavor == METHOD_CALL;
-            assert Long.bitCount(recipe) <= type.parameterCount();
 
             final MethodHandle fallback = FALLBACK.bindTo(this)
               .asCollector(Object[].class, type.parameterCount())
@@ -101,10 +101,10 @@ public final class DefBootstrap {
         /**
          * Does a slow lookup against the whitelist.
          */
-        private MethodHandle lookup(int flavor, Class<?> clazz, String name, Object[] args, long recipe) throws Throwable {
+        private MethodHandle lookup(int flavor, Class<?> clazz, String name, Object[] args) throws Throwable {
             switch(flavor) {
                 case METHOD_CALL:
-                    return Def.lookupMethod(lookup, clazz, name, args, recipe);
+                    return Def.lookupMethod(lookup, type(), clazz, name, args, (Long) this.args[0]);
                 case LOAD:
                     return Def.lookupGetter(clazz, name);
                 case STORE:
@@ -115,6 +115,8 @@ public final class DefBootstrap {
                     return Def.lookupArrayStore(clazz);
                 case ITERATOR:
                     return Def.lookupIterator(clazz);
+                case REFERENCE:
+                    return Def.lookupReference(lookup, (String) this.args[0], clazz, name);
                 default: throw new AssertionError();
             }
         }
@@ -128,7 +130,7 @@ public final class DefBootstrap {
             final MethodType type = type();
             final Object receiver = args[0];
             final Class<?> receiverClass = receiver.getClass();
-            final MethodHandle target = lookup(flavor, receiverClass, name, args, recipe).asType(type);
+            final MethodHandle target = lookup(flavor, receiverClass, name, args).asType(type);
 
             if (depth >= MAX_DEPTH) {
                 // revert to a vtable call
@@ -170,8 +172,36 @@ public final class DefBootstrap {
      * <p>
      * see https://docs.oracle.com/javase/specs/jvms/se7/html/jvms-6.html#jvms-6.5.invokedynamic
      */
-    public static CallSite bootstrap(Lookup lookup, String name, MethodType type, int flavor, long recipe) {
+    public static CallSite bootstrap(Lookup lookup, String name, MethodType type, int flavor, Object... args) {
-        return new PIC(lookup, name, type, flavor, recipe);
+        // validate arguments
+        switch(flavor) {
+            case METHOD_CALL:
+                if (args.length != 1) {
+                    throw new BootstrapMethodError("Invalid number of parameters for method call");
+                }
+                if (args[0] instanceof Long == false) {
+                    throw new BootstrapMethodError("Illegal parameter for method call: " + args[0]);
+                }
+                long recipe = (Long) args[0];
+                if (Long.bitCount(recipe) > type.parameterCount()) {
+                    throw new BootstrapMethodError("Illegal recipe for method call: too many bits");
+                }
+                break;
+            case REFERENCE:
+                if (args.length != 1) {
+                    throw new BootstrapMethodError("Invalid number of parameters for reference call");
+                }
+                if (args[0] instanceof String == false) {
+                    throw new BootstrapMethodError("Illegal parameter for reference call: " + args[0]);
+                }
+                break;
+            default:
+                if (args.length > 0) {
+                    throw new BootstrapMethodError("Illegal static bootstrap parameters for flavor: " + flavor);
+                }
+                break;
+        }
+        return new PIC(lookup, name, type, flavor, args);
     }
 
 }

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

@@ -890,8 +890,7 @@ public final class Definition {
                             throw new AssertionError(e);
                         }
                     }
-                    owner.methods.put(methodKey,
+                    owner.methods.put(methodKey, method);
-                        new Method(method.name, owner, method.rtn, method.arguments, method.method, method.modifiers, method.handle));
                 }
             }
 

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

@@ -1,5 +1,7 @@
 package org.elasticsearch.painless;
 
+import java.util.function.Function;
+
 /*
  * Licensed to Elasticsearch under one or more contributor
  * license agreements. See the NOTICE file distributed with
@@ -63,4 +65,9 @@ public class FeatureTest {
     public static boolean overloadedStatic(boolean whatToReturn) {
         return whatToReturn;
     }
+    
+    /** method taking two functions! */
+    public Object twoFunctionsOfX(Function<Object,Object> f, Function<Object,Object> g) {
+        return f.apply(g.apply(x));
+    }
 }

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

@@ -53,8 +53,9 @@ public class FunctionRef {
      * @param expected interface type to implement.
      * @param type the left hand side of a method reference expression
      * @param call the right hand side of a method reference expression
+     * @param captures captured arguments
      */
-    public FunctionRef(Definition.Type expected, String type, String call) {
+    public FunctionRef(Definition.Type expected, String type, String call, Class<?>... captures) {
         boolean isCtorReference = "new".equals(call);
         // check its really a functional interface
         // for e.g. Comparable
@@ -66,7 +67,7 @@ public class FunctionRef {
         // e.g. compareTo
         invokedName = method.name;
         // e.g. (Object)Comparator
-        invokedType = MethodType.methodType(expected.clazz);
+        invokedType = MethodType.methodType(expected.clazz, captures);
         // e.g. (Object,Object)int
         interfaceMethodType = method.handle.type().dropParameterTypes(0, 1);
         // lookup requested method
@@ -80,7 +81,15 @@ public class FunctionRef {
             Definition.Method staticImpl = struct.staticMethods.get(new Definition.MethodKey(call, method.arguments.size()));
             if (staticImpl == null) {
                 // otherwise a virtual impl
-                impl = struct.methods.get(new Definition.MethodKey(call, method.arguments.size()-1));
+                final int arity;
+                if (captures.length > 0) {
+                    // receiver captured
+                    arity = method.arguments.size();
+                } else {
+                    // receiver passed
+                    arity = method.arguments.size() - 1;
+                }
+                impl = struct.methods.get(new Definition.MethodKey(call, arity));
             } else {
                 impl = staticImpl;
             }
@@ -98,12 +107,19 @@ public class FunctionRef {
         } else {
             tag = Opcodes.H_INVOKEVIRTUAL;
         }
+        if (impl.owner.clazz.isInterface()) {
             implMethodASM = new Handle(tag, struct.type.getInternalName(), impl.name, impl.method.getDescriptor());
+        } else {
+            implMethodASM = new Handle(tag, impl.owner.type.getInternalName(), impl.name, impl.method.getDescriptor());
+        }
         implMethod = impl.handle;
         if (isCtorReference) {
             samMethodType = MethodType.methodType(interfaceMethodType.returnType(), impl.handle.type().parameterArray());
         } else if (Modifier.isStatic(impl.modifiers)) {
             samMethodType = impl.handle.type();
+        } else if (captures.length > 0) {
+            // drop the receiver, we capture it
+            samMethodType = impl.handle.type().dropParameterTypes(0, 1);
         } else {
             // ensure the receiver type is exact and not a superclass type
             samMethodType = impl.handle.type().changeParameterType(0, struct.clazz);

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

@@ -22,7 +22,6 @@ package org.elasticsearch.painless;
 import org.elasticsearch.painless.Definition.Cast;
 import org.elasticsearch.painless.Definition.Sort;
 import org.elasticsearch.painless.Definition.Type;
-import org.objectweb.asm.ClassVisitor;
 import org.objectweb.asm.ClassWriter;
 import org.objectweb.asm.Label;
 import org.objectweb.asm.Opcodes;

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

@@ -119,7 +119,7 @@ final class ScriptImpl implements ExecutableScript, LeafSearchScript {
     public Object run() {
         try {
             return executable.execute(variables, scorer, doc, aggregationValue);
-        } catch (PainlessError | BootstrapMethodError | Exception t) {
+        } catch (PainlessError | BootstrapMethodError | IllegalAccessError | Exception t) {
             throw convertToScriptException(t);
         }
     }

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

@@ -68,7 +68,7 @@ public final class WriterConstants {
 
     /** dynamic callsite bootstrap signature */
     public final static MethodType DEF_BOOTSTRAP_TYPE =
-        MethodType.methodType(CallSite.class, MethodHandles.Lookup.class, String.class, MethodType.class, int.class, long.class);
+        MethodType.methodType(CallSite.class, MethodHandles.Lookup.class, String.class, MethodType.class, int.class, Object[].class);
     public final static Handle DEF_BOOTSTRAP_HANDLE =
         new Handle(Opcodes.H_INVOKESTATIC, Type.getInternalName(DefBootstrap.class),
             "bootstrap", DEF_BOOTSTRAP_TYPE.toMethodDescriptorString());

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

@@ -100,6 +100,7 @@ import org.elasticsearch.painless.node.AStatement;
 import org.elasticsearch.painless.node.EBinary;
 import org.elasticsearch.painless.node.EBool;
 import org.elasticsearch.painless.node.EBoolean;
+import org.elasticsearch.painless.node.ECapturingFunctionRef;
 import org.elasticsearch.painless.node.EChain;
 import org.elasticsearch.painless.node.EComp;
 import org.elasticsearch.painless.node.EConditional;
@@ -447,15 +448,19 @@ 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.ID() != null) {
+            if (ctx.NEW() != null) {
-            methodText = ctx.ID().getText();
-        } else if (ctx.NEW() != null ){
                 methodText = ctx.NEW().getText();
             } else {
-            throw location(ctx).createError(new IllegalStateException("Illegal tree structure."));
+                methodText = ctx.ID(0).getText();
             }
             return new EFunctionRef(location(ctx), ctx.TYPE().getText(), methodText);
+        } else {
+            // capturing object::method
+            return new ECapturingFunctionRef(location(ctx), ctx.ID(0).getText(), ctx.ID(1).getText());
+        }
     }
 
     @Override

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

@@ -0,0 +1,119 @@
+/*
+ * 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.DefBootstrap;
+import org.elasticsearch.painless.Definition;
+import org.elasticsearch.painless.FunctionRef;
+import org.elasticsearch.painless.Location;
+import org.elasticsearch.painless.MethodWriter;
+import org.elasticsearch.painless.Locals;
+import org.elasticsearch.painless.Locals.Variable;
+import org.objectweb.asm.Opcodes;
+import org.objectweb.asm.Type;
+
+import static org.elasticsearch.painless.WriterConstants.DEF_BOOTSTRAP_HANDLE;
+import static org.elasticsearch.painless.WriterConstants.LAMBDA_BOOTSTRAP_HANDLE;
+
+import java.lang.invoke.LambdaMetafactory;
+
+/**
+ * Represents a capturing function reference.
+ */
+public class ECapturingFunctionRef extends AExpression {
+    public final String type;
+    public final String call;
+    
+    private FunctionRef ref;
+    Variable captured;
+    private boolean defInterface;
+
+    public ECapturingFunctionRef(Location location, String type, String call) {
+        super(location);
+
+        this.type = type;
+        this.call = call;
+    }
+
+    @Override
+    void analyze(Locals variables) {
+        captured = variables.getVariable(location, type);
+        if (expected == null) {
+            defInterface = true;
+            actual = Definition.getType("String");
+        } else {
+            defInterface = false;
+            // static case
+            if (captured.type.sort != Definition.Sort.DEF) {
+                try {
+                    ref = new FunctionRef(expected, captured.type.name, call, captured.type.clazz);
+                } catch (IllegalArgumentException e) {
+                    throw createError(e);
+                }
+            }
+            actual = expected;
+        }
+    }
+
+    @Override
+    void write(MethodWriter writer) {
+        writer.writeDebugInfo(location);
+        if (defInterface && captured.type.sort == Definition.Sort.DEF) {
+            // dynamic interface, dynamic implementation
+            writer.push("D" + type + "." + call + ",1");
+            writer.visitVarInsn(captured.type.type.getOpcode(Opcodes.ILOAD), captured.slot);
+        } else if (defInterface) {
+            // dynamic interface, typed implementation
+            writer.push("S" + captured.type.name + "." + call + ",1");
+            writer.visitVarInsn(captured.type.type.getOpcode(Opcodes.ILOAD), captured.slot);
+        } else if (ref == null) {
+            // typed interface, dynamic implementation
+            writer.visitVarInsn(captured.type.type.getOpcode(Opcodes.ILOAD), captured.slot);
+            String descriptor = Type.getMethodType(expected.type, captured.type.type).getDescriptor();
+            writer.invokeDynamic(call, descriptor, DEF_BOOTSTRAP_HANDLE, (Object)DefBootstrap.REFERENCE, expected.name);
+        } else {
+            // typed interface, typed implementation
+            writer.visitVarInsn(captured.type.type.getOpcode(Opcodes.ILOAD), captured.slot);
+            // convert MethodTypes to asm Type for the constant pool.
+            String invokedType = ref.invokedType.toMethodDescriptorString();
+            Type samMethodType = Type.getMethodType(ref.samMethodType.toMethodDescriptorString());
+            Type interfaceType = Type.getMethodType(ref.interfaceMethodType.toMethodDescriptorString());
+            if (ref.needsBridges()) {
+                writer.invokeDynamic(ref.invokedName, 
+                                     invokedType, 
+                                     LAMBDA_BOOTSTRAP_HANDLE, 
+                                     samMethodType, 
+                                     ref.implMethodASM, 
+                                     samMethodType, 
+                                     LambdaMetafactory.FLAG_BRIDGES, 
+                                     1, 
+                                     interfaceType);
+            } else {
+                writer.invokeDynamic(ref.invokedName, 
+                                     invokedType, 
+                                     LAMBDA_BOOTSTRAP_HANDLE, 
+                                     samMethodType, 
+                                     ref.implMethodASM, 
+                                     samMethodType, 
+                                     0);
+            }
+        }
+    }
+}

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

@@ -64,7 +64,7 @@ public class EFunctionRef extends AExpression {
     @Override
     void write(MethodWriter writer) {
         if (ref == null) {
-            writer.push(type + "." + call);
+            writer.push("S" + type + "." + call + ",0");
         } else {
             writer.writeDebugInfo(location);
             // convert MethodTypes to asm Type for the constant pool.

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

@@ -62,7 +62,7 @@ final class LDefArray extends ALink implements IDefLink {
         writer.writeDebugInfo(location);
 
         String desc = Type.getMethodDescriptor(after.type, Definition.DEF_TYPE.type, index.actual.type);
-        writer.invokeDynamic("arrayLoad", desc, DEF_BOOTSTRAP_HANDLE, (Object)DefBootstrap.ARRAY_LOAD, 0);
+        writer.invokeDynamic("arrayLoad", desc, DEF_BOOTSTRAP_HANDLE, (Object)DefBootstrap.ARRAY_LOAD);
     }
 
     @Override
@@ -70,6 +70,6 @@ final class LDefArray extends ALink implements IDefLink {
         writer.writeDebugInfo(location);
 
         String desc = Type.getMethodDescriptor(Definition.VOID_TYPE.type, Definition.DEF_TYPE.type, index.actual.type, after.type);
-        writer.invokeDynamic("arrayStore", desc, DEF_BOOTSTRAP_HANDLE, (Object)DefBootstrap.ARRAY_STORE, 0);
+        writer.invokeDynamic("arrayStore", desc, DEF_BOOTSTRAP_HANDLE, (Object)DefBootstrap.ARRAY_STORE);
     }
 }

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

@@ -55,11 +55,15 @@ final class LDefCall extends ALink implements IDefLink {
         }
 
         recipe = 0;
+        int totalCaptures = 0;
         for (int argument = 0; argument < arguments.size(); ++argument) {
             AExpression expression = arguments.get(argument);
 
             if (expression instanceof EFunctionRef) {
-                recipe |= (1L << argument); // mark argument as deferred reference
+                recipe |= (1L << (argument + totalCaptures)); // mark argument as deferred reference
+            } else if (expression instanceof ECapturingFunctionRef) {
+                recipe |= (1L << (argument + totalCaptures)); // mark argument as deferred reference
+                totalCaptures++;
             }
             expression.internal = true;
             expression.analyze(locals);
@@ -90,6 +94,10 @@ final class LDefCall extends ALink implements IDefLink {
 
         for (AExpression argument : arguments) {
             signature.append(argument.actual.type.getDescriptor());
+            if (argument instanceof ECapturingFunctionRef) {
+                ECapturingFunctionRef capturingRef = (ECapturingFunctionRef) argument;
+                signature.append(capturingRef.captured.type.type.getDescriptor());
+            }
             argument.write(writer);
         }
 

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

@@ -59,7 +59,7 @@ final class LDefField extends ALink implements IDefLink {
         writer.writeDebugInfo(location);
 
         String desc = Type.getMethodDescriptor(after.type, Definition.DEF_TYPE.type);
-        writer.invokeDynamic(value, desc, DEF_BOOTSTRAP_HANDLE, (Object)DefBootstrap.LOAD, 0);
+        writer.invokeDynamic(value, desc, DEF_BOOTSTRAP_HANDLE, (Object)DefBootstrap.LOAD);
     }
 
     @Override
@@ -67,6 +67,6 @@ final class LDefField extends ALink implements IDefLink {
         writer.writeDebugInfo(location);
 
         String desc = Type.getMethodDescriptor(Definition.VOID_TYPE.type, Definition.DEF_TYPE.type, after.type);
-        writer.invokeDynamic(value, desc, DEF_BOOTSTRAP_HANDLE, (Object)DefBootstrap.STORE, 0);
+        writer.invokeDynamic(value, desc, DEF_BOOTSTRAP_HANDLE, (Object)DefBootstrap.STORE);
     }
 }

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

@@ -195,7 +195,7 @@ public class SEach extends AStatement {
         if (method == null) {
             Type itr = Definition.getType("Iterator");
             String desc = org.objectweb.asm.Type.getMethodDescriptor(itr.type, Definition.DEF_TYPE.type);
-            writer.invokeDynamic("iterator", desc, DEF_BOOTSTRAP_HANDLE, (Object)DefBootstrap.ITERATOR, (Object)0);
+            writer.invokeDynamic("iterator", desc, DEF_BOOTSTRAP_HANDLE, (Object)DefBootstrap.ITERATOR);
         } else if (java.lang.reflect.Modifier.isInterface(method.owner.clazz.getModifiers())) {
             writer.invokeInterface(method.owner.type, method.method);
         } else {

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

@@ -35,6 +35,7 @@
  * {@link org.elasticsearch.painless.node.EBinary} - Represents a binary math expression.
  * {@link org.elasticsearch.painless.node.EBool} - Represents a boolean expression.
  * {@link org.elasticsearch.painless.node.EBoolean} - Represents a boolean constant.
+ * {@link org.elasticsearch.painless.node.ECapturingFunctionRef} - Represents a function reference (capturing).
  * {@link org.elasticsearch.painless.node.ECast} - Represents an implicit cast in most cases.  (Internal only.)
  * {@link org.elasticsearch.painless.node.EChain} - Represents the entirety of a variable/method chain for read/write operations.
  * {@link org.elasticsearch.painless.node.EComp} - Represents a comparison expression.
@@ -42,7 +43,7 @@
  * {@link org.elasticsearch.painless.node.EConstant} - Represents a constant.  (Internal only.)
  * {@link org.elasticsearch.painless.node.EDecimal} - Represents a decimal constant.
  * {@link org.elasticsearch.painless.node.EExplicit} - Represents an explicit cast.
- * {@link org.elasticsearch.painless.node.EFunctionRef} - Represents a function reference.
+ * {@link org.elasticsearch.painless.node.EFunctionRef} - Represents a function reference (non-capturing).
  * {@link org.elasticsearch.painless.node.ENull} - Represents a null constant.
  * {@link org.elasticsearch.painless.node.ENumeric} - Represents a non-decimal numeric constant.
  * {@link org.elasticsearch.painless.node.EUnary} - Represents a unary math expression.

modules/lang-painless/src/main/resources/org/elasticsearch/painless/org.elasticsearch.txt

@@ -125,4 +125,5 @@ class org.elasticsearch.painless.FeatureTest -> org.elasticsearch.painless.Featu
   void setY(int)
   boolean overloadedStatic()
   boolean overloadedStatic(boolean)
+  Object twoFunctionsOfX(Function,Function)
 }

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

@@ -33,7 +33,7 @@ public class DefBootstrapTests extends ESTestCase {
         CallSite site = DefBootstrap.bootstrap(MethodHandles.publicLookup(), 
                                                   "toString", 
                                                   MethodType.methodType(String.class, Object.class), 
-                                                  DefBootstrap.METHOD_CALL, 0);
+                                                  DefBootstrap.METHOD_CALL, 0L);
         MethodHandle handle = site.dynamicInvoker();
         assertDepthEquals(site, 0);
 
@@ -50,7 +50,7 @@ public class DefBootstrapTests extends ESTestCase {
         CallSite site = DefBootstrap.bootstrap(MethodHandles.publicLookup(), 
                                                   "toString", 
                                                   MethodType.methodType(String.class, Object.class), 
-                                                  DefBootstrap.METHOD_CALL, 0);
+                                                  DefBootstrap.METHOD_CALL, 0L);
         MethodHandle handle = site.dynamicInvoker();
         assertDepthEquals(site, 0);
 
@@ -72,7 +72,7 @@ public class DefBootstrapTests extends ESTestCase {
         CallSite site = DefBootstrap.bootstrap(MethodHandles.publicLookup(), 
                                                   "toString", 
                                                   MethodType.methodType(String.class, Object.class), 
-                                                  DefBootstrap.METHOD_CALL, 0);
+                                                  DefBootstrap.METHOD_CALL, 0L);
         MethodHandle handle = site.dynamicInvoker();
         assertDepthEquals(site, 0);
 

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

@@ -19,6 +19,8 @@
 
 package org.elasticsearch.painless;
 
+import java.util.Comparator;
+
 public class FunctionRefTests extends ScriptTestCase {
 
     public void testStaticMethodReference() {
@@ -57,6 +59,58 @@ public class FunctionRefTests extends ScriptTestCase {
                  "return stats.getSum()"));
     }
     
+    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);"));
+    }
+    
+    public void testCapturingMethodReferenceDefImpl() {
+        assertEquals("5", exec("def x = Integer.valueOf(5); return Optional.empty().orElseGet(x::toString);"));
+        assertEquals("[]", exec("def l = new ArrayList(); return Optional.empty().orElseGet(l::toString);"));
+    }
+    
+    public void testCapturingMethodReferenceDefInterface() {
+        assertEquals("5", exec("Integer x = Integer.valueOf(5); def opt = Optional.empty(); return opt.orElseGet(x::toString);"));
+        assertEquals("[]", exec("List l = new ArrayList(); def opt = Optional.empty(); return opt.orElseGet(l::toString);"));
+    }
+    
+    public void testCapturingMethodReferenceDefEverywhere() {
+        assertEquals("5", exec("def x = Integer.valueOf(5); def opt = Optional.empty(); return opt.orElseGet(x::toString);"));
+        assertEquals("[]", exec("def l = new ArrayList(); def opt = Optional.empty(); return opt.orElseGet(l::toString);"));
+    }
+    
+    public void testCapturingMethodReferenceMultipleLambdas() {
+        assertEquals("testingcdefg", exec(
+                "String x = 'testing';" +
+                "String y = 'abcdefg';" + 
+                "org.elasticsearch.painless.FeatureTest test = new org.elasticsearch.painless.FeatureTest(2,3);" + 
+                "return test.twoFunctionsOfX(x::concat, y::substring);"));
+    }
+    
+    public void testCapturingMethodReferenceMultipleLambdasDefImpls() {
+        assertEquals("testingcdefg", exec(
+                "def x = 'testing';" +
+                "def y = 'abcdefg';" + 
+                "org.elasticsearch.painless.FeatureTest test = new org.elasticsearch.painless.FeatureTest(2,3);" + 
+                "return test.twoFunctionsOfX(x::concat, y::substring);"));
+    }
+    
+    public void testCapturingMethodReferenceMultipleLambdasDefInterface() {
+        assertEquals("testingcdefg", exec(
+                "String x = 'testing';" +
+                "String y = 'abcdefg';" + 
+                "def test = new org.elasticsearch.painless.FeatureTest(2,3);" + 
+                "return test.twoFunctionsOfX(x::concat, y::substring);"));
+    }
+    
+    public void testCapturingMethodReferenceMultipleLambdasDefEverywhere() {
+        assertEquals("testingcdefg", exec(
+                "def x = 'testing';" +
+                "def y = 'abcdefg';" + 
+                "def test = new org.elasticsearch.painless.FeatureTest(2,3);" + 
+                "return test.twoFunctionsOfX(x::concat, y::substring);"));
+    }
+
     public void testMethodMissing() {
         IllegalArgumentException expected = expectScriptThrows(IllegalArgumentException.class, () -> {
             exec("List l = new ArrayList(); l.add(2); l.add(1); l.sort(Integer::bogus); return l.get(0);");