Last active
September 30, 2023 06:26
-
-
Save forax/176a61cc9707c884bfb059c31811dac9 to your computer and use it in GitHub Desktop.
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
import java.lang.annotation.Annotation; | |
import java.lang.annotation.ElementType; | |
import java.lang.annotation.Retention; | |
import java.lang.annotation.RetentionPolicy; | |
import java.lang.annotation.Target; | |
import java.lang.invoke.MethodHandle; | |
import java.lang.invoke.MethodHandles; | |
import java.lang.invoke.MethodHandles.Lookup; | |
import java.lang.invoke.MethodType; | |
import java.lang.invoke.MutableCallSite; | |
import java.lang.reflect.Method; | |
import java.lang.reflect.Modifier; | |
import java.util.Arrays; | |
import java.util.List; | |
import java.util.stream.Stream; | |
import static java.lang.invoke.MethodHandles.exactInvoker; | |
import static java.lang.invoke.MethodHandles.foldArguments; | |
import static java.lang.invoke.MethodHandles.lookup; | |
import static java.lang.invoke.MethodType.methodType; | |
public class DestructuredVisitor { | |
@Target(ElementType.PARAMETER) | |
@Retention(RetentionPolicy.RUNTIME) | |
public @interface Signature { | |
Class<?>[] value(); | |
} | |
private record MethodData(MethodHandle target, List<Signature> signatures) { } | |
private static final ScopedValue<MethodData> CURRENT_METHOD_DATA = ScopedValue.newInstance(); | |
private final Lookup lookup; | |
private final ClassValue<MethodData> cache = new ClassValue<>() { | |
@Override | |
protected MethodData computeValue(Class<?> type) { | |
return CURRENT_METHOD_DATA.get(); | |
} | |
}; | |
private DestructuredVisitor(Lookup lookup) { | |
this.lookup = lookup; | |
} | |
private static List<Signature> toSignatures(Annotation[][] parameterAnnotations) { | |
return Arrays.stream(parameterAnnotations) | |
.map(annotations -> | |
Arrays.stream(annotations) | |
.flatMap(annotation -> annotation instanceof Signature signature ? Stream.of(signature) : null) | |
.findFirst().orElse(null)) | |
.toList(); | |
} | |
public static DestructuredVisitor of(Lookup lookup, List<Method> methods) { | |
var destructuredVisitor = new DestructuredVisitor(lookup); | |
for(var method: methods) { | |
if (!Modifier.isStatic(method.getModifiers())) { | |
throw new IllegalArgumentException("method " + method + " is not static"); | |
} | |
MethodHandle mh; | |
try { | |
mh = lookup.unreflect(method); | |
} catch (IllegalAccessException e) { | |
throw (IllegalAccessError) new IllegalAccessError().initCause(e); | |
} | |
var methodType = mh.type(); | |
if (methodType.parameterCount() == 0) { | |
throw new IllegalArgumentException("method " + method + " should have at least one parameter"); | |
} | |
var methodData = new MethodData(mh, toSignatures(method.getParameterAnnotations())); | |
// inject the method into the cache | |
ScopedValue.runWhere(CURRENT_METHOD_DATA, methodData, () -> destructuredVisitor.cache.get(methodType.parameterType(0))); | |
} | |
return destructuredVisitor; | |
} | |
private static MethodType toMethodType(Class<?>[] signatureClasses) { | |
return methodType( | |
signatureClasses[0], | |
Arrays.stream(signatureClasses).skip(1).toArray(Class<?>[]::new)); | |
} | |
private MethodHandle appendInliningCaches(MethodHandle target, List<Signature> signatures) { | |
var type = target.type(); | |
for(var i = type.parameterCount(); --i >= 0;) { | |
var signature = signatures.get(i); | |
if (signature != null) { | |
var inliningCacheMH = new InliningCache(toMethodType(signature.value())).dynamicInvoker(); | |
target = MethodHandles.insertArguments(target, i, inliningCacheMH); | |
} | |
} | |
return target; | |
} | |
private final class InliningCache extends MutableCallSite { | |
private static final MethodHandle FALLBACK, POINTER_CHECK; | |
static { | |
var lookup = lookup(); | |
try { | |
FALLBACK = lookup.findVirtual(InliningCache.class, "fallback", | |
methodType(MethodHandle.class, Object.class)); | |
POINTER_CHECK = lookup.findStatic(InliningCache.class, "pointerCheck", | |
methodType(boolean.class, Class.class, Object.class)); | |
} catch (NoSuchMethodException | IllegalAccessException e) { | |
throw new AssertionError(e); | |
} | |
} | |
public InliningCache(MethodType type) { | |
super(type); | |
setTarget(foldArguments(exactInvoker(type), FALLBACK.bindTo(this).asType(methodType(MethodHandle.class, type.parameterType(0))))); | |
} | |
private static boolean pointerCheck(Class<?> type, Object receiver) { | |
return type == receiver.getClass(); | |
} | |
private MethodHandle fallback(Object receiver) { | |
var receiverClass = receiver.getClass(); | |
var methodData = cache.get(receiverClass); | |
var target = appendInliningCaches(methodData.target, methodData.signatures).asType(type()); | |
var test = POINTER_CHECK.bindTo(receiverClass); | |
var guard = MethodHandles.guardWithTest( | |
test.asType(methodType(boolean.class, type().parameterType(0))), | |
target, | |
new InliningCache(type()).dynamicInvoker() | |
); | |
setTarget(guard); | |
return target; | |
} | |
} | |
public MethodHandle createDispatch(Class<?>... signatureClasses) { | |
return new InliningCache(toMethodType(signatureClasses)).dynamicInvoker(); | |
} | |
// --- demo | |
static final class Demo { | |
interface Vehicle {} | |
record Car(int passenger) implements Vehicle { } | |
record CarrierTruck(Vehicle vehicle) implements Vehicle { } | |
static String toJSON(CarrierTruck truck, int depth, | |
@Signature({String.class, Vehicle.class, int.class}) MethodHandle dispatch) throws Throwable { | |
return STR.""" | |
{ | |
"vehicle" : \{ (String) dispatch.invokeExact(truck.vehicle, depth + 2) } | |
} | |
""".indent(depth); | |
} | |
static String toJSON(Car car, int depth) { | |
return STR.""" | |
{ | |
"passenger" : \{ car.passenger } | |
} | |
""".indent(depth); | |
} | |
} | |
private static final MethodHandle DISPATCH = DestructuredVisitor.of(lookup(), | |
Arrays.stream(Demo.class.getDeclaredMethods()) | |
.filter(m -> m.getName().equals("toJSON")) | |
.toList()) | |
.createDispatch(String.class, Demo.Vehicle.class, int.class); | |
public static void main(String[] args) throws Throwable { | |
var truck1 = new Demo.CarrierTruck(new Demo.Car(4)); | |
var result1 = (String) DISPATCH.invokeExact((Demo.Vehicle) truck1, 0); | |
System.out.println(result1); | |
var truck2 = new Demo.CarrierTruck(new Demo.CarrierTruck(new Demo.Car(7))); | |
var result2 = (String) DISPATCH.invokeExact((Demo.Vehicle) truck2, 0); | |
System.out.println(result2); | |
} | |
} |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment