Add typed emit functionality

modules/lang-painless/spi/src/main/java/org/elasticsearch/painless/spi/annotation/InjectScriptAnnotation.java

@@ -0,0 +1,27 @@
+/*
+ * 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.spi.annotation;
+
+/**
+ * Inject the script itself into a method call. Only allowed on augmentations.
+ */
+public class InjectScriptAnnotation {
+    public static final String NAME = "inject_script";
+}

modules/lang-painless/spi/src/main/java/org/elasticsearch/painless/spi/annotation/InjectScriptAnnotationParser.java

@@ -0,0 +1,37 @@
+/*
+ * 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.spi.annotation;
+
+import java.util.Map;
+
+public class InjectScriptAnnotationParser implements WhitelistAnnotationParser {
+
+    public static final InjectScriptAnnotationParser INSTANCE = new InjectScriptAnnotationParser();
+
+    private InjectScriptAnnotationParser() {}
+
+    @Override
+    public Object parse(Map<String, String> arguments) {
+        if (false == arguments.isEmpty()) {
+            throw new IllegalArgumentException("[@inject_script] can no have parameters");
+        }
+        return new InjectScriptAnnotation();
+    }
+}

modules/lang-painless/spi/src/main/java/org/elasticsearch/painless/spi/annotation/WhitelistAnnotationParser.java

@@ -36,7 +36,8 @@ public interface WhitelistAnnotationParser {
                     new AbstractMap.SimpleEntry<>(NoImportAnnotation.NAME, NoImportAnnotationParser.INSTANCE),
                     new AbstractMap.SimpleEntry<>(DeprecatedAnnotation.NAME, DeprecatedAnnotationParser.INSTANCE),
                     new AbstractMap.SimpleEntry<>(NonDeterministicAnnotation.NAME, NonDeterministicAnnotationParser.INSTANCE),
-                    new AbstractMap.SimpleEntry<>(InjectConstantAnnotation.NAME, InjectConstantAnnotationParser.INSTANCE)
+                    new AbstractMap.SimpleEntry<>(InjectConstantAnnotation.NAME, InjectConstantAnnotationParser.INSTANCE),
+                    new AbstractMap.SimpleEntry<>(InjectScriptAnnotation.NAME, InjectScriptAnnotationParser.INSTANCE)
             ).collect(Collectors.toMap(Map.Entry::getKey, Map.Entry::getValue))
     );
 

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

@@ -22,6 +22,7 @@
 import org.elasticsearch.painless.lookup.PainlessLookup;
 import org.elasticsearch.painless.lookup.PainlessLookupUtility;
 import org.elasticsearch.painless.lookup.PainlessMethod;
+import org.elasticsearch.painless.spi.annotation.InjectScriptAnnotation;
 import org.elasticsearch.painless.symbol.FunctionTable;
 import org.elasticsearch.script.JodaCompatibleZonedDateTime;
 
@@ -196,25 +197,37 @@ static MethodHandle lookupMethod(PainlessLookup painlessLookup, FunctionTable fu
             MethodHandles.Lookup methodHandlesLookup, MethodType callSiteType, Class<?> receiverClass, String name, Object[] args)
             throws Throwable {
 
-         String recipeString = (String) args[0];
+         int argsOffset = 0; 
+         String recipeString = (String) args[argsOffset++];
+         boolean injectedScript = ((Integer) args[argsOffset++] > 0);
          int numArguments = callSiteType.parameterCount();
+         int arityDiff = 1;
+         if (injectedScript) {
+             arityDiff++;
+         }
          // simple case: no lambdas
          if (recipeString.isEmpty()) {
-             PainlessMethod painlessMethod = painlessLookup.lookupRuntimePainlessMethod(receiverClass, name, numArguments - 1);
+             PainlessMethod painlessMethod = painlessLookup.lookupRuntimePainlessMethod(receiverClass, name, numArguments - arityDiff);
 
              if (painlessMethod == null) {
                  throw new IllegalArgumentException("dynamic method " +
-                         "[" + typeToCanonicalTypeName(receiverClass) + ", " + name + "/" + (numArguments - 1) + "] not found");
+                         "[" + typeToCanonicalTypeName(receiverClass) + ", " + name + "/" + (numArguments - arityDiff) + "] not found");
              }
 
              MethodHandle handle = painlessMethod.methodHandle;
-             Object[] injections = PainlessLookupUtility.buildInjections(painlessMethod, constants);
 
+             Object[] injections = PainlessLookupUtility.buildInjections(painlessMethod, constants);
              if (injections.length > 0) {
-                 // method handle contains the "this" pointer so start injections at 1
+                 // method handle contains the receiver pointer so start injections at 1
                  handle = MethodHandles.insertArguments(handle, 1, injections);
              }
 
+             if (injectedScript && painlessMethod.annotations.containsKey(InjectScriptAnnotation.class) == false) {
+                 // Strip the `Script` parameter if we don't need it
+                 Class<?> typeToInject = painlessLookup.typeOfScriptToInjectForMethod(name, numArguments - arityDiff);
+                 handle = MethodHandles.dropArguments(handle, 1, typeToInject);
+             }
+
              return handle;
          }
 
@@ -227,7 +240,7 @@ static MethodHandle lookupMethod(PainlessLookup painlessLookup, FunctionTable fu
          // 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 = callSiteType.parameterCount() - 1;
-         int upTo = 1;
+         int upTo = argsOffset;
          for (int i = 1; i < numArguments; i++) {
              if (lambdaArgs.get(i - 1)) {
                  String signature = (String) args[upTo++];
@@ -247,14 +260,13 @@ static MethodHandle lookupMethod(PainlessLookup painlessLookup, FunctionTable fu
 
         MethodHandle handle = method.methodHandle;
         Object[] injections = PainlessLookupUtility.buildInjections(method, constants);
-
         if (injections.length > 0) {
             // method handle contains the "this" pointer so start injections at 1
             handle = MethodHandles.insertArguments(handle, 1, injections);
         }
 
          int replaced = 0;
-         upTo = 1;
+         upTo = argsOffset;
          for (int i = 1; i < numArguments; i++) {
              // its a functional reference, replace the argument with an impl
              if (lambdaArgs.get(i - 1)) {
@@ -360,6 +372,7 @@ private static MethodHandle lookupReferenceInternal(
                 ref.delegateMethodType,
                 ref.isDelegateInterface ? 1 : 0,
                 ref.isDelegateAugmented ? 1 : 0,
+                ref.injectScript ? 1 : 0,
                 ref.delegateInjections
         );
         return callSite.dynamicInvoker().asType(MethodType.methodType(clazz, ref.factoryMethodType.parameterArray()));

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

@@ -445,7 +445,7 @@ public static CallSite bootstrap(PainlessLookup painlessLookup, FunctionTable fu
         switch(flavor) {
             // "function-call" like things get a polymorphic cache
             case METHOD_CALL:
-                if (args.length == 0) {
+                if (args.length < 2) {
                     throw new BootstrapMethodError("Invalid number of parameters for method call");
                 }
                 if (args[0] instanceof String == false) {
@@ -456,7 +456,7 @@ public static CallSite bootstrap(PainlessLookup painlessLookup, FunctionTable fu
                 if (numLambdas > type.parameterCount()) {
                     throw new BootstrapMethodError("Illegal recipe for method call: too many bits");
                 }
-                if (args.length != numLambdas + 1) {
+                if (args.length != numLambdas + 2) {
                     throw new BootstrapMethodError("Illegal number of parameters: expected " + numLambdas + " references");
                 }
                 return new PIC(painlessLookup, functions, constants, methodHandlesLookup, name, type, initialDepth, flavor, args);

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

@@ -23,6 +23,7 @@
 import org.elasticsearch.painless.lookup.PainlessLookup;
 import org.elasticsearch.painless.lookup.PainlessLookupUtility;
 import org.elasticsearch.painless.lookup.PainlessMethod;
+import org.elasticsearch.painless.spi.annotation.InjectScriptAnnotation;
 import org.elasticsearch.painless.symbol.FunctionTable;
 import org.elasticsearch.painless.symbol.FunctionTable.LocalFunction;
 
@@ -84,6 +85,7 @@ public static FunctionRef create(PainlessLookup painlessLookup, FunctionTable fu
             String delegateMethodName;
             MethodType delegateMethodType;
             Object[] delegateInjections;
+            boolean injectScript;
 
             Class<?> delegateMethodReturnType;
             List<Class<?>> delegateMethodParameters;
@@ -113,6 +115,7 @@ public static FunctionRef create(PainlessLookup painlessLookup, FunctionTable fu
                 delegateMethodName = localFunction.getFunctionName();
                 delegateMethodType = localFunction.getMethodType();
                 delegateInjections = new Object[0];
+                injectScript = false;
 
                 delegateMethodReturnType = localFunction.getReturnType();
                 delegateMethodParameters = localFunction.getTypeParameters();
@@ -136,6 +139,7 @@ public static FunctionRef create(PainlessLookup painlessLookup, FunctionTable fu
                 delegateMethodName = PainlessLookupUtility.CONSTRUCTOR_NAME;
                 delegateMethodType = painlessConstructor.methodType;
                 delegateInjections = new Object[0];
+                injectScript = false;
 
                 delegateMethodReturnType = painlessConstructor.javaConstructor.getDeclaringClass();
                 delegateMethodParameters = painlessConstructor.typeParameters;
@@ -177,6 +181,7 @@ public static FunctionRef create(PainlessLookup painlessLookup, FunctionTable fu
                 delegateMethodName = painlessMethod.javaMethod.getName();
                 delegateMethodType = painlessMethod.methodType;
                 delegateInjections = PainlessLookupUtility.buildInjections(painlessMethod, constants);
+                injectScript = painlessMethod.annotations.containsKey(InjectScriptAnnotation.class);
 
                 delegateMethodReturnType = painlessMethod.returnType;
 
@@ -204,9 +209,14 @@ public static FunctionRef create(PainlessLookup painlessLookup, FunctionTable fu
                     delegateMethodType.dropParameterTypes(numberOfCaptures, delegateMethodType.parameterCount()));
             delegateMethodType = delegateMethodType.dropParameterTypes(0, numberOfCaptures);
 
+            if (injectScript) {
+                factoryMethodType = factoryMethodType.insertParameterTypes(0, delegateMethodType.parameterType(1));
+            }
+
             return new FunctionRef(interfaceMethodName, interfaceMethodType,
                     delegateClassName, isDelegateInterface, isDelegateAugmented,
                     delegateInvokeType, delegateMethodName, delegateMethodType, delegateInjections,
+                    injectScript,
                     factoryMethodType
             );
         } catch (IllegalArgumentException iae) {
@@ -236,13 +246,16 @@ public static FunctionRef create(PainlessLookup painlessLookup, FunctionTable fu
     public final MethodType delegateMethodType;
     /** injected constants */
     public final Object[] delegateInjections;
