HiddenClassLambdaTest.java

import java.lang.constant.ClassDesc;
import java.lang.constant.DirectMethodHandleDesc;
import java.lang.constant.DirectMethodHandleDesc.Kind;
import java.lang.constant.MethodHandleDesc;
import java.lang.invoke.CallSite;
import java.lang.invoke.ConstantCallSite;
import java.lang.invoke.LambdaMetafactory;
import java.lang.invoke.MethodHandle;
import java.lang.invoke.MethodHandles;
import java.lang.invoke.MethodType;
import java.lang.invoke.TypeDescriptor;

import org.objectweb.asm.ClassReader;
import org.objectweb.asm.ClassVisitor;
import org.objectweb.asm.ClassWriter;
import org.objectweb.asm.Handle;
import org.objectweb.asm.MethodVisitor;

import static java.lang.invoke.MethodType.methodType;
import static org.objectweb.asm.Opcodes.*;
import static java.lang.constant.ConstantDescs.*;

public class HiddenClassLambdaTest {
    /** This class is to be loaded and executed as hidden class */
    public static final class LambdaRunner implements Runnable {
        @Override public void run() {
            Runnable runnable = () -> System.out.println("Success");
            runnable.run();
        }
    }
    
    public static void main(String[] args) throws Throwable {
        // Path to the class file of the nested class defined above
        byte[] classFileContents = HiddenClassLambdaTest.class
                .getResourceAsStream("HiddenClassLambdaTest$LambdaRunner.class").readAllBytes();
        
        classFileContents = processLambdas(classFileContents);
        MethodHandles.Lookup hiddenClass = MethodHandles.lookup().defineHiddenClass(classFileContents, true);
        Runnable lambdaRunnerInstance = (Runnable) (
                hiddenClass.findConstructor(hiddenClass.lookupClass(), methodType(void.class))
                ).asType(methodType(Runnable.class)).invokeExact();
        lambdaRunnerInstance.run();
    }
    
    public static CallSite metafactory(MethodHandles.Lookup l, String name, MethodType mt,
            MethodType interfaceType, MethodHandle mh, MethodType dynamicMethodType) throws Throwable {
        MethodHandle invoker = MethodHandles.exactInvoker(mh.type());
        if (mt.parameterCount() == 0) {
            // Non-capturing lambda
            mt = mt.appendParameterTypes(MethodHandle.class);
            CallSite cs = LambdaMetafactory.metafactory(l, name, mt, interfaceType, invoker, dynamicMethodType);
            Object instance = cs.dynamicInvoker().asType(methodType(Object.class, MethodHandle.class)).invokeExact(mh);
            return new ConstantCallSite(MethodHandles.constant(mt.returnType(), instance));
        } else {
            // capturing
            MethodType lambdaMt = mt.insertParameterTypes(0, MethodHandle.class);
            CallSite cs = LambdaMetafactory.metafactory(l, name, lambdaMt, interfaceType, invoker, dynamicMethodType);
            return new ConstantCallSite(cs.dynamicInvoker().bindTo(mh));
        }
    }
    
    public static CallSite altMetafactory(MethodHandles.Lookup l, String name, MethodType mt, Object... args) {
        throw new UnsupportedOperationException("Not Implemented");
    }
    
    
    private static byte[] processLambdas(byte[] bytes) {
        ClassReader cr = new ClassReader(bytes);
        ClassWriter cw = new ClassWriter(cr, 0);
        
        ClassVisitor cv = cw;
        cv = new LambdaTransformer(cv);
        
        cr.accept(cv, 0);
        return cw.toByteArray();
    }
    
    private static class LambdaTransformer extends ClassVisitor {
        LambdaTransformer(ClassVisitor parent) {
            super(ASM9, parent);
        }
        
        @Override
        public MethodVisitor visitMethod(int access, String name, String descriptor,
                String signature, String[] exceptions) {
            return new LambdaMethodTransformer(super.visitMethod(access, name, descriptor, signature, exceptions));
        }
        
        private static class LambdaMethodTransformer extends MethodVisitor {
            public LambdaMethodTransformer(MethodVisitor parent) {
                super(ASM9, parent);
            }
            
            private static final ClassDesc CD_LambdaMetafactory = LambdaMetafactory.class.describeConstable().orElseThrow();
            private static final ClassDesc CD_HiddenLambdaTest = HiddenClassLambdaTest.class.describeConstable().orElseThrow();
            
            private static final DirectMethodHandleDesc LMF_FACTORY = ofCallsiteBootstrap(
                    CD_LambdaMetafactory, "metafactory", CD_CallSite, CD_MethodType, CD_MethodHandle, CD_MethodType);
            private static final DirectMethodHandleDesc LMF_ALTFACTORY = ofCallsiteBootstrap(
                    CD_LambdaMetafactory, "altMetafactory", CD_Object.arrayType());
            
            private static final Handle MY_FACTORY = toASM(ofCallsiteBootstrap(CD_HiddenLambdaTest, "metafactory", CD_CallSite, CD_MethodType, CD_MethodHandle, CD_MethodType));
            private static final Handle MY_ALTFACTORY = toASM(ofCallsiteBootstrap(CD_HiddenLambdaTest, "altMetafactory", CD_CallSite, CD_Object.arrayType()));
            
            @Override
            public void visitInvokeDynamicInsn(String name, String descriptor,
                    Handle bootstrapMethodHandle, Object... bootstrapMethodArguments) {
                MethodHandleDesc h = fromASM(bootstrapMethodHandle);
                if (h.equals(LMF_FACTORY)) {
                    super.visitInvokeDynamicInsn(name, descriptor,
                            MY_FACTORY, bootstrapMethodArguments);
                } else if (h.equals(LMF_ALTFACTORY)) {
                    super.visitInvokeDynamicInsn(name, descriptor, MY_ALTFACTORY, bootstrapMethodArguments);
                } else {
                    super.visitInvokeDynamicInsn(name, descriptor, bootstrapMethodHandle, bootstrapMethodArguments);
                }
            }
            
            private static MethodHandleDesc fromASM(Handle h) {
                return MethodHandleDesc.of(Kind.valueOf(h.getTag(), h.isInterface()),
                        ClassDesc.ofDescriptor("L" + h.getOwner() + ";"),
                        h.getName(), h.getDesc());
            }
            
            private static Handle toASM(DirectMethodHandleDesc desc) {
                return new Handle(desc.refKind(), toInternal(desc.owner()), desc.methodName(), desc.lookupDescriptor(), desc.isOwnerInterface());
            }
            
            private static String toInternal(TypeDescriptor.OfField<?> desc) {
                String d = desc.descriptorString();
                if (d.charAt(0) != 'L') {
                    throw new IllegalArgumentException("Not a valid internal type: " + d);
                }
                return d.substring(1, d.length() - 1); // Strip "L" + ";"
            }
        }
    }
}