Proxy 2.0 in one file

gistfile1.java

package java.lang.invoke;

import java.lang.invoke.MethodHandle;
import java.lang.invoke.MethodHandles.Lookup;
import java.lang.invoke.MethodType;
import java.lang.reflect.Method;
import java.util.Arrays;

import jdk.internal.org.objectweb.asm.ClassWriter;
import jdk.internal.org.objectweb.asm.FieldVisitor;
import jdk.internal.org.objectweb.asm.Handle;
import jdk.internal.org.objectweb.asm.MethodVisitor;
import jdk.internal.org.objectweb.asm.Type;
import sun.misc.Unsafe;
import static jdk.internal.org.objectweb.asm.Opcodes.*;

@SuppressWarnings("restriction")
public class Proxy2 {
  private static final Unsafe UNSAFE = Unsafe.getUnsafe();

  @FunctionalInterface
  public interface ProxyHandler {
    public default boolean overrideMethod(Method method) {
      return !method.isDefault();
    }
    public CallSite bootstrap(Lookup lookup, Method method) throws Throwable;
  }
  
  private static String internalName(Class<?> type) {
    return type.getName().replace('.', '/');
  }
  
  public static MethodHandle createProxyFactory(MethodType methodType, ProxyHandler handler) {
    Class<?> interfaze = methodType.returnType();
    
    ClassWriter writer = new ClassWriter(ClassWriter.COMPUTE_FRAMES|ClassWriter.COMPUTE_MAXS);
    writer.visit(V1_8, ACC_PUBLIC|ACC_SUPER, "java/lang/invoke/Foo", null, "java/lang/Object", new String[]{ internalName(interfaze) });
    
    String initDesc;
    {
      initDesc = methodType.changeReturnType(void.class).toMethodDescriptorString();
      MethodVisitor init = writer.visitMethod(ACC_PUBLIC, "<init>", initDesc, null, null);
      String factoryDesc = methodType.toMethodDescriptorString();
      MethodVisitor factory = writer.visitMethod(ACC_PUBLIC|ACC_STATIC, "0-^-0", factoryDesc, null, null);
      init.visitCode();
      factory.visitCode();
      factory.visitTypeInsn(NEW, "java/lang/invoke/Foo");
      factory.visitInsn(DUP);
      
      int slot = 1;
      for(int i = 0; i < methodType.parameterCount(); i++) {
        Class<?> boundType = methodType.parameterType(i);
        String fieldName = "arg" + i;
        FieldVisitor fv = writer.visitField(ACC_PRIVATE|ACC_FINAL, fieldName, Type.getDescriptor(boundType), null, null);
        fv.visitEnd();
        
        int loadOp = Type.getType(boundType).getOpcode(ILOAD);
        init.visitVarInsn(ALOAD, 0);
        init.visitVarInsn(loadOp, slot);
        init.visitFieldInsn(PUTFIELD, "java/lang/invoke/Foo", fieldName, Type.getDescriptor(boundType));
        
        factory.visitVarInsn(loadOp, slot - 1);
        
        slot += (boundType == long.class || boundType == double.class)? 2: 1;
      }
      
      init.visitInsn(RETURN);
      factory.visitMethodInsn(INVOKESPECIAL, "java/lang/invoke/Foo", "<init>", initDesc, false);
      factory.visitInsn(ARETURN);
      
      init.visitMaxs(-1, -1);
      init.visitEnd();
      factory.visitMaxs(-1, -1);
      factory.visitEnd();
    }
    
    String handlerPlaceHolder = "<<HANDLER_HOLDER>>";
    int handlerHolderCPIndex = writer.newConst(handlerPlaceHolder);
    
    Method[] methods = interfaze.getMethods();
    int[] methodHolderCPIndexes = new int[methods.length];
    for(int methodIndex = 0; methodIndex < methods.length; methodIndex++) {
      Method method = methods[methodIndex];
      if (method.isDefault() && !handler.overrideMethod(method)) {
        continue;
      }
      String methodDesc = Type.getMethodDescriptor(method);
      MethodVisitor mv = writer.visitMethod(ACC_PUBLIC, method.getName(), methodDesc, null,
          Arrays.stream(method.getExceptionTypes()).map(Proxy2::internalName).toArray(String[]::new));
      mv.visitCode();
      mv.visitVarInsn(ALOAD, 0);
      for(int i = 0; i < methodType.parameterCount(); i++) {
        Class<?> boundType = methodType.parameterType(i);
        mv.visitVarInsn(ALOAD, 0);
        mv.visitFieldInsn(GETFIELD, "java/lang/invoke/Foo", "arg" + i, Type.getDescriptor(boundType));
      }
      int slot = 1;
      for(Class<?> parameterType: method.getParameterTypes()) {
        mv.visitVarInsn(Type.getType(parameterType).getOpcode(ILOAD), slot);
        slot += (parameterType == long.class || parameterType == double.class)? 2: 1;
      }
      String methodPlaceHolder = "<<METHOD_HOLDER " + methodIndex + ">>";
      methodHolderCPIndexes[methodIndex] = writer.newConst(methodPlaceHolder);
      mv.visitInvokeDynamicInsn(method.getName(),
          "(Ljava/lang/Object;" + initDesc.substring(1, initDesc.length() - 2) + methodDesc.substring(1),
          BSM, handlerPlaceHolder, methodPlaceHolder);
      mv.visitInsn(Type.getReturnType(method).getOpcode(IRETURN));
      mv.visitMaxs(-1, -1);
      mv.visitEnd();
    }
    writer.visitEnd();
    byte[] data = writer.toByteArray();
    
    int constantPoolSize = writer.newConst("<<SENTINEL>>");
    
    Object[] patches = new Object[constantPoolSize];
    patches[handlerHolderCPIndex] = handler;
    for(int i = 0; i < methodHolderCPIndexes.length; i++) {
      patches[methodHolderCPIndexes[i]] = methods[i];
    }
    Class<?> clazz = UNSAFE.defineAnonymousClass(Proxy2.class, data, patches);
    try {
      return MethodHandles.publicLookup().findStatic(clazz, "0-^-0", methodType);
    } catch (NoSuchMethodException | IllegalAccessException e) {
      throw new AssertionError(e);
    }
  }
  
