DestructuredVisitor.java

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);
  }
}