+    /** Should the script be injected into the method arguments? */
+    public final boolean injectScript;
     /** factory (CallSite) method signature */
     public final MethodType factoryMethodType;
 
     private FunctionRef(
             String interfaceMethodName, MethodType interfaceMethodType,
             String delegateClassName, boolean isDelegateInterface, boolean isDelegateAugmented,
             int delegateInvokeType, String delegateMethodName, MethodType delegateMethodType, Object[] delegateInjections,
+            boolean injectScript,
             MethodType factoryMethodType) {
 
         this.interfaceMethodName = interfaceMethodName;
@@ -254,6 +267,7 @@ private FunctionRef(
         this.delegateMethodName = delegateMethodName;
         this.delegateMethodType = delegateMethodType;
         this.delegateInjections = delegateInjections;
+        this.injectScript = injectScript;
         this.factoryMethodType = factoryMethodType;
     }
 }

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

@@ -31,11 +31,11 @@
 import java.lang.invoke.LambdaConversionException;
 import java.lang.invoke.MethodHandle;
 import java.lang.invoke.MethodHandles;
+import java.lang.invoke.MethodHandles.Lookup;
 import java.lang.invoke.MethodType;
 import java.security.AccessController;
 import java.security.PrivilegedAction;
 
-import static java.lang.invoke.MethodHandles.Lookup;
 import static org.elasticsearch.painless.WriterConstants.CLASS_VERSION;
 import static org.elasticsearch.painless.WriterConstants.CTOR_METHOD_NAME;
 import static org.elasticsearch.painless.WriterConstants.DELEGATE_BOOTSTRAP_HANDLE;
