From 5c86baa171053ea49990d0959726fdc401d11aa7 Mon Sep 17 00:00:00 2001
From: Rafael Winterhalter <rafael.wth@gmail.com>
Date: Fri, 29 Apr 2016 15:48:54 +0200
Subject: [PATCH] Added method call hook to add all arguments of an
 instrumented method to an invocation.

---
 .../bytebuddy/implementation/MethodCall.java  | 45 +++++++++++++++++++
 .../implementation/MethodCallTest.java        | 18 +++++++-
 2 files changed, 62 insertions(+), 1 deletion(-)

diff --git a/byte-buddy-dep/src/main/java/net/bytebuddy/implementation/MethodCall.java b/byte-buddy-dep/src/main/java/net/bytebuddy/implementation/MethodCall.java
index c9a36b8c51d..88b2eb6390f 100644
--- a/byte-buddy-dep/src/main/java/net/bytebuddy/implementation/MethodCall.java
+++ b/byte-buddy-dep/src/main/java/net/bytebuddy/implementation/MethodCall.java
@@ -365,6 +365,21 @@ public MethodCall withArgument(int... index) {
                 typing);
     }
 
+    /**
+     * Adds all arguments of the instrumented method as arguments to the invoked method to this method call.
+     *
+     * @return A method call that hands all arguments arguments of the instrumented method to the invoked method.
+     */
+    public MethodCall withAllArguments() {
+        return new MethodCall(methodLocator,
+                targetHandler,
+                CompoundList.of(this.argumentLoaders, ArgumentLoader.ForMethodParameter.OfInstrumentedMethod.INSTANCE),
+                methodInvoker,
+                terminationHandler,
+                assigner,
+                typing);
+    }
+
     /**
      * Assigns the {@code this} reference to the next parameter.
      *
@@ -1184,6 +1199,36 @@ public String toString() {
                         '}';
             }
 
+            /**
+             * A factory for argument loaders that supplies all arguments of the instrumented method as arguments.
+             */
+            protected enum OfInstrumentedMethod implements ArgumentLoader.Factory {
+
+                /**
+                 * The singleton instance.
+                 */
+                INSTANCE;
+
+                @Override
+                public InstrumentedType prepare(InstrumentedType instrumentedType) {
+                    return instrumentedType;
+                }
+
+                @Override
+                public List<ArgumentLoader> make(TypeDescription instrumentedType, MethodDescription instrumentedMethod) {
+                    List<ArgumentLoader> argumentLoaders = new ArrayList<ArgumentLoader>(instrumentedMethod.getParameters().size());
+                    for (ParameterDescription parameterDescription : instrumentedMethod.getParameters()) {
+                        argumentLoaders.add(new ForMethodParameter(parameterDescription.getIndex(), instrumentedMethod));
+                    }
+                    return argumentLoaders;
+                }
+
+                @Override
+                public String toString() {
+                    return "MethodCall.ArgumentLoader.ForMethodParameter.OfInstrumentedMethod." + name();
+                }
+            }
+
             /**
              * A factory for an argument loader that supplies a method parameter as an argument.
              */
diff --git a/byte-buddy-dep/src/test/java/net/bytebuddy/implementation/MethodCallTest.java b/byte-buddy-dep/src/test/java/net/bytebuddy/implementation/MethodCallTest.java
index 0095c5935b5..ca7bdb1ccf9 100644
--- a/byte-buddy-dep/src/test/java/net/bytebuddy/implementation/MethodCallTest.java
+++ b/byte-buddy-dep/src/test/java/net/bytebuddy/implementation/MethodCallTest.java
@@ -254,7 +254,7 @@ public void testWithExplicitArgumentFieldNonAssignable() throws Exception {
     }
 
     @Test
-    public void testWithParameter() throws Exception {
+    public void testWithArgument() throws Exception {
         DynamicType.Loaded<MethodCallWithExplicitArgument> loaded = implement(MethodCallWithExplicitArgument.class,
                 MethodCall.invokeSuper().withArgument(0));
         assertThat(loaded.getLoadedAuxiliaryTypes().size(), is(0));
@@ -268,6 +268,21 @@ public void testWithParameter() throws Exception {
         assertThat(instance.foo(BAR), is(BAR));
     }
 
+    @Test
+    public void testWithAllArguments() throws Exception {
+        DynamicType.Loaded<MethodCallWithExplicitArgument> loaded = implement(MethodCallWithExplicitArgument.class,
+                MethodCall.invokeSuper().withAllArguments());
+        assertThat(loaded.getLoadedAuxiliaryTypes().size(), is(0));
+        assertThat(loaded.getLoaded().getDeclaredMethods().length, is(1));
+        assertThat(loaded.getLoaded().getDeclaredMethod(FOO, String.class), not(nullValue(Method.class)));
+        assertThat(loaded.getLoaded().getDeclaredConstructors().length, is(1));
+        assertThat(loaded.getLoaded().getDeclaredFields().length, is(0));
+        MethodCallWithExplicitArgument instance = loaded.getLoaded().newInstance();
+        assertThat(instance.getClass(), not(CoreMatchers.<Class<?>>is(MethodCallWithExplicitArgument.class)));
+        assertThat(instance, instanceOf(MethodCallWithExplicitArgument.class));
+        assertThat(instance.foo(BAR), is(BAR));
+    }
+
     @Test
     public void testWithInstanceField() throws Exception {
         DynamicType.Loaded<MethodCallWithExplicitArgument> loaded = implement(MethodCallWithExplicitArgument.class,
@@ -653,6 +668,7 @@ public Class<?> create() {
         ObjectPropertyAssertion.of(MethodCall.ArgumentLoader.ForLongConstant.class).apply();
         ObjectPropertyAssertion.of(MethodCall.ArgumentLoader.ForMethodParameter.class).apply();
         ObjectPropertyAssertion.of(MethodCall.ArgumentLoader.ForMethodParameter.Factory.class).apply();
+        ObjectPropertyAssertion.of(MethodCall.ArgumentLoader.ForMethodParameter.OfInstrumentedMethod.class).apply();
         ObjectPropertyAssertion.of(MethodCall.ArgumentLoader.ForShortConstant.class).apply();
         ObjectPropertyAssertion.of(MethodCall.ArgumentLoader.ForTextConstant.class).apply();
         ObjectPropertyAssertion.of(MethodCall.ArgumentLoader.ForClassConstant.class).apply();