Added further advice implementation.

raphw · Mar 4, 2016 · b209c41 · b209c41
1 parent 23529bc
commit b209c41
Showing 1 changed file with 33 additions and 17 deletions.
diff --git a/byte-buddy-dep/src/main/java/net/bytebuddy/asm/advice/Advice.java b/byte-buddy-dep/src/main/java/net/bytebuddy/asm/advice/Advice.java
@@ -148,11 +148,20 @@ public MethodVisitor visitMethod(int modifiers, String internalName, String desc
 
             protected class TransferingVisitor extends MethodVisitor {
 
-                private final MethodDescription inlinedMethod;
+                private final Map<Integer, ParameterDescription> parameters;
+
+                private final int offsetCorrection;
+
+                private final Label endOfMethod;
 
                 protected TransferingVisitor(MethodVisitor methodVisitor, MethodDescription inlinedMethod) {
                     super(Opcodes.ASM5, methodVisitor);
-                    this.inlinedMethod = inlinedMethod;
+                    parameters = new HashMap<Integer, ParameterDescription>();
+                    for (ParameterDescription parameter : inlinedMethod.getParameters()) {
+                        parameters.put(parameter.getOffset(), parameter);
+                    }
+                    offsetCorrection = instrumentedMethod.getStackSize() - inlinedMethod.getStackSize();
+                    endOfMethod = new Label();
                 }
 
                 @Override
@@ -204,44 +213,50 @@ public void visitEnd() {
                 public void visitMaxs(int maxStack, int maxLocals) {
                     AsmAdvice.this.maxStack = maxStack + stackIncrement;
                     AsmAdvice.this.maxLocals = maxLocals;
+                    super.visitLabel(endOfMethod);
                 }
 
                 @Override
                 public void visitInsn(int opcode) {
                     switch (opcode) {
                         case Opcodes.RETURN:
+                            super.visitJumpInsn(Opcodes.GOTO, endOfMethod);
                             break;
                         case Opcodes.IRETURN:
                         case Opcodes.FRETURN:
-                        case Opcodes.ARETURN:
                             super.visitInsn(Opcodes.POP);
+                            super.visitJumpInsn(Opcodes.GOTO, endOfMethod);
                             break;
                         case Opcodes.LRETURN:
                         case Opcodes.DRETURN:
                             super.visitInsn(Opcodes.POP2);
+                            super.visitJumpInsn(Opcodes.GOTO, endOfMethod);
                             break;
+                        case Opcodes.ARETURN:
                         default:
                             super.visitInsn(opcode);
                     }
                 }
 
                 @Override
-                public void visitVarInsn(int opcode, int variableOffset) {
+                public void visitVarInsn(int opcode, int offset) {
-                    for (ParameterDescription parameter : inlinedMethod.getParameters()) {
+                    ParameterDescription parameterDescription = parameters.get(offset);
-                        int parameterOffset = parameter.getOffset();
+                    if (parameterDescription != null) {
-                        if (parameterOffset == variableOffset) {
+                        AnnotationDescription.Loadable<Argument> argument = parameterDescription.getDeclaredAnnotations().ofType(Argument.class);
-                            AnnotationDescription.Loadable<Argument> argument = parameter.getDeclaredAnnotations().ofType(Argument.class);
+                        int targetIndex = argument == null
-                            int instrumentedIndex = argument == null
+                                ? parameterDescription.getIndex()
-                                    ? parameter.getIndex()
+                                : argument.loadSilent().value();
-                                    : argument.loadSilent().value();
+                        ParameterList<?> targetParameters = instrumentedMethod.getParameters();
-                            ParameterList<?> instrumentedParameters = instrumentedMethod.getParameters();
+                        if (targetIndex >= targetParameters.size()) {
-                            if (instrumentedIndex > instrumentedParameters.size()) {
+                            throw new IllegalStateException(instrumentedMethod + " does not define a parameter with index " + targetIndex);
-                                throw new IllegalStateException(instrumentedMethod + " does not define a parameter with index " + instrumentedIndex);
+                        } else if (!targetParameters.get(targetIndex).getType().asErasure().isAssignableTo(parameterDescription.getType().asErasure())) {
-                            }
+                            throw new IllegalStateException("Cannot assign " + targetParameters.get(targetIndex) + " to " + parameterDescription);
-                            variableOffset = instrumentedParameters.get(instrumentedIndex).getOffset();
                         }
+                        offset = targetParameters.get(targetIndex).getOffset();
+                    } else {
+                        offset += offsetCorrection;
                     }
-                    super.visitVarInsn(opcode, variableOffset);
+                    super.visitVarInsn(opcode, offset);
                 }
             }
         }
@@ -265,6 +280,7 @@ public void visitVarInsn(int opcode, int variableOffset) {
     @Retention(RetentionPolicy.RUNTIME)
     @Target(ElementType.PARAMETER)
     public @interface Argument {
+
         int value();
     }
 }