Skip to content

Commit

Permalink
Add typed emit functionality
Browse files Browse the repository at this point in the history
This creates `.emit()` methods on a few primitive things for `keyword`
and `date` style runtime fields. Now you can do stuff like:
```
"d": {
  "type": "date",
  "script": "'2020-01-18T17:41:34.000Z'.emit()"
}
```

We get to piggy back of painless's dynamic dispatch code to regonize
that you are in a `date` and trying to emit a `String` so we can "do the
right thing" without any extra runtime cost. In this case we just parse
the date using the `formatter` on the field. And since the example above
doesn't use a formatter we get ISO8601, the date format of kings.

Similarly, you can do stuff like:
```
"s": {
  "type": "keyword",
  "script": """
    for (int i = 0; i < 100; i++) {
      i.emit();
    }
  """
}
```

Painless *knows* `i` is an `int` and will call an emit method that emits
its string value.

Also! Assuming we get the syntax proposed in elastic#68088, because this is a
chain of method invocations you can do something like:
```
"i": {
  "type": "long",
  "script": """
    grok('%{NUMBER:i} %{NUMBER:j}').extract(doc['message'].value)?.i?.emit()
  """
}
```

This should be read as "if grok matches the message and extracts a value
for `i` then emit it to the runtime field." If either the grok doesn't
match or doesn't extract an `i` value then nothing will be emitted. As
an extra nice thing - we'll automatilly convert whatever the `grok`
expression outputs into a `long`. All of it handled for us using
painless's standard dynamic dispatch code.
  • Loading branch information
nik9000 committed Jan 29, 2021
1 parent a5add32 commit e344828
Show file tree
Hide file tree
Showing 28 changed files with 1,024 additions and 214 deletions.
Original file line number Diff line number Diff line change
@@ -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";
}
Original file line number Diff line number Diff line change
@@ -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();
}
}
Original file line number Diff line number Diff line change
Expand Up @@ -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))
);

Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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;

Expand Down Expand Up @@ -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;
}

Expand All @@ -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++];
Expand All @@ -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)) {
Expand Down Expand Up @@ -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()));
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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) {
Expand All @@ -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);
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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;

Expand Down Expand Up @@ -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;
Expand Down Expand Up @@ -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();
Expand All @@ -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;
Expand Down Expand Up @@ -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;

Expand Down Expand Up @@ -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) {
Expand Down Expand Up @@ -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;
Expand All @@ -254,6 +267,7 @@ private FunctionRef(
this.delegateMethodName = delegateMethodName;
this.delegateMethodType = delegateMethodType;
this.delegateInjections = delegateInjections;
this.injectScript = injectScript;
this.factoryMethodType = factoryMethodType;
}
}
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down Expand Up @@ -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();
Expand All @@ -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);

Expand Down Expand Up @@ -382,6 +396,7 @@ private static void generateInterfaceMethod(
boolean isDelegateInterface,
boolean isDelegateAugmented,
Capture[] captures,
boolean injectScript,
Object... injections)
throws LambdaConversionException {

Expand All @@ -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
Expand All @@ -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.
Expand Down Expand Up @@ -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));
}
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -515,14 +515,15 @@ public void invokeMethodCall(PainlessMethod painlessMethod) {
}

public void invokeLambdaCall(FunctionRef functionRef) {
Object[] args = new Object[7 + functionRef.delegateInjections.length];
Object[] args = new Object[8 + 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;
args[7] = functionRef.injectScript ? 1 : 0;
System.arraycopy(functionRef.delegateInjections, 0, args, 7, functionRef.delegateInjections.length);

invokeDynamic(
Expand Down
Loading

0 comments on commit e344828

Please sign in to comment.