Added method mapper

src/main/java/net/lenni0451/reflect/proxy/ProxyBuilder.java

@@ -16,6 +16,7 @@
 import java.lang.reflect.Field;
 import java.lang.reflect.Method;
 import java.util.Arrays;
+import java.util.function.Function;
 import java.util.function.Predicate;
 
 import static net.lenni0451.reflect.bytecode.BytecodeUtils.*;
@@ -35,11 +36,20 @@ public class ProxyBuilder {
     @Nullable
     private Class<?>[] interfaces;
     private Predicate<Method> methodFilter = m -> true;
+    private Function<Method, Method> methodMapper = Function.identity();
     private InvocationHandler invocationHandler = InvocationHandler.forwarding();
     private ProxyClassDefiner classDefiner = ProxyClassDefiner.loader(ProxyBuilder.class.getClassLoader());
 
     private Class<?> proxyClass;
 
+    /**
+     * @return The super class of the proxy class
+     */
+    @Nullable
+    public Class<?> getSuperClass() {
+        return this.superClass;
+    }
+
     /**
      * Set the super class of the proxy class.<br>
      * The super class must be public and not final.
@@ -56,11 +66,11 @@ public ProxyBuilder setSuperClass(@Nullable final Class<?> superClass) {
     }
 
     /**
-     * @return The super class of the proxy class
+     * @return The interfaces of the proxy class
      */
     @Nullable
-    public Class<?> getSuperClass() {
-        return this.superClass;
+    public Class<?>[] getInterfaces() {
+        return this.interfaces;
     }
 
     /**
@@ -101,15 +111,16 @@ public ProxyBuilder setInterfaces(@Nullable final Class<?>... interfaces) {
     }
 
     /**
-     * @return The interfaces of the proxy class
+     * @return The current method filter
      */
-    @Nullable
-    public Class<?>[] getInterfaces() {
-        return this.interfaces;
+    public Predicate<Method> getMethodFilter() {
+        return this.methodFilter;
     }
 
     /**
-     * Set the filter for the methods that should be overridden by the proxy class.
+     * Set the filter for the methods that should be overridden by the proxy class.<br>
+     * Filtered methods will not be overridden by the proxy class.<br>
+     * Abstract methods will always be overridden and cannot be filtered.
      *
      * @param methodFilter The method filter
      * @return This builder
@@ -121,21 +132,23 @@ public ProxyBuilder setMethodFilter(@Nonnull final Predicate<Method> methodFilte
     }
 
     /**
-     * @return The current method filter
+     * @return The current method mapper
      */
-    public Predicate<Method> getMethodFilter() {
-        return this.methodFilter;
+    public Function<Method, Method> getMethodMapper() {
+        return this.methodMapper;
     }
 
     /**
-     * Set the invocation handler for the proxy class.
+     * Set the mapper for the methods that should be overridden by the proxy class.<br>
+     * The mapper can be used to change the owner of an overridden method.<br>
+     * The owner must be a super class of the original owner.
      *
-     * @param invocationHandler The invocation handler
+     * @param methodMapper The method mapper
      * @return This builder
      */
-    public ProxyBuilder setInvocationHandler(@Nonnull final InvocationHandler invocationHandler) {
+    public ProxyBuilder setMethodMapper(@Nonnull final Function<Method, Method> methodMapper) {
         this.reset();
-        this.invocationHandler = invocationHandler;
+        this.methodMapper = methodMapper;
         return this;
     }
 
@@ -147,14 +160,14 @@ public InvocationHandler getInvocationHandler() {
     }
 
     /**
-     * Set the class definer for the proxy class.
+     * Set the invocation handler for the proxy class.
      *
-     * @param classDefiner The class definer
+     * @param invocationHandler The invocation handler
      * @return This builder
      */
-    public ProxyBuilder setClassDefiner(@Nonnull final ProxyClassDefiner classDefiner) {
+    public ProxyBuilder setInvocationHandler(@Nonnull final InvocationHandler invocationHandler) {
         this.reset();
-        this.classDefiner = classDefiner;
+        this.invocationHandler = invocationHandler;
         return this;
     }
 
@@ -165,6 +178,18 @@ public ProxyClassDefiner getClassDefiner() {
         return this.classDefiner;
     }
 
+    /**
+     * Set the class definer for the proxy class.
+     *
+     * @param classDefiner The class definer
+     * @return This builder
+     */
+    public ProxyBuilder setClassDefiner(@Nonnull final ProxyClassDefiner classDefiner) {
+        this.reset();
+        this.classDefiner = classDefiner;
+        return this;
+    }
+
     /**
      * Build the proxy class.
      *
@@ -173,21 +198,22 @@ public ProxyClassDefiner getClassDefiner() {
     public ProxyClass build() {
         if (this.proxyClass == null) {
             Reference<Method[]> methodsReference = new Reference<>();
-            BuiltClass builtClass = this.buildClass(methodsReference);
+            Reference<Method[]> originalMethodsReference = new Reference<>();
+            BuiltClass builtClass = this.buildClass(methodsReference, originalMethodsReference);
             this.proxyClass = this.classDefiner.defineProxyClass(builtClass, this.superClass, this.interfaces);
 
             //Set the static fields
             Field methods = Fields.getDeclaredField(this.proxyClass, METHODS_FIELD);
             Fields.setObject(null, methods, methodsReference.value);
 
-            Class<ProxyMethod>[] proxyMethodClasses = this.buildProxyMethodClasses(methodsReference.value);
+            Class<ProxyMethod>[] proxyMethodClasses = this.buildProxyMethodClasses(methodsReference.value, originalMethodsReference.value);
             Field proxyMethodClassesField = Fields.getDeclaredField(this.proxyClass, PROXY_METHOD_CLASSES_FIELD);
             Fields.setObject(null, proxyMethodClassesField, proxyMethodClasses);
         }
         return new ProxyClass(this.proxyClass, this.invocationHandler);
     }
 
-    private BuiltClass buildClass(final Reference<Method[]> methodsReference) {
+    private BuiltClass buildClass(final Reference<Method[]> methodsReference, final Reference<Method[]> originalMethodsReference) {
         String className = "net/lenni0451/reflect/proxy/ProxyImpl$" + System.nanoTime();
         Class<?>[] interfaces = this.interfaces;
         if (interfaces == null) {
@@ -208,6 +234,7 @@ private BuiltClass buildClass(final Reference<Method[]> methodsReference) {
 
                     Method[] methods = ProxyUtils.getOverridableMethod(this.superClass, this.interfaces, this.methodFilter);
                     methodsReference.value = methods;
+                    originalMethodsReference.value = ProxyUtils.mapMethods(methods, this.methodMapper);
                     this.addFields(cb, methods);
                     this.addMethods(cb, methods);
                     this.addDefaultMethods(cb);
@@ -329,10 +356,10 @@ private void addDefaultMethods(final ClassBuilder cb) {
         });
     }
 
-    private Class<ProxyMethod>[] buildProxyMethodClasses(final Method[] methods) {
+    private Class<ProxyMethod>[] buildProxyMethodClasses(final Method[] methods, final Method[] originalMethods) {
         Class<ProxyMethod>[] proxyMethodClasses = new Class[methods.length];
         for (int i = 0; i < methods.length; i++) {
-            proxyMethodClasses[i] = ProxyMethodBuilder.buildProxyMethodClass(this.proxyClass, methods[i]);
+            proxyMethodClasses[i] = ProxyMethodBuilder.buildProxyMethodClass(this.proxyClass, methods[i], originalMethods[i]);
         }
         return proxyMethodClasses;
     }

src/main/java/net/lenni0451/reflect/proxy/ProxyMethodBuilder.java

@@ -7,6 +7,7 @@
 import net.lenni0451.reflect.proxy.impl.ProxyMethod;
 import net.lenni0451.reflect.proxy.impl.ProxyRuntime;
 import org.jetbrains.annotations.ApiStatus;
+import org.jetbrains.annotations.Nullable;
 
 import java.lang.invoke.MethodHandle;
 import java.lang.reflect.Method;
@@ -22,14 +23,14 @@ class ProxyMethodBuilder {
 
     private static final BytecodeBuilder BUILDER = BytecodeBuilder.get();
 
-    public static Class<ProxyMethod> buildProxyMethodClass(final Class<?> proxyClass, final Method method) {
+    public static Class<ProxyMethod> buildProxyMethodClass(final Class<?> proxyClass, final Method method, @Nullable final Method originalMethod) {
         BuiltClass builtClass = BUILDER.class_(BUILDER.opcode("ACC_PUBLIC"), slash(proxyClass) + "$ProxyMethodImpl", null, slash(Object.class), new String[]{slash(ProxyMethod.class)}, cb -> {
             addFields(proxyClass, cb);
-            addStaticBlock(method, cb);
+            addStaticBlock(method, originalMethod, cb);
             addConstructor(proxyClass, cb);
             addMethodGetter(cb);
             addInvokeWith(method, cb);
-            addInvokeSuper(proxyClass, method, cb);
+            addInvokeSuper(proxyClass, originalMethod == null ? method : originalMethod, cb);
             addCancel(method, cb);
         });
 
@@ -43,10 +44,15 @@ private static void addFields(final Class<?> proxyClass, final ClassBuilder cb)
         cb.field(BUILDER.opcode("ACC_PRIVATE", "ACC_FINAL"), "method", desc(Method.class), null, null);
     }
 
-    private static void addStaticBlock(final Method method, final ClassBuilder cb) {
+    private static void addStaticBlock(final Method method, @Nullable final Method originalMethod, final ClassBuilder cb) {
         cb.method(BUILDER.opcode("ACC_PUBLIC", "ACC_STATIC"), "<clinit>", mdesc(void.class), null, null, mb -> {
+            mb.typeLdc(BUILDER, method.getDeclaringClass()); //Instance class
+            if (originalMethod != null) {
+                mb.typeLdc(BUILDER, originalMethod.getDeclaringClass()); //Super class
+            } else {
+                mb.insn(BUILDER.opcode("DUP"));
+            }
             mb
-                    .typeLdc(BUILDER, method.getDeclaringClass()) //Instance class
                     .ldc(method.getName()) //Method name
                     .intPush(BUILDER, method.getParameterCount())
                     .type(BUILDER.opcode("ANEWARRAY"), slash(Class.class)); //Parameter array
@@ -60,7 +66,7 @@ private static void addStaticBlock(final Method method, final ClassBuilder cb) {
             }
             mb
                     .typeLdc(BUILDER, method.getReturnType()) //Return type
-                    .method(BUILDER.opcode("INVOKESTATIC"), slash(ProxyRuntime.class), "getMethodHandles", mdesc(MethodHandle[].class, Class.class, String.class, Class[].class, Class.class), false);
+                    .method(BUILDER.opcode("INVOKESTATIC"), slash(ProxyRuntime.class), "getMethodHandles", mdesc(MethodHandle[].class, Class.class, Class.class, String.class, Class[].class, Class.class), false);
             mb //Store the method handles
                     .insn(BUILDER.opcode("DUP"))
                     .intPush(BUILDER, 0)

src/main/java/net/lenni0451/reflect/proxy/ProxyUtils.java

@@ -12,6 +12,7 @@
 import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
+import java.util.function.Function;
 import java.util.function.Predicate;
 
 /**
@@ -64,4 +65,20 @@ public static void getOverridableMethod(final Class<?> clazz, final Map<String,
         for (Class<?> inter : clazz.getInterfaces()) getOverridableMethod(inter, methods);
     }
 
+    public static Method[] mapMethods(final Method[] methods, final Function<Method, Method> methodMapper) {
+        Method[] originalMethods = new Method[methods.length];
+        for (int i = 0; i < methods.length; i++) {
+            Method method = methods[i];
+            Method mappedMethod = methodMapper.apply(method);
+            if (method.equals(mappedMethod)) continue;
+
+            if (!mappedMethod.getDeclaringClass().isAssignableFrom(method.getDeclaringClass())) {
+                throw new IllegalArgumentException("The method '" + method.getName() + "' in class '" + method.getDeclaringClass().getName() + "' is not assignable to the method in the proxy class");
+            }
+            methods[i] = mappedMethod;
+            originalMethods[i] = method;
+        }
+        return originalMethods;
+    }
+
 }

src/main/java/net/lenni0451/reflect/proxy/impl/ProxyRuntime.java

@@ -28,11 +28,11 @@ public class ProxyRuntime {
      * @throws NoSuchMethodException  If the method does not exist
      * @throws IllegalAccessException If the method is not accessible
      */
-    public static MethodHandle[] getMethodHandles(final Class<?> owner, final String name, final Class<?>[] parameters, final Class<?> returnType) throws NoSuchMethodException, IllegalAccessException {
+    public static MethodHandle[] getMethodHandles(final Class<?> owner, final Class<?> superOwner, final String name, final Class<?>[] parameters, final Class<?> returnType) throws NoSuchMethodException, IllegalAccessException {
         MethodHandle[] methodHandles = new MethodHandle[2];
         methodHandles[0] = TRUSTED_LOOKUP.findVirtual(owner, name, MethodType.methodType(returnType, parameters));
         try {
-            methodHandles[1] = TRUSTED_LOOKUP.findSpecial(owner, name, MethodType.methodType(returnType, parameters), owner);
+            methodHandles[1] = TRUSTED_LOOKUP.findSpecial(superOwner, name, MethodType.methodType(returnType, parameters), superOwner);
         } catch (Throwable ignored) {
         }
         return methodHandles;

src/test/java/net/lenni0451/reflect/proxy/ProxyTest.java

@@ -1,10 +1,13 @@
 package net.lenni0451.reflect.proxy;
 
+import net.lenni0451.reflect.Methods;
 import net.lenni0451.reflect.proxy.impl.Proxy;
 import net.lenni0451.reflect.proxy.test.*;
 import org.junit.jupiter.api.BeforeEach;
 import org.junit.jupiter.api.Test;
 
+import java.lang.reflect.Method;
+
 import static org.junit.jupiter.api.Assertions.*;
 
 class ProxyTest {
@@ -99,4 +102,42 @@ void testSuperImplementsProxy() {
         assertDoesNotThrow(instance::get);
     }
 
+    @Test
+    void testInvokeWrongObject() {
+        ProxyClass proxyClass = new ProxyBuilder()
+                .setSuperClass(Class2.class)
+                .setInvocationHandler((thiz, proxyMethod, args) -> proxyMethod.invokeWith(new Class1(), args))
+                .build();
+        Class2 proxy = proxyClass.allocateInstance();
+        //The method found by the proxy builder belongs to Class2 (the proxy super class)
+        //Invoking it with an instance of Class1 should throw an exception because Class1 does not extend Class2, even tho it is the other way round
+        //The cast is done by the proxy to allow calling invokeExact of the method handle
+        assertThrows(ClassCastException.class, proxy::test);
+    }
+
+    @Test
+    void testMethodMapper() {
+        ProxyClass proxyClass = new ProxyBuilder()
+                .setSuperClass(Class2.class)
+                .setMethodMapper(method -> {
+                    Method class1Method = Methods.getDeclaredMethod(Class1.class, method.getName(), method.getParameterTypes());
+                    if (class1Method != null) return class1Method;
+                    return method;
+                })
+                .setInvocationHandler((thiz, proxyMethod, args) -> proxyMethod.invokeWith(new Class1(), args))
+                .build();
+        Class2 proxy = proxyClass.allocateInstance();
+        //The invocation is the same as in testInvokeWrongObject but this time the method mapper maps the method to one owned by Class1
+        //This should work because Class2 extends Class1 and the correct method should be invoked
+        assertEquals(1, proxy.test());
+
+        ((Proxy) proxy).setInvocationHandler(InvocationHandler.forwarding());
+        //Here it should return 2 because the super method is still the one from Class2
+        assertEquals(2, proxy.test());
+
+        ((Proxy) proxy).setInvocationHandler((thiz, proxyMethod, args) -> proxyMethod.invokeWith(new Class2(), args));
+        //This should also work because it's forwarded to Class2
+        assertEquals(2, proxy.test());
+    }
+
 }