Merge pull request #18907 from rmuir/fix_horrible_capture

Fix horrible capture

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

@@ -215,25 +215,27 @@ public final class Def {
      * @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
-     * @param recipe bitset marking functional parameters
+     * @param args bootstrap args passed to callsite
      * @return pointer to matching method to invoke. never returns null.
      * @throws IllegalArgumentException if no matching whitelisted method was found.
      * @throws Throwable if a method reference cannot be converted to an functional interface
      */
      static MethodHandle lookupMethod(Lookup lookup, MethodType callSiteType, 
-             Class<?> receiverClass, String name, Object args[], long recipe) throws Throwable {
+             Class<?> receiverClass, String name, Object args[]) throws Throwable {
+         long recipe = (Long) args[0];
+         int numArguments = callSiteType.parameterCount();
          // simple case: no lambdas
          if (recipe == 0) {
-             return lookupMethodInternal(receiverClass, name, args.length - 1).handle;
+             return lookupMethodInternal(receiverClass, name, numArguments - 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++) {
+         int arity = callSiteType.parameterCount() - 1;
+         int upTo = 1;
+         for (int i = 0; i < numArguments; i++) {
              if ((recipe & (1L << (i - 1))) != 0) {
-                 String signature = (String) args[i];
+                 String signature = (String) args[upTo++];
                  int numCaptures = Integer.parseInt(signature.substring(signature.indexOf(',')+1));
                  arity -= numCaptures;
              }
@@ -245,11 +247,12 @@ public final class Def {
          MethodHandle handle = method.handle;
 
          int replaced = 0;
-         for (int i = 1; i < args.length; i++) {
+         upTo = 1;
+         for (int i = 1; i < numArguments; i++) {
              // its a functional reference, replace the argument with an impl
              if ((recipe & (1L << (i - 1))) != 0) {
                  // decode signature of form 'type.call,2' 
-                 String signature = (String) args[i];
+                 String signature = (String) args[upTo++];
                  int separator = signature.indexOf('.');
                  int separator2 = signature.indexOf(',');
                  String type = signature.substring(1, separator);

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

@@ -125,10 +125,10 @@ public final class DefBootstrap {
         /**
          * Does a slow lookup against the whitelist.
          */
-        private MethodHandle lookup(int flavor, String name, Class<?> receiver, Object[] callArgs) throws Throwable {
+        private MethodHandle lookup(int flavor, String name, Class<?> receiver) throws Throwable {
             switch(flavor) {
                 case METHOD_CALL:
-                    return Def.lookupMethod(lookup, type(), receiver, name, callArgs, (Long) this.args[0]);
+                    return Def.lookupMethod(lookup, type(), receiver, name, args);
                 case LOAD:
                     return Def.lookupGetter(receiver, name);
                 case STORE:
@@ -140,7 +140,7 @@ public final class DefBootstrap {
                 case ITERATOR:
                     return Def.lookupIterator(receiver);
                 case REFERENCE:
-                    return Def.lookupReference(lookup, (String) this.args[0], receiver, name);
+                    return Def.lookupReference(lookup, (String) args[0], receiver, name);
                 default: throw new AssertionError();
             }
         }
@@ -148,17 +148,15 @@ public final class DefBootstrap {
         /**
          * Creates the {@link MethodHandle} for the megamorphic call site
          * using {@link ClassValue} and {@link MethodHandles#exactInvoker(MethodType)}:
-         * <p>
-         * TODO: Remove the variable args and just use {@code type()}!
          */
-        private MethodHandle createMegamorphicHandle(final Object[] callArgs) throws Throwable {
+        private MethodHandle createMegamorphicHandle() throws Throwable {
             final MethodType type = type();
             final ClassValue<MethodHandle> megamorphicCache = new ClassValue<MethodHandle>() {
                 @Override
                 protected MethodHandle computeValue(Class<?> receiverType) {
                     // it's too stupid that we cannot throw checked exceptions... (use rethrow puzzler):
                     try {
-                        return lookup(flavor, name, receiverType, callArgs).asType(type);
+                        return lookup(flavor, name, receiverType).asType(type);
                     } catch (Throwable t) {
                         Def.rethrow(t);
                         throw new AssertionError();
@@ -180,13 +178,13 @@ public final class DefBootstrap {
         Object fallback(final Object[] callArgs) throws Throwable {
             if (depth >= MAX_DEPTH) {
                 // we revert the whole cache and build a new megamorphic one
-                final MethodHandle target = this.createMegamorphicHandle(callArgs);
+                final MethodHandle target = this.createMegamorphicHandle();
                 
                 setTarget(target);
                 return target.invokeWithArguments(callArgs);                    
             } else {
                 final Class<?> receiver = callArgs[0].getClass();
-                final MethodHandle target = lookup(flavor, name, receiver, callArgs).asType(type());
+                final MethodHandle target = lookup(flavor, name, receiver).asType(type());
     
                 MethodHandle test = CHECK_CLASS.bindTo(receiver);
                 MethodHandle guard = MethodHandles.guardWithTest(test, target, getTarget());
@@ -398,16 +396,20 @@ public final class DefBootstrap {
         switch(flavor) {
             // "function-call" like things get a polymorphic cache
             case METHOD_CALL:
-                if (args.length != 1) {
+                if (args.length == 0) {
                     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()) {
+                int numLambdas = Long.bitCount(recipe);
+                if (numLambdas > type.parameterCount()) {
                     throw new BootstrapMethodError("Illegal recipe for method call: too many bits");
                 }
+                if (args.length != numLambdas + 1) {
+                    throw new BootstrapMethodError("Illegal number of parameters: expected " + numLambdas + " references");
+                }
                 return new PIC(lookup, name, type, flavor, args);
             case LOAD:
             case STORE:

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

@@ -43,7 +43,7 @@ public class ECapturingFunctionRef extends AExpression {
     
     private FunctionRef ref;
     Variable captured;
-    private boolean defInterface;
+    String defPointer;
 
     public ECapturingFunctionRef(Location location, String type, String call) {
         super(location);
@@ -56,10 +56,16 @@ public class ECapturingFunctionRef extends AExpression {
     void analyze(Locals variables) {
         captured = variables.getVariable(location, type);
         if (expected == null) {
-            defInterface = true;
+            if (captured.type.sort == Definition.Sort.DEF) {
+                // dynamic implementation
+                defPointer = "D" + type + "." + call + ",1";
+            } else {
+                // typed implementation
+                defPointer = "S" + captured.type.name + "." + call + ",1";
+            }
             actual = Definition.getType("String");
         } else {
-            defInterface = false;
+            defPointer = null;
             // static case
             if (captured.type.sort != Definition.Sort.DEF) {
                 try {
@@ -75,13 +81,10 @@ public class ECapturingFunctionRef extends AExpression {
     @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");
+        if (defPointer != null) {
+            // dynamic interface: push captured parameter on stack
+            // TODO: don't do this: its just to cutover :)
+            writer.push((String)null);
             writer.visitVarInsn(captured.type.type.getOpcode(Opcodes.ILOAD), captured.slot);
         } else if (ref == null) {
             // typed interface, dynamic implementation

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

@@ -40,6 +40,7 @@ public class EFunctionRef extends AExpression {
     public final String call;
 
     private FunctionRef ref;
+    String defPointer;
 
     public EFunctionRef(Location location, String type, String call) {
         super(location);
@@ -53,7 +54,9 @@ public class EFunctionRef extends AExpression {
         if (expected == null) {
             ref = null;
             actual = Definition.getType("String");
+            defPointer = "S" + type + "." + call + ",0";
         } else {
+            defPointer = null;
             try {
                 if ("this".equals(type)) {
                     // user's own function
@@ -81,9 +84,7 @@ public class EFunctionRef extends AExpression {
 
     @Override
     void write(MethodWriter writer) {
-        if (ref == null) {
-            writer.push("S" + type + "." + call + ",0");
-        } else {
+        if (ref != null) {
             writer.writeDebugInfo(location);
             // convert MethodTypes to asm Type for the constant pool.
             String invokedType = ref.invokedType.toMethodDescriptorString();
@@ -108,6 +109,9 @@ public class EFunctionRef extends AExpression {
                                      samMethodType,
                                      0);
             }
+        } else {
+            // TODO: don't do this: its just to cutover :)
+            writer.push((String)null);
         }
     }
 }

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

@@ -25,6 +25,7 @@ import org.elasticsearch.painless.DefBootstrap;
 import org.elasticsearch.painless.Locals;
 import org.elasticsearch.painless.MethodWriter;
 
+import java.util.ArrayList;
 import java.util.List;
 
 import static org.elasticsearch.painless.WriterConstants.DEF_BOOTSTRAP_HANDLE;
@@ -37,6 +38,7 @@ final class LDefCall extends ALink implements IDefLink {
     final String name;
     final List<AExpression> arguments;
     long recipe;
+    List<String> pointers = new ArrayList<>();
 
     LDefCall(Location location, String name, List<AExpression> arguments) {
         super(location, -1);
@@ -59,14 +61,18 @@ final class LDefCall extends ALink implements IDefLink {
         for (int argument = 0; argument < arguments.size(); ++argument) {
             AExpression expression = arguments.get(argument);
 
+            expression.internal = true;
+            expression.analyze(locals);
+
             if (expression instanceof EFunctionRef) {
+                pointers.add(((EFunctionRef)expression).defPointer);
                 recipe |= (1L << (argument + totalCaptures)); // mark argument as deferred reference
             } else if (expression instanceof ECapturingFunctionRef) {
+                pointers.add(((ECapturingFunctionRef)expression).defPointer);
                 recipe |= (1L << (argument + totalCaptures)); // mark argument as deferred reference
                 totalCaptures++;
             }
-            expression.internal = true;
-            expression.analyze(locals);
+
             expression.expected = expression.actual;
             arguments.set(argument, expression.cast(locals));
         }
@@ -105,7 +111,11 @@ final class LDefCall extends ALink implements IDefLink {
         // return value
         signature.append(after.type.getDescriptor());
 
-        writer.invokeDynamic(name, signature.toString(), DEF_BOOTSTRAP_HANDLE, DefBootstrap.METHOD_CALL, recipe);
+        List<Object> args = new ArrayList<>();
+        args.add(DefBootstrap.METHOD_CALL);
+        args.add(recipe);
+        args.addAll(pointers);
+        writer.invokeDynamic(name, signature.toString(), DEF_BOOTSTRAP_HANDLE, args.toArray());
     }
 
     @Override