  private static final Handle BSM =
      new Handle(H_INVOKESTATIC, "java/lang/invoke/Proxy2", "bootstrap",
          MethodType.methodType(CallSite.class, Lookup.class, String.class, MethodType.class, ProxyHandler.class, Method.class).toMethodDescriptorString());
  
  // should be package-private but invokedynamic doesn't honor anonymous host class visibility 
  public static CallSite bootstrap(Lookup lookup, String name, MethodType methodType, ProxyHandler handler, Method method) throws Throwable {
    return handler.bootstrap(lookup, method);
  }
}

gistfile2.txt

import java.lang.invoke.ConstantCallSite;
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.Proxy2;
import java.lang.reflect.Method;
import java.util.function.Function;

public class Main {
  public static void main(String[] args) throws Throwable {
    MethodHandle identity = MethodHandles.identity(Object.class);
    
    MethodHandle mh = Proxy2.createProxyFactory(MethodType.methodType(Function.class), 
      (Lookup lookup, Method method) -> {
        System.out.println("require implementation of " + method);
        return new ConstantCallSite(MethodHandles.dropArguments(identity, 0, Object.class));
      });
    Function<Object,Object> fun = (Function<Object,Object>)mh.invokeExact();
    System.out.println(fun.apply("hello"));
    
    for(int i = 0; i < 1_000_000; i++) {
      fun.apply(i);
    }
  }
}

gistfile3.java

import java.lang.invoke.ConstantCallSite;
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.Proxy2;
import java.lang.reflect.Method;
import java.lang.reflect.UndeclaredThrowableException;
import java.util.function.IntUnaryOperator;

public class Main2 {
  private static final ClassValue<MethodHandle> FACTORY_CACHE = 
      new ClassValue<MethodHandle>() {
        @Override
        protected MethodHandle computeValue(Class<?> type) {
          return Proxy2.createProxyFactory(MethodType.methodType(type, type), 
              (Lookup lookup, Method method) -> {
                MethodHandle target = lookup.unreflect(method);
                return new ConstantCallSite(MethodHandles.dropArguments(target, 0, Object.class));
              });
        }
      };
  
  private static <T> T proxy(T delegate, Class<T> interfaze) {
    try {
      return interfaze.cast(FACTORY_CACHE.get(interfaze).invoke(delegate));
    } catch(RuntimeException | Error e) {
      throw e;
    } catch (Throwable e) {
      throw new UndeclaredThrowableException(e);
    }
  }
  
  public static void main(String[] args) {
    IntUnaryOperator delegate = x -> x * 2;
    IntUnaryOperator op = proxy(delegate, IntUnaryOperator.class);
    System.out.println(op.applyAsInt(3));
  }
}