@@ -210,6 +210,7 @@ public static CallSite lambdaBootstrap(
             MethodType delegateMethodType,
             int isDelegateInterface,
             int isDelegateAugmented,
+            int injectScript,
             Object... injections)
             throws LambdaConversionException {
         Compiler.Loader loader = (Compiler.Loader)lookup.lookupClass().getClassLoader();
@@ -233,9 +234,22 @@ public static CallSite lambdaBootstrap(
             delegateInvokeType = H_INVOKESTATIC;
         }
 
-        generateInterfaceMethod(cw, factoryMethodType, lambdaClassType, interfaceMethodName,
-            interfaceMethodType, delegateClassType, delegateInvokeType,
-            delegateMethodName, delegateMethodType, isDelegateInterface == 1, isDelegateAugmented == 1, captures, injections);
+        generateInterfaceMethod(
+            cw,
+            factoryMethodType,
+            lambdaClassType,
+            interfaceMethodName,
+            interfaceMethodType,
+            delegateClassType,
+            delegateInvokeType,
+            delegateMethodName,
+            delegateMethodType,
+            isDelegateInterface == 1,
+            isDelegateAugmented == 1,
+            captures,
+            injectScript == 1,
+            injections
+        );
 
         endLambdaClass(cw);
 
@@ -382,6 +396,7 @@ private static void generateInterfaceMethod(
             boolean isDelegateInterface,
             boolean isDelegateAugmented,
             Capture[] captures,
+            boolean injectScript,
             Object... injections)
             throws LambdaConversionException {
 
@@ -394,15 +409,24 @@ private static void generateInterfaceMethod(
         iface.visitCode();
 
         // Loads any captured variables onto the stack.
-        for (int captureCount = 0; captureCount < captures.length; ++captureCount) {
+        int capturesBeforeArgs = captures.length;
+        if (injectScript) {
+            capturesBeforeArgs--;
+        }
+        for (int captureCount = 0; captureCount < capturesBeforeArgs; ++captureCount) {
             iface.loadThis();
-            iface.getField(
-                lambdaClassType, captures[captureCount].name, captures[captureCount].type);
+            iface.getField(lambdaClassType, captures[captureCount].name, captures[captureCount].type);
         }
 
         // Loads any passed in arguments onto the stack.
         iface.loadArgs();
 
+        if (injectScript) {
+            iface.loadThis();
+            Capture scriptCapture = captures[captures.length - 1];
+            iface.getField(lambdaClassType, scriptCapture.name, scriptCapture.type);
+        }
+
         // Handles the case for a lambda function or a static reference method.
         // interfaceMethodType and delegateMethodType both have the captured types
         // inserted into their type signatures.  This later allows the delegate
@@ -411,10 +435,24 @@ private static void generateInterfaceMethod(
         // Example: Integer::parseInt
         // Example: something.each(x -> x + 1)
         if (delegateInvokeType == H_INVOKESTATIC) {
-            interfaceMethodType =
-                interfaceMethodType.insertParameterTypes(0, factoryMethodType.parameterArray());
-            delegateMethodType =
-                delegateMethodType.insertParameterTypes(0, factoryMethodType.parameterArray());
+            Class<?>[] inject = factoryMethodType.parameterArray();
+            if (injectScript) {
+                inject = new Class<?>[inject.length - 1];
+                System.arraycopy(factoryMethodType.parameterArray(), 0, inject, 0, inject.length);
+                Class<?> scriptParam = factoryMethodType.parameterType(inject.length);
+                interfaceMethodType = interfaceMethodType.insertParameterTypes(1, scriptParam);
+            }
+            interfaceMethodType = interfaceMethodType.insertParameterTypes(0, inject);
+            delegateMethodType = delegateMethodType.insertParameterTypes(0, inject);
+            if (injectScript) {
+                /*
+                 * The script needs to be shifted after the receiver.
+                 */
+                factoryMethodType = factoryMethodType.dropParameterTypes(
+                    factoryMethodType.parameterCount() - 1,
+                    factoryMethodType.parameterCount()
+                );
+            }
         } else if (delegateInvokeType == H_INVOKEVIRTUAL ||
             delegateInvokeType == H_INVOKEINTERFACE) {
             // Handles the case for a virtual or interface reference method with no captures.
@@ -482,7 +520,7 @@ private static Class<?> createLambdaClass(
 
         byte[] classBytes = cw.toByteArray();
         // DEBUG:
-        // new ClassReader(classBytes).accept(new TraceClassVisitor(new PrintWriter(System.out)), ClassReader.SKIP_DEBUG);
+//         new ClassReader(classBytes).accept(new TraceClassVisitor(new PrintWriter(System.out)), ClassReader.SKIP_DEBUG);
         return AccessController.doPrivileged((PrivilegedAction<Class<?>>)() ->
             loader.defineLambda(lambdaClassType.getClassName(), classBytes));
     }

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

@@ -515,15 +515,17 @@ public void invokeMethodCall(PainlessMethod painlessMethod) {
     }
 
     public void invokeLambdaCall(FunctionRef functionRef) {
-        Object[] args = new Object[7 + functionRef.delegateInjections.length];
+        int fixedArgs = 8;
+        Object[] args = new Object[fixedArgs + functionRef.delegateInjections.length];
         args[0] = Type.getMethodType(functionRef.interfaceMethodType.toMethodDescriptorString());
         args[1] = functionRef.delegateClassName;
         args[2] = functionRef.delegateInvokeType;
         args[3] = functionRef.delegateMethodName;
         args[4] = Type.getMethodType(functionRef.delegateMethodType.toMethodDescriptorString());
         args[5] = functionRef.isDelegateInterface ? 1 : 0;
         args[6] = functionRef.isDelegateAugmented ? 1 : 0;
-        System.arraycopy(functionRef.delegateInjections, 0, args, 7, functionRef.delegateInjections.length);
+        args[7] = functionRef.injectScript ? 1 : 0;
+        System.arraycopy(functionRef.delegateInjections, 0, args, fixedArgs, functionRef.delegateInjections.length);
 
         invokeDynamic(
                 functionRef.interfaceMethodName,