Snap for 4662252 from a8266e2df1957adfa83283707bfdc821c591f5c3 to pi-release

Change-Id: Ib2aac92ecc6d897c453eade2601ebbdc0b08ab58

21 files changed, 1383 insertions, 784 deletions

diff --git a/Android.bp b/Android.bp
index 245d3e8..61b24ec 100644
--- a/Android.bp
+++ b/Android.bp
@@ -41,7 +41,7 @@ java_library_static {
 // Build dispatcher for Dexmaker's inline MockMaker
 java_library_static {
     name: "dexmaker-inline-mockmaker-dispatcher",
-    sdk_version: "25",
+    sdk_version: "current",
     srcs: ["dexmaker-mockito-inline-dispatcher/src/main/java/**/*.java"],
 }
 
@@ -107,7 +107,7 @@ cc_library_shared {
 // Build Dexmaker's inline MockMaker, a plugin to Mockito
 java_library_static {
     name: "dexmaker-inline-mockmaker",
-    sdk_version: "25",
+    sdk_version: "current",
     srcs: ["dexmaker-mockito-inline/src/main/java/**/*.java"],
     java_resource_dirs: ["dexmaker-mockito-inline/src/main/resources"],
     libs: [
@@ -115,6 +115,12 @@ java_library_static {
         "mockito-api",
     ],
     required: ["libdexmakerjvmtiagent"],
+
+    errorprone: {
+        javacflags: [
+            "-Xep:CollectionIncompatibleType:WARN"
+        ],
+    }
 }
 
 java_import {
diff --git a/README.version b/README.version
index 55abaa3..2d91670 100644
--- a/README.version
+++ b/README.version
@@ -1,5 +1,5 @@
 URL: https://github.com/linkedin/dexmaker/
-Version: master (fce01046a9519f8c1e5fd826fe5169eb600710ad)
+Version: master (5fb49bba98647d7a0aeea0cbf91fd670c3ff552a)
 License: Apache 2.0
 Description:
 Dexmaker is a Java-language API for doing compile time or runtime code generation targeting the Dalvik VM. Unlike cglib or ASM, this library creates Dalvik .dex files instead of Java .class files.
@@ -10,5 +10,3 @@ It includes a stock code generator for class proxies. If you just want to do AOP
 
 Local Modifications:
         Allow to share classloader via dexmaker.share_classloader system property (I8c2490c3ec8e8582dc41c486f8f7a406bd635ebb)
-        Dynamically register transformation hook (I2556e749806bed4f80697d5cba38a8d8a2f7ce6a)
-	Use new attach API (I2be3ac3218c11f2ccf9a675fffd35f4fab548070)
diff --git a/dexmaker-mockito-inline-tests/src/androidTest/java/com/android/dx/mockito/inline/tests/MultipleJvmtiAgentsInterference.java b/dexmaker-mockito-inline-tests/src/androidTest/java/com/android/dx/mockito/inline/tests/MultipleJvmtiAgentsInterference.java
index 9231202..bdc0626 100644
--- a/dexmaker-mockito-inline-tests/src/androidTest/java/com/android/dx/mockito/inline/tests/MultipleJvmtiAgentsInterference.java
+++ b/dexmaker-mockito-inline-tests/src/androidTest/java/com/android/dx/mockito/inline/tests/MultipleJvmtiAgentsInterference.java
@@ -17,27 +17,18 @@
 package com.android.dx.mockito.inline.tests;
 
 import android.os.Build;
+import android.os.Debug;
 
+import org.junit.AfterClass;
 import org.junit.BeforeClass;
-import org.junit.Ignore;
 import org.junit.Test;
 
-import java.io.File;
-import java.io.FileInputStream;
-import java.io.FileOutputStream;
-import java.io.InputStream;
-import java.io.OutputStream;
-
-import dalvik.system.BaseDexClassLoader;
-
 import static org.junit.Assert.assertNull;
 import static org.junit.Assume.assumeTrue;
 import static org.mockito.Mockito.mock;
 
-@Ignore("Leaving the transformation hook enabled causes a lot of unnecessary transformations. " +
-        "This is too expensive. Hence for now, we cannot support multiple agents")
 public class MultipleJvmtiAgentsInterference {
-    private static final String AGENT_LIB_NAME = "multiplejvmtiagentsinterferenceagent";
+    private static final String AGENT_LIB_NAME = "libmultiplejvmtiagentsinterferenceagent.so";
 
     public class TestClass {
         public String returnA() {
@@ -50,29 +41,8 @@ public class MultipleJvmtiAgentsInterference {
         // TODO (moltmann@google.com): Replace with proper check for >= P
         assumeTrue(Build.VERSION.CODENAME.equals("P"));
 
-        // Currently Debug.attachJvmtiAgent requires a file in the right directory
-        File copiedAgent = File.createTempFile("testagent", ".so");
-        copiedAgent.deleteOnExit();
-
-        try (InputStream is = new FileInputStream(((BaseDexClassLoader)
-                MultipleJvmtiAgentsInterference.class.getClassLoader()).findLibrary
-                (AGENT_LIB_NAME))) {
-            try (OutputStream os = new FileOutputStream(copiedAgent)) {
-                byte[] buffer = new byte[64 * 1024];
-
-                while (true) {
-                    int numRead = is.read(buffer);
-                    if (numRead == -1) {
-                        break;
-                    }
-                    os.write(buffer, 0, numRead);
-                }
-            }
-        }
-
-        // TODO (moltmann@google.com): Replace with regular method call once the API becomes public
-        Class.forName("android.os.Debug").getMethod("attachJvmtiAgent", String.class, String
-                .class).invoke(null, copiedAgent.getAbsolutePath(), null);
+        Debug.attachJvmtiAgent(AGENT_LIB_NAME, null,
+                MultipleJvmtiAgentsInterference.class.getClassLoader());
     }
 
     @Test
@@ -90,5 +60,11 @@ public class MultipleJvmtiAgentsInterference {
         assertNull(t.returnA());
     }
 
+    @AfterClass
+    public static void DisableRetransfromHook() {
+        disableRetransformHook();
+    }
+
     private native int nativeRetransformClasses(Class<?>[] classes);
+    private static native int disableRetransformHook();
 }
diff --git a/dexmaker-mockito-inline-tests/src/main/jni/multiplejvmtiagentsinterferenceagent/agent.cc b/dexmaker-mockito-inline-tests/src/main/jni/multiplejvmtiagentsinterferenceagent/agent.cc
index a293fe7..b1c5455 100644
--- a/dexmaker-mockito-inline-tests/src/main/jni/multiplejvmtiagentsinterferenceagent/agent.cc
+++ b/dexmaker-mockito-inline-tests/src/main/jni/multiplejvmtiagentsinterferenceagent/agent.cc
@@ -22,9 +22,9 @@
 
 #include "jvmti.h"
 
-#include <dex_ir.h>
-#include <writer.h>
-#include <reader.h>
+#include <slicer/dex_ir.h>
+#include <slicer/writer.h>
+#include <slicer/reader.h>
 
 using namespace dex;
 
@@ -148,4 +148,15 @@ namespace com_android_dx_mockito_inline_tests {
 
         return error;
     }
-}
\ No newline at end of file
+
+    // Disable hook to not slow down test
+    extern "C" JNIEXPORT jint JNICALL
+    Java_com_android_dx_mockito_inline_tests_MultipleJvmtiAgentsInterference_disableRetransformHook(
+            JNIEnv *env,
+            jclass ignored) {
+        return localJvmtiEnv->SetEventNotificationMode(JVMTI_DISABLE,
+                                                       JVMTI_EVENT_CLASS_FILE_LOAD_HOOK,
+                                                       NULL);
+
+    }
+}
diff --git a/dexmaker-mockito-inline/CMakeLists.txt b/dexmaker-mockito-inline/CMakeLists.txt
index cd26d58..b700cd4 100644
--- a/dexmaker-mockito-inline/CMakeLists.txt
+++ b/dexmaker-mockito-inline/CMakeLists.txt
@@ -20,7 +20,7 @@ add_library(slicer
             STATIC
             ${slicer_sources})
 
-include_directories(external/jdk external/slicer/)
+include_directories(external/jdk external/slicer/export/)
 
 target_link_libraries(slicer z)
 
diff --git a/dexmaker-mockito-inline/build.gradle b/dexmaker-mockito-inline/build.gradle
index 54e85ec..853f061 100644
--- a/dexmaker-mockito-inline/build.gradle
+++ b/dexmaker-mockito-inline/build.gradle
@@ -1,7 +1,7 @@
 apply plugin: 'com.android.library'
 
 android {
-    compileSdkVersion 25
+    compileSdkVersion 'android-P'
     buildToolsVersion "25.0.0"
 
     lintOptions {
@@ -28,6 +28,6 @@ repositories {
 
 dependencies {
     compile project(':dexmaker')
-    compile 'org.mockito:mockito-core:2.12.0'
+    compile 'org.mockito:mockito-core:2.15.0', { exclude group: "net.bytebuddy" }
 }
 
diff --git a/dexmaker-mockito-inline/src/main/java/com/android/dx/mockito/inline/ClassTransformer.java b/dexmaker-mockito-inline/src/main/java/com/android/dx/mockito/inline/ClassTransformer.java
index f719304..c702b2f 100644
--- a/dexmaker-mockito-inline/src/main/java/com/android/dx/mockito/inline/ClassTransformer.java
+++ b/dexmaker-mockito-inline/src/main/java/com/android/dx/mockito/inline/ClassTransformer.java
@@ -17,15 +17,15 @@
 package com.android.dx.mockito.inline;
 
 import org.mockito.exceptions.base.MockitoException;
-import org.mockito.internal.util.concurrent.WeakConcurrentMap;
-import org.mockito.internal.util.concurrent.WeakConcurrentSet;
 
 import java.lang.reflect.InvocationTargetException;
 import java.lang.reflect.Method;
 import java.lang.reflect.Modifier;
 import java.security.ProtectionDomain;
 import java.util.Arrays;
+import java.util.Collections;
 import java.util.HashSet;
+import java.util.Map;
 import java.util.Random;
 import java.util.Set;
 
@@ -65,7 +65,7 @@ class ClassTransformer {
     private final JvmtiAgent agent;
 
     /** Types that have already be transformed */
-    private final WeakConcurrentSet<Class<?>> mockedTypes;
+    private final Set<Class<?>> mockedTypes;
 
     /**
      * A unique identifier that is baked into the transformed classes. The entry hooks will then
@@ -96,9 +96,9 @@ class ClassTransformer {
      *              mocked or not.
      */
     ClassTransformer(JvmtiAgent agent, Class dispatcherClass,
-                     WeakConcurrentMap<Object, InvocationHandlerAdapter> mocks) {
+                     Map<Object, InvocationHandlerAdapter> mocks) {
         this.agent = agent;
-        mockedTypes = new WeakConcurrentSet<>(WeakConcurrentSet.Cleaner.INLINE);
+        mockedTypes = Collections.synchronizedSet(new HashSet<Class<?>>());
         identifier = Long.toString(random.nextLong());
         MockMethodAdvice advice = new MockMethodAdvice(mocks);
 
@@ -186,5 +186,16 @@ class ClassTransformer {
         }
     }
 
+    /**
+     * Check if the class should be transformed.
+     *
+     * @param classBeingRedefined The class that might need to transformed
+     *
+     * @return {@code true} iff the class needs to be transformed
+     */
+    boolean shouldTransform(Class<?> classBeingRedefined) {
+        return classBeingRedefined != null && mockedTypes.contains(classBeingRedefined);
+    }
+
     private native byte[] nativeRedefine(String identifier, byte[] original);
 }
diff --git a/dexmaker-mockito-inline/src/main/java/com/android/dx/mockito/inline/InlineDexmakerMockMaker.java b/dexmaker-mockito-inline/src/main/java/com/android/dx/mockito/inline/InlineDexmakerMockMaker.java
index d5be235..319ff0d 100644
--- a/dexmaker-mockito-inline/src/main/java/com/android/dx/mockito/inline/InlineDexmakerMockMaker.java
+++ b/dexmaker-mockito-inline/src/main/java/com/android/dx/mockito/inline/InlineDexmakerMockMaker.java
@@ -16,24 +16,34 @@
 
 package com.android.dx.mockito.inline;
 
+import android.os.AsyncTask;
+import android.os.Build;
+import android.util.ArraySet;
+
 import com.android.dx.stock.ProxyBuilder;
 import com.android.dx.stock.ProxyBuilder.MethodSetEntry;
 
+import org.mockito.Mockito;
 import org.mockito.exceptions.base.MockitoException;
-import org.mockito.internal.configuration.plugins.Plugins;
 import org.mockito.internal.creation.instance.Instantiator;
-import org.mockito.internal.util.Platform;
-import org.mockito.internal.util.concurrent.WeakConcurrentMap;
 import org.mockito.invocation.MockHandler;
 import org.mockito.mock.MockCreationSettings;
+import org.mockito.plugins.InstantiatorProvider;
 import org.mockito.plugins.MockMaker;
 
 import java.io.IOException;
 import java.io.InputStream;
+import java.lang.ref.Reference;
+import java.lang.ref.ReferenceQueue;
+import java.lang.ref.WeakReference;
 import java.lang.reflect.Method;
 import java.lang.reflect.Modifier;
 import java.lang.reflect.Proxy;
+import java.util.AbstractMap;
+import java.util.Collection;
+import java.util.HashMap;
 import java.util.HashSet;
+import java.util.Map;
 import java.util.Set;
 
 /**
@@ -111,7 +121,7 @@ public final class InlineDexmakerMockMaker implements MockMaker {
      * modified, some are not. This list helps the {@link MockMethodAdvice} help figure out if a
      * object's method calls should be intercepted.
      */
-    private final WeakConcurrentMap<Object, InvocationHandlerAdapter> mocks;
+    private final Map<Object, InvocationHandlerAdapter> mocks;
 
     /**
      * Class doing the actual byte code transformation.
@@ -126,10 +136,11 @@ public final class InlineDexmakerMockMaker implements MockMaker {
             throw new RuntimeException(
                     "Could not initialize inline mock maker.\n"
                     + "\n"
-                    + Platform.describe(), INITIALIZATION_ERROR);
+                    + "Release: Android " + Build.VERSION.RELEASE + " " + Build.VERSION.INCREMENTAL
+                    + "Device: " + Build.BRAND + " " + Build.MODEL, INITIALIZATION_ERROR);
         }
 
-        mocks = new WeakConcurrentMap.WithInlinedExpunction<>();
+        mocks = new MockMap();
         classTransformer = new ClassTransformer(AGENT, DISPATCHER_CLASS, mocks);
     }
 
@@ -212,7 +223,8 @@ public final class InlineDexmakerMockMaker implements MockMaker {
 
             Class<? extends T> proxyClass;
 
-            Instantiator instantiator = Plugins.getInstantiatorProvider().getInstantiator(settings);
+            Instantiator instantiator = Mockito.framework().getPlugins()
+                    .getDefaultPlugin(InstantiatorProvider.class).getInstantiator(settings);
 
             if (subclassingRequired) {
                 try {
@@ -299,4 +311,336 @@ public final class InlineDexmakerMockMaker implements MockMaker {
 
         return mocks.get(instance);
     }
+
+    /**
+     * A map mock -> adapter that holds weak references to the mocks and cleans them up when a
+     * stale reference is found.
+     */
+    private static class MockMap extends ReferenceQueue<Object>
+            implements Map<Object, InvocationHandlerAdapter> {
+        private static final int MIN_CLEAN_INTERVAL_MILLIS = 16000;
+        private static final int MAX_GET_WITHOUT_CLEAN = 16384;
+
+        private final Object lock = new Object();
+        private static StrongKey cachedKey;
+
+        private HashMap<WeakKey, InvocationHandlerAdapter> adapters = new HashMap<>();
+
+        /**
+         * The time we issues the last cleanup
+         */
+        long mLastCleanup = 0;
+
+        /**
+         * If {@link #cleanStaleReferences} is currently cleaning stale references out of
+         * {@link #adapters}
+         */
+        private boolean isCleaning = false;
+
+        /**
+         * The number of time {@link #get} was called without cleaning up stale references.
+         * {@link #get} is a method that is called often.
+         *
+         * We need to do periodic cleanups as we might never look at mocks at higher indexes and
+         * hence never realize that their references are stale.
+         */
+        private int getCount = 0;
+
+        /**
+         * Try to get a recycled cached key.
+         *
+         * @param obj the reference the key wraps
+         *
+         * @return The recycled cached key or a new one
+         */
+        private StrongKey createStrongKey(Object obj) {
+            synchronized (lock) {
+                if (cachedKey == null) {
+                    cachedKey = new StrongKey();
+                }
+
+                cachedKey.obj = obj;
+                StrongKey newKey = cachedKey;
+                cachedKey = null;
+
+                return newKey;
+            }
+        }
+
+        /**
+         * Recycle a key. The key should not be used afterwards
+         *
+         * @param key The key to recycle
+         */
+        private void recycleStrongKey(StrongKey key) {
+            synchronized (lock) {
+                cachedKey = key;
+            }
+        }
+
+        @Override
+        public int size() {
+            return adapters.size();
+        }
+
+        @Override
+        public boolean isEmpty() {
+            return adapters.isEmpty();
+        }
+
+        @Override
+        public boolean containsKey(Object mock) {
+            synchronized (lock) {
+                StrongKey key = createStrongKey(mock);
+                boolean containsKey = adapters.containsKey(key);
+                recycleStrongKey(key);
+
+                return containsKey;
+            }
+        }
+
+        @Override
+        public boolean containsValue(Object adapter) {
+            synchronized (lock) {
+                return adapters.containsValue(adapter);
+            }
+        }
+
+        @Override
+        public InvocationHandlerAdapter get(Object mock) {
+            synchronized (lock) {
+                if (getCount > MAX_GET_WITHOUT_CLEAN) {
+                    cleanStaleReferences();
+                    getCount = 0;
+                } else {
+                    getCount++;
+                }
+
+                StrongKey key = createStrongKey(mock);
+                InvocationHandlerAdapter adapter = adapters.get(key);
+                recycleStrongKey(key);
+
+                return adapter;
+            }
+        }
+
+        /**
+         * Remove entries that reference a stale mock from {@link #adapters}.
+         */
+        private void cleanStaleReferences() {
+            synchronized (lock) {
+                if (!isCleaning) {
+                    if (System.currentTimeMillis() - MIN_CLEAN_INTERVAL_MILLIS < mLastCleanup) {
+                        return;
+                    }
+
+                    isCleaning = true;
+
+                    AsyncTask.THREAD_POOL_EXECUTOR.execute(new Runnable() {
+                        @Override
+                        public void run() {
+                            synchronized (lock) {
+                                while (true) {
+                                    Reference<?> ref = MockMap.this.poll();
+                                    if (ref == null) {
+                                        break;
+                                    }
+
+                                    adapters.remove(ref);
+                                }
+
+                                mLastCleanup = System.currentTimeMillis();
+                                isCleaning = false;
+                            }
+                        }
+                    });
+                }
+            }
+        }
+
+        @Override
+        public InvocationHandlerAdapter put(Object mock, InvocationHandlerAdapter adapter) {
+            synchronized (lock) {
+                InvocationHandlerAdapter oldValue = remove(mock);
+                adapters.put(new WeakKey(mock), adapter);
+
+                return oldValue;
+            }
+        }
+
+        @Override
+        public InvocationHandlerAdapter remove(Object mock) {
+            synchronized (lock) {
+                StrongKey key = createStrongKey(mock);
+                InvocationHandlerAdapter adapter = adapters.remove(key);
+                recycleStrongKey(key);
+
+                return adapter;
+            }
+        }
+
+        @Override
+        public void putAll(Map<?, ? extends InvocationHandlerAdapter> map) {
+            synchronized (lock) {
+                for (Entry<?, ? extends InvocationHandlerAdapter> entry : map.entrySet()) {
+                    put(entry.getKey(), entry.getValue());
+                }
+            }
+        }
+
+        @Override
+        public void clear() {
+            synchronized (lock) {
+                adapters.clear();
+            }
+        }
+
+        @Override
+        public Set<Object> keySet() {
+            synchronized (lock) {
+                Set<Object> mocks = new ArraySet<>(adapters.size());
+
+                boolean hasStaleReferences = false;
+                for (WeakKey key : adapters.keySet()) {
+                    Object mock = key.get();
+
+                    if (mock == null) {
+                        hasStaleReferences = true;
+                    } else {
+                        mocks.add(mock);
+                    }
+                }
+
+                if (hasStaleReferences) {
+                    cleanStaleReferences();
+                }
+
+                return mocks;
+            }
+        }
+
+        @Override
+        public Collection<InvocationHandlerAdapter> values() {
+            synchronized (lock) {
+                return adapters.values();
+            }
+        }
+
+        @Override
+        public Set<Entry<Object, InvocationHandlerAdapter>> entrySet() {
+            synchronized (lock) {
+                Set<Entry<Object, InvocationHandlerAdapter>> entries = new ArraySet<>(
+                        adapters.size());
+
+                boolean hasStaleReferences = false;
+                for (Entry<WeakKey, InvocationHandlerAdapter> entry : adapters.entrySet()) {
+                    Object mock = entry.getKey().get();
+
+                    if (mock == null) {
+                        hasStaleReferences = true;
+                    } else {
+                        entries.add(new AbstractMap.SimpleEntry<>(mock, entry.getValue()));
+                    }
+                }
+
+                if (hasStaleReferences) {
+                    cleanStaleReferences();
+                }
+
+                return entries;
+            }
+        }
+
+        /**
+         * A weakly referencing wrapper to a mock.
+         *
+         * Only equals other weak or strong keys where the mock is the same.
+         */
+        private class WeakKey extends WeakReference<Object> {
+            private final int hashCode;
+
+            private WeakKey(/*@NonNull*/ Object obj) {
+                super(obj, MockMap.this);
+
+                // Cache the hashcode as the referenced object might disappear
+                hashCode = System.identityHashCode(obj);
+            }
+
+            @Override
+            public boolean equals(Object other) {
+                if (other == this) {
+                    return true;
+                }
+
+                if (other == null) {
+                    return false;
+                }
+
+                // Checking hashcode is cheap
+                if (other.hashCode() != hashCode) {
+                    return false;
+                }
+
+                Object obj = get();
+
+                if (obj == null) {
+                    cleanStaleReferences();
+                    return false;
+                }
+
+                if (other instanceof WeakKey) {
+                    Object otherObj = ((WeakKey) other).get();
+
+                    if (otherObj == null) {
+                        cleanStaleReferences();
+                        return false;
+                    }
+
+                    return obj == otherObj;
+                } else if (other instanceof StrongKey) {
+                    Object otherObj = ((StrongKey) other).obj;
+                    return obj == otherObj;
+                } else {
+                    return false;
+                }
+            }
+
+            @Override
+            public int hashCode() {
+                return hashCode;
+            }
+        }
+
+        /**
+         * A strongly referencing wrapper to a mock.
+         *
+         * Only equals other weak or strong keys where the mock is the same.
+         */
+        private class StrongKey {
+            /*@NonNull*/ private Object obj;
+
+            @Override
+            public boolean equals(Object other) {
+                if (other instanceof WeakKey) {
+                    Object otherObj = ((WeakKey) other).get();
+
+                    if (otherObj == null) {
+                        cleanStaleReferences();
+                        return false;
+                    }
+
+                    return obj == otherObj;
+                } else if (other instanceof StrongKey) {
+                    return this.obj == ((StrongKey)other).obj;
+                } else {
+                    return false;
+                }
+            }
+
+            @Override
+            public int hashCode() {
+                return System.identityHashCode(obj);
+            }
+        }
+    }
 }
diff --git a/dexmaker-mockito-inline/src/main/java/com/android/dx/mockito/inline/InterceptedInvocation.java b/dexmaker-mockito-inline/src/main/java/com/android/dx/mockito/inline/InterceptedInvocation.java
deleted file mode 100644
index a6d11b5..0000000
--- a/dexmaker-mockito-inline/src/main/java/com/android/dx/mockito/inline/InterceptedInvocation.java
+++ /dev/null
@@ -1,189 +0,0 @@
-/*
- * Copyright (c) 2016 Mockito contributors
- * This program is made available under the terms of the MIT License.
- */
-
-package com.android.dx.mockito.inline;
-
-import org.mockito.internal.debugging.LocationImpl;
-import org.mockito.internal.exceptions.VerificationAwareInvocation;
-import org.mockito.internal.invocation.ArgumentsProcessor;
-import org.mockito.internal.invocation.MockitoMethod;
-import org.mockito.internal.reporting.PrintSettings;
-import org.mockito.invocation.Invocation;
-import org.mockito.invocation.Location;
-import org.mockito.invocation.StubInfo;
-
-import java.io.Serializable;
-import java.lang.reflect.Method;
-import java.util.Arrays;
-
-import static org.mockito.internal.exceptions.Reporter.cannotCallAbstractRealMethod;
-
-/**
- * {@link Invocation} used when intercepting methods from an method entry hook.
- */
-class InterceptedInvocation implements Invocation, VerificationAwareInvocation {
-    /** The mocked instance */
-    private final Object mock;
-
-    /** The method invoked */
-    private final MockitoMethod method;
-
-    /** expanded arguments to the method */
-    private final Object[] arguments;
-
-    /** raw arguments to the method */
-    private final Object[] rawArguments;
-
-    /** The super method */
-    private final SuperMethod superMethod;
-
-    /** sequence number of the invocation (different for each invocation) */
-    private final int sequenceNumber;
-
-    /** the location of the invocation (i.e. the stack trace) */
-    private final Location location;
-
-    /** Was this invocation {@link #markVerified() marked as verified} */
-    private boolean verified;
-
-    /** Should this be {@link #ignoreForVerification()} ignored for verification?} */
-    private boolean isIgnoredForVerification;
-
-    /** The stubinfo is this was {@link #markStubbed(StubInfo) markes as stubbed}*/
-    private StubInfo stubInfo;
-
-    /**
-     * Create a new invocation.
-     *
-     * @param mock mocked instance
-     * @param method method invoked
-     * @param arguments arguments to the method
-     * @param superMethod super method
-     * @param sequenceNumber sequence number of the invocation
-     */
-    InterceptedInvocation(Object mock, MockitoMethod method, Object[] arguments,
-                          SuperMethod superMethod, int sequenceNumber) {
-        this.mock = mock;
-        this.method = method;
-        this.arguments = ArgumentsProcessor.expandArgs(method, arguments);
-        this.rawArguments = arguments;
-        this.superMethod = superMethod;
-        this.sequenceNumber = sequenceNumber;
-        location = new LocationImpl();
-    }
-
-    @Override
-    public boolean isVerified() {
-        return verified || isIgnoredForVerification;
-    }
-
-    @Override
-    public int getSequenceNumber() {
-        return sequenceNumber;
-    }
-
-    @Override
-    public Location getLocation() {
-        return location;
-    }
-
-    @Override
-    public Object[] getRawArguments() {
-        return rawArguments;
-    }
-
-    @Override
-    public Class<?> getRawReturnType() {
-        return method.getReturnType();
-    }
-
-    @Override
-    public void markVerified() {
-        verified = true;
-    }
-
-    @Override
-    public StubInfo stubInfo() {
-        return stubInfo;
-    }
-
-    @Override
-    public void markStubbed(StubInfo stubInfo) {
-        this.stubInfo = stubInfo;
-    }
-
-    @Override
-    public boolean isIgnoredForVerification() {
-        return isIgnoredForVerification;
-    }
-
-    @Override
-    public void ignoreForVerification() {
-        isIgnoredForVerification = true;
-    }
-
-    @Override
-    public Object getMock() {
-        return mock;
-    }
-
-    @Override
-    public Method getMethod() {
-        return method.getJavaMethod();
-    }
-
-    @Override
-    public Object[] getArguments() {
-        return arguments;
-    }
-
-    @Override
-    @SuppressWarnings("unchecked")
-    public <T> T getArgument(int index) {
-        return (T) arguments[index];
-    }
-
-    @Override
-    public Object callRealMethod() throws Throwable {
-        if (!superMethod.isInvokable()) {
-            throw cannotCallAbstractRealMethod();
-        }
-        return superMethod.invoke();
-    }
-
-    @Override
-    public int hashCode() {
-        // TODO SF we need to provide hash code implementation so that there are no unexpected,
-        //         slight perf issues
-        return 1;
-    }
-
-    @Override
-    public boolean equals(Object o) {
-        if (o == null || !o.getClass().equals(this.getClass())) {
-            return false;
-        }
-        InterceptedInvocation other = (InterceptedInvocation) o;
-        return this.mock.equals(other.mock)
-                && this.method.equals(other.method)
-                && this.equalArguments(other.arguments);
-    }
-
-    private boolean equalArguments(Object[] arguments) {
-        return Arrays.equals(arguments, this.arguments);
-    }
-
-    @Override
-    public String toString() {
-        return new PrintSettings().print(ArgumentsProcessor.argumentsToMatchers(getArguments()),
-                this);
-    }
-
-    interface SuperMethod extends Serializable {
-        boolean isInvokable();
-
-        Object invoke() throws Throwable;
-    }
-}
diff --git a/dexmaker-mockito-inline/src/main/java/com/android/dx/mockito/inline/InvocationHandlerAdapter.java b/dexmaker-mockito-inline/src/main/java/com/android/dx/mockito/inline/InvocationHandlerAdapter.java
index 5d01a1e..afeedb2 100644
--- a/dexmaker-mockito-inline/src/main/java/com/android/dx/mockito/inline/InvocationHandlerAdapter.java
+++ b/dexmaker-mockito-inline/src/main/java/com/android/dx/mockito/inline/InvocationHandlerAdapter.java
@@ -16,22 +16,17 @@
 
 package com.android.dx.mockito.inline;
 
-import org.mockito.internal.creation.DelegatingMethod;
-import org.mockito.internal.debugging.LocationImpl;
-import org.mockito.internal.exceptions.VerificationAwareInvocation;
-import org.mockito.internal.invocation.ArgumentsProcessor;
-import org.mockito.internal.progress.SequenceNumber;
-import org.mockito.invocation.Invocation;
-import org.mockito.invocation.Location;
+import com.android.dx.stock.ProxyBuilder;
+
+import org.mockito.Mockito;
+import org.mockito.invocation.InvocationFactory.RealMethodBehavior;
 import org.mockito.invocation.MockHandler;
-import org.mockito.invocation.StubInfo;
 import org.mockito.mock.MockCreationSettings;
 
 import java.lang.reflect.InvocationHandler;
 import java.lang.reflect.Method;
-import java.lang.reflect.Modifier;
 
-import static org.mockito.internal.exceptions.Reporter.cannotCallAbstractRealMethod;
+import static org.mockito.Mockito.withSettings;
 
 /**
  * Handles proxy and entry hook method invocations added by
@@ -63,46 +58,65 @@ final class InvocationHandlerAdapter implements InvocationHandler {
      *
      * @param mock mocked object
      * @param method method that was called
-     * @param args arguments to the method
+     * @param rawArgs arguments to the method
      * @param superMethod The super method
      *
      * @return mocked result
      * @throws Throwable An exception if thrown
      */
-    Object interceptEntryHook(Object mock, Method method, Object[] args,
-                              InterceptedInvocation.SuperMethod superMethod) throws Throwable {
-        return handler.handle(new InterceptedInvocation(mock, new DelegatingMethod(method), args,
-                superMethod, SequenceNumber.next()));
+    Object interceptEntryHook(final Object mock, final Method method, final Object[] rawArgs,
+                              final SuperMethod superMethod) throws Throwable {
+        // args can be null if the method invoked has no arguments, but Mockito expects a non-null
+        Object[] args = rawArgs;
+        if (rawArgs == null) {
+            args = new Object[0];
+        }
+
+        return handler.handle(Mockito.framework().getInvocationFactory().createInvocation(mock,
+                withSettings().build(mock.getClass()), method, new RealMethodBehavior() {
+                    @Override
+                    public Object call() throws Throwable {
+                        return superMethod.invoke();
+                    }
+                }, args));
     }
 
     /**
      * Intercept a method call. Called <u>before</u> a method is called by the proxied method.
      *
-     * <p>This does the same as {@link #interceptEntryHook(Object, Method, Object[],
-     * InterceptedInvocation.SuperMethod)} but this handles proxied methods. We only proxy abstract
-     * methods.
+     * <p>This does the same as {@link #interceptEntryHook(Object, Method, Object[], SuperMethod)}
+     * but this handles proxied methods. We only proxy abstract methods.
      *
      * @param proxy proxies object
      * @param method method that was called
-     * @param argsIn arguments to the method
+     * @param rawArgs arguments to the method
      *
      * @return mocked result
      * @throws Throwable An exception if thrown
      */
     @Override
-    public Object invoke(final Object proxy, final Method method, Object[] argsIn) throws
+    public Object invoke(final Object proxy, final Method method, final Object[] rawArgs) throws
             Throwable {
         // args can be null if the method invoked has no arguments, but Mockito expects a non-null
-        // array
-        final Object[] args = argsIn != null ? argsIn : new Object[0];
+        Object[] args = rawArgs;
+        if (rawArgs == null) {
+            args = new Object[0];
+        }
+
         if (isEqualsMethod(method)) {
             return proxy == args[0];
         } else if (isHashCodeMethod(method)) {
             return System.identityHashCode(proxy);
         }
 
-        return handler.handle(new ProxyInvocation(proxy, method, args, new DelegatingMethod
-                (method), SequenceNumber.next(), new LocationImpl()));
+        return handler.handle(Mockito.framework().getInvocationFactory().createInvocation(proxy,
+                withSettings().build(proxy.getClass().getSuperclass()), method,
+                new RealMethodBehavior() {
+                    @Override
+                    public Object call() throws Throwable {
+                        return ProxyBuilder.callSuper(proxy, method, rawArgs);
+                    }
+                }, args));
     }
 
     /**
@@ -122,106 +136,9 @@ final class InvocationHandlerAdapter implements InvocationHandler {
     }
 
     /**
-     * Invocation on a proxy
+     * Interface used to describe a supermethod that can be called.
      */
-    private class ProxyInvocation implements Invocation, VerificationAwareInvocation {
-        private final Object proxy;
-        private final Method method;
-        private final Object[] rawArgs;
-        private final int sequenceNumber;
-        private final Location location;
-        private final Object[] args;
-
-        private StubInfo stubInfo;
-        private boolean isIgnoredForVerification;
-        private boolean verified;
-
-        private ProxyInvocation(Object proxy, Method method, Object[] rawArgs, DelegatingMethod
-                                mockitoMethod, int sequenceNumber, Location location) {
-            this.rawArgs = rawArgs;
-            this.proxy = proxy;
-            this.method = method;
-            this.sequenceNumber = sequenceNumber;
-            this.location = location;
-            args = ArgumentsProcessor.expandArgs(mockitoMethod, rawArgs);
-        }
-
-        @Override
-        public Object getMock() {
-            return proxy;
-        }
-
-        @Override
-        public Method getMethod() {
-            return method;
-        }
-
-        @Override
-        public Object[] getArguments() {
-            return args;
-        }
-
-        @Override
-        public <T> T getArgument(int index) {
-            return (T)args[index];
-        }
-
-        @Override
-        public Object callRealMethod() throws Throwable {
-            if (Modifier.isAbstract(method.getModifiers())) {
-                throw cannotCallAbstractRealMethod();
-            }
-            return method.invoke(proxy, rawArgs);
-        }
-
-        @Override
-        public boolean isVerified() {
-            return verified || isIgnoredForVerification;
-        }
-
-        @Override
-        public int getSequenceNumber() {
-            return sequenceNumber;
-        }
-
-        @Override
-        public Location getLocation() {
-            return location;
-        }
-
-        @Override
-        public Object[] getRawArguments() {
-            return rawArgs;
-        }
-
-        @Override
-        public Class<?> getRawReturnType() {
-            return method.getReturnType();
-        }
-
-        @Override
-        public void markVerified() {
-            verified = true;
-        }
-
-        @Override
-        public StubInfo stubInfo() {
-            return stubInfo;
-        }
-
-        @Override
-        public void markStubbed(StubInfo stubInfo) {
-            this.stubInfo = stubInfo;
-        }
-
-        @Override
-        public boolean isIgnoredForVerification() {
-            return isIgnoredForVerification;
-        }
-
-        @Override
-        public void ignoreForVerification() {
-            isIgnoredForVerification = true;
-        }
+    interface SuperMethod {
+        Object invoke() throws Throwable;
     }
 }
diff --git a/dexmaker-mockito-inline/src/main/java/com/android/dx/mockito/inline/JvmtiAgent.java b/dexmaker-mockito-inline/src/main/java/com/android/dx/mockito/inline/JvmtiAgent.java
index aca7a0c..5cd79d1 100644
--- a/dexmaker-mockito-inline/src/main/java/com/android/dx/mockito/inline/JvmtiAgent.java
+++ b/dexmaker-mockito-inline/src/main/java/com/android/dx/mockito/inline/JvmtiAgent.java
@@ -17,14 +17,13 @@
 package com.android.dx.mockito.inline;
 
 import android.os.Build;
+import android.os.Debug;
 
 import java.io.File;
-import java.io.FileInputStream;
 import java.io.FileOutputStream;
 import java.io.IOException;
 import java.io.InputStream;
 import java.io.OutputStream;
-import java.lang.reflect.InvocationTargetException;
 import java.security.ProtectionDomain;
 import java.util.ArrayList;
 
@@ -57,37 +56,13 @@ class JvmtiAgent {
             throw new IOException("Requires Android P. Build is " + Build.VERSION.CODENAME);
         }
 
-        Throwable loadJvmtiException = null;
-
         ClassLoader cl = JvmtiAgent.class.getClassLoader();
         if (!(cl instanceof BaseDexClassLoader)) {
             throw new IOException("Could not load jvmti plugin as JvmtiAgent class was not loaded "
                     + "by a BaseDexClassLoader");
         }
 
-        try {
-            /*
-             * TODO (moltmann@google.com): Replace with regular method call once the API becomes
-             *                             public
-             */
-            Class.forName("android.os.Debug").getMethod("attachJvmtiAgent", String.class,
-                    String.class, ClassLoader.class).invoke(null, AGENT_LIB_NAME,
-                    null, cl);
-        } catch (InvocationTargetException e) {
-            loadJvmtiException = e.getCause();
-        } catch (IllegalAccessException | ClassNotFoundException | NoSuchMethodException e) {
-            loadJvmtiException = e;
-        }
-
-        if (loadJvmtiException != null) {
-            if (loadJvmtiException instanceof IOException) {
-                throw new IOException(cl.toString(), loadJvmtiException);
-            } else {
-                throw new IOException("Could not load jvmti plugin",
-                        loadJvmtiException);
-            }
-        }
-
+        Debug.attachJvmtiAgent(AGENT_LIB_NAME, null, cl);
         nativeRegisterTransformerHook();
     }
 
@@ -145,6 +120,18 @@ class JvmtiAgent {
         }
     }
 
+    // called by JNI
+    @SuppressWarnings("unused")
+    public boolean shouldTransform(Class<?> classBeingRedefined) {
+        for (ClassTransformer transformer : transformers) {
+            if (transformer.shouldTransform(classBeingRedefined)) {
+                return true;
+            }
+        }
+
+        return false;
+    }
+
     /**
      * Register a transformer. These are called for each class when a transformation was triggered
      * via {@link #requestTransformClasses(Class[])}.
diff --git a/dexmaker-mockito-inline/src/main/java/com/android/dx/mockito/inline/MockMethodAdvice.java b/dexmaker-mockito-inline/src/main/java/com/android/dx/mockito/inline/MockMethodAdvice.java
index 4cf2ac8..dfe242f 100644
--- a/dexmaker-mockito-inline/src/main/java/com/android/dx/mockito/inline/MockMethodAdvice.java
+++ b/dexmaker-mockito-inline/src/main/java/com/android/dx/mockito/inline/MockMethodAdvice.java
@@ -5,13 +5,11 @@
 
 package com.android.dx.mockito.inline;
 
-import org.mockito.internal.exceptions.stacktrace.ConditionalStackTraceFilter;
-import org.mockito.internal.util.concurrent.WeakConcurrentMap;
-
 import java.lang.reflect.InvocationTargetException;
 import java.lang.reflect.Method;
 import java.lang.reflect.Modifier;
 import java.util.ArrayList;
+import java.util.Map;
 import java.util.concurrent.Callable;
 import java.util.regex.Matcher;
 import java.util.regex.Pattern;
@@ -21,14 +19,14 @@ import java.util.regex.Pattern;
  * be ignored.
  */
 class MockMethodAdvice {
-    private final WeakConcurrentMap<Object, InvocationHandlerAdapter> interceptors;
+    private final Map<Object, InvocationHandlerAdapter> interceptors;
 
     /** Pattern to decompose a instrumentedMethodWithTypeAndSignature */
     private final Pattern methodPattern = Pattern.compile("(.*)#(.*)\\((.*)\\)");
 
     private final SelfCallInfo selfCallInfo = new SelfCallInfo();
 
-    MockMethodAdvice(WeakConcurrentMap<Object, InvocationHandlerAdapter> interceptors) {
+    MockMethodAdvice(Map<Object, InvocationHandlerAdapter> interceptors) {
         this.interceptors = interceptors;
     }
 
@@ -48,12 +46,7 @@ class MockMethodAdvice {
         try {
             return origin.invoke(instance, arguments);
         } catch (InvocationTargetException exception) {
-            Throwable cause = exception.getCause();
-
-            new ConditionalStackTraceFilter().filter(hideRecursiveCall(cause,
-                    new Throwable().getStackTrace().length, origin.getDeclaringClass()));
-
-            throw cause;
+            throw exception.getCause();
         }
     }
 
@@ -259,7 +252,7 @@ class MockMethodAdvice {
     /**
      * Used to call the read (non mocked) method.
      */
-    private static class SuperMethodCall implements InterceptedInvocation.SuperMethod {
+    private static class SuperMethodCall implements InvocationHandlerAdapter.SuperMethod {
         private final SelfCallInfo selfCallInfo;
         private final Method origin;
         private final Object instance;
@@ -273,11 +266,6 @@ class MockMethodAdvice {
             this.arguments = arguments;
         }
 
-        @Override
-        public boolean isInvokable() {
-            return true;
-        }
-
         /**
          * Call the read (non mocked) method.
          *
diff --git a/dexmaker-mockito-inline/src/main/jni/dexmakerjvmtiagent/agent.cc b/dexmaker-mockito-inline/src/main/jni/dexmakerjvmtiagent/agent.cc
index 358e375..cfa10a5 100644
--- a/dexmaker-mockito-inline/src/main/jni/dexmakerjvmtiagent/agent.cc
+++ b/dexmaker-mockito-inline/src/main/jni/dexmakerjvmtiagent/agent.cc
@@ -25,13 +25,13 @@
 
 #include "jvmti.h"
 
-#include <dex_ir.h>
-#include <code_ir.h>
-#include <dex_ir_builder.h>
-#include <dex_utf8.h>
-#include <writer.h>
-#include <reader.h>
-#include <instrumentation.h>
+#include <slicer/dex_ir.h>
+#include <slicer/code_ir.h>
+#include <slicer/dex_ir_builder.h>
+#include <slicer/dex_utf8.h>
+#include <slicer/writer.h>
+#include <slicer/reader.h>
+#include <slicer/instrumentation.h>
 
 using namespace dex;
 using namespace lir;
@@ -70,6 +70,19 @@ Transform(jvmtiEnv* jvmti_env,
           jint* newClassDataLen,
           unsigned char** newClassData) {
     if (sTransformer != NULL) {
+        // Even reading the classData array is expensive as the data is only generated when the
+        // memory is touched. Hence call JvmtiAgent#shouldTransform to check if we need to transform
+        // the class.
+        jclass cls = env->GetObjectClass(sTransformer);
+        jmethodID shouldTransformMethod = env->GetMethodID(cls, "shouldTransform",
+                                                           "(Ljava/lang/Class;)Z");
+
+        jboolean shouldTransform = env->CallBooleanMethod(sTransformer, shouldTransformMethod,
+                                                          classBeingRedefined);
+        if (!shouldTransform) {
+            return;
+        }
+
         // Isolate byte code of class class. This is needed as Android usually gives us more
         // than the class we need.
         Reader reader(classData, classDataLen);
@@ -97,16 +110,15 @@ Transform(jvmtiEnv* jvmti_env,
         jstring nameStr = env->NewStringUTF(name);
 
         // Call JvmtiAgent#runTransformers
-        jclass cls = env->GetObjectClass(sTransformer);
-        jmethodID runTransformers = env->GetMethodID(cls, "runTransformers",
-                                                     "(Ljava/lang/ClassLoader;"
-                                                     "Ljava/lang/String;"
-                                                     "Ljava/lang/Class;"
-                                                     "Ljava/security/ProtectionDomain;"
-                                                     "[B)[B");
+        jmethodID runTransformersMethod = env->GetMethodID(cls, "runTransformers",
+                                                           "(Ljava/lang/ClassLoader;"
+                                                           "Ljava/lang/String;"
+                                                           "Ljava/lang/Class;"
+                                                           "Ljava/security/ProtectionDomain;"
+                                                           "[B)[B");
 
         jbyteArray transformedArr = (jbyteArray) env->CallObjectMethod(sTransformer,
-                                                                       runTransformers,
+                                                                       runTransformersMethod,
                                                                        loader, nameStr,
                                                                        classBeingRedefined,
                                                                        protectionDomain,
@@ -779,19 +791,6 @@ Java_com_android_dx_mockito_inline_ClassTransformer_nativeRedefine(JNIEnv* env,
     return transformedArr;
 }
 
-// Register the ClassFileLoadHook hook. This causes Transform to be called for every class load and
-// for every trigger of RetransformClasses
-static jvmtiError registerClassFileLoadHook() {
-    return localJvmtiEnv->SetEventNotificationMode(JVMTI_ENABLE, JVMTI_EVENT_CLASS_FILE_LOAD_HOOK,
-                                                   NULL);
-}
-
-// Unregister the ClassFileLoadHook hook.
-static jvmtiError unregisterClassFileLoadHook() {
-    return localJvmtiEnv->SetEventNotificationMode(JVMTI_DISABLE, JVMTI_EVENT_CLASS_FILE_LOAD_HOOK,
-                                                   NULL);
-}
-
 // Initializes the agent
 extern "C" jint Agent_OnAttach(JavaVM* vm,
                                char* options,
@@ -819,6 +818,12 @@ extern "C" jint Agent_OnAttach(JavaVM* vm,
         return error;
     }
 
+    error = localJvmtiEnv->SetEventNotificationMode(JVMTI_ENABLE, JVMTI_EVENT_CLASS_FILE_LOAD_HOOK,
+                                                    NULL);
+    if (error != JVMTI_ERROR_NONE) {
+        return error;
+    }
+
     return JVMTI_ERROR_NONE;
 }
 
@@ -861,16 +866,8 @@ Java_com_android_dx_mockito_inline_JvmtiAgent_nativeRetransformClasses(JNIEnv* e
         transformedClasses[i] = (jclass) env->NewGlobalRef(env->GetObjectArrayElement(classes, i));
     }
 
-    jvmtiError error = registerClassFileLoadHook();
-    if (error == JVMTI_ERROR_NONE) {
-        error = localJvmtiEnv->RetransformClasses(numTransformedClasses,
-                                                  transformedClasses);
-
-        jvmtiError unregisterError = unregisterClassFileLoadHook();
-        if (error == JVMTI_ERROR_NONE && unregisterError != JVMTI_ERROR_NONE) {
-            error = unregisterError;
-        }
-    }
+    jvmtiError error = localJvmtiEnv->RetransformClasses(numTransformedClasses,
+                                                         transformedClasses);
 
     for (int i = 0; i < numTransformedClasses; i++) {
         env->DeleteGlobalRef(transformedClasses[i]);
diff --git a/dexmaker-mockito-tests/build.gradle b/dexmaker-mockito-tests/build.gradle
index f0befe0..a08e254 100644
--- a/dexmaker-mockito-tests/build.gradle
+++ b/dexmaker-mockito-tests/build.gradle
@@ -20,6 +20,7 @@ android {
 
 repositories {
     jcenter()
+    google()
 }
 
 dependencies {
diff --git a/dexmaker-mockito-tests/src/androidTest/java/com/android/dx/mockito/tests/GeneralMocking.java b/dexmaker-mockito-tests/src/androidTest/java/com/android/dx/mockito/tests/GeneralMocking.java
index 6f824f2..36d5612 100644
--- a/dexmaker-mockito-tests/src/androidTest/java/com/android/dx/mockito/tests/GeneralMocking.java
+++ b/dexmaker-mockito-tests/src/androidTest/java/com/android/dx/mockito/tests/GeneralMocking.java
@@ -18,13 +18,24 @@ package com.android.dx.mockito.tests;
 
 import android.support.test.runner.AndroidJUnit4;
 
+import org.junit.Ignore;
 import org.junit.Test;
 import org.junit.runner.RunWith;
+import org.mockito.exceptions.base.MockitoException;
 import org.mockito.exceptions.verification.NoInteractionsWanted;
 
+import java.util.ArrayList;
+import java.util.Arrays;
+
+import static org.junit.Assert.assertArrayEquals;
 import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertNotNull;
 import static org.junit.Assert.assertNull;
 import static org.junit.Assert.assertTrue;
+import static org.junit.Assert.fail;
+import static org.mockito.ArgumentMatchers.anyString;
+import static org.mockito.ArgumentMatchers.eq;
+import static org.mockito.ArgumentMatchers.isNull;
 import static org.mockito.Mockito.mock;
 import static org.mockito.Mockito.spy;
 import static org.mockito.Mockito.verifyNoMoreInteractions;
@@ -36,10 +47,32 @@ public class GeneralMocking {
         public String returnA() {
             return "A";
         }
+
+        public String throwThrowable() throws Throwable {
+            throw new Throwable();
+        }
+
+        public String throwOutOfMemoryError() throws OutOfMemoryError {
+            throw new OutOfMemoryError();
+        }
+
+        public void throwNullPointerException() {
+            throw new NullPointerException();
+        }
+
+        public String concat(String a, String b) {
+            return a + b;
+        }
+    }
+
+    public static class TestSubClass extends TestClass {
+
     }
 
     public interface TestInterface {
         String returnA();
+
+        String concat(String a, String b);
     }
 
     @Test
@@ -103,7 +136,147 @@ public class GeneralMocking {
                 assertTrue(e.getMessage(),
                         e.getMessage().contains(here.getStackTrace()[0].getMethodName()));
             }
+        }
+    }
+
+    @Test
+    public void spyThrowingMethod() throws Exception {
+        TestClass t = spy(TestClass.class);
+
+        try {
+            t.throwThrowable();
+        } catch (Throwable e) {
+            assertEquals("throwThrowable", e.getStackTrace()[0].getMethodName());
+            return;
+        }
+
+        fail();
+    }
 
+    @Test()
+    public void spyErrorMethod() throws Exception {
+        TestClass t = spy(TestClass.class);
+
+        try {
+            t.throwOutOfMemoryError();
+            fail();
+        } catch (OutOfMemoryError e) {
+            assertEquals("throwOutOfMemoryError", e.getStackTrace()[0].getMethodName());
+        }
+    }
+
+    @Test()
+    public void spyExceptingMethod() throws Exception {
+        TestClass t = spy(TestClass.class);
+
+        try {
+            t.throwNullPointerException();
+            fail();
+        } catch (NullPointerException e) {
+            assertEquals("throwNullPointerException", e.getStackTrace()[0].getMethodName());
         }
     }
+
+
+    @Test
+    public void callAbstractRealMethod() throws Exception {
+        TestInterface t = mock(TestInterface.class);
+
+        try {
+            when(t.returnA()).thenCallRealMethod();
+            fail();
+        } catch (MockitoException e) {
+            assertEquals("callAbstractRealMethod", e.getStackTrace()[0].getMethodName());
+        }
+    }
+
+    @Test
+    public void callInterfaceWithoutMatcher() throws Exception {
+        TestInterface t = mock(TestInterface.class);
+
+        when(t.concat("a", "b")).thenReturn("match");
+
+        assertEquals("match", t.concat("a", "b"));
+        assertNull(t.concat("b", "a"));
+    }
+
+    @Test
+    public void callInterfaceWithMatcher() throws Exception {
+        TestInterface t = mock(TestInterface.class);
+
+        when(t.concat(eq("a"), anyString())).thenReturn("match");
+
+        assertEquals("match", t.concat("a", "b"));
+        assertNull(t.concat("b", "a"));
+    }
+
+    @Test
+    public void callInterfaceWithNullMatcher() throws Exception {
+        TestInterface t = mock(TestInterface.class);
+
+        when(t.concat(eq("a"), (String) isNull())).thenReturn("match");
+
+        assertEquals("match", t.concat("a", null));
+        assertNull(t.concat("a", "b"));
+    }
+
+    @Test
+    public void callClassWithoutMatcher() throws Exception {
+        TestClass t = spy(TestClass.class);
+
+        when(t.concat("a", "b")).thenReturn("match");
+
+        assertEquals("match", t.concat("a", "b"));
+        assertEquals("ba", t.concat("b", "a"));
+    }
+
+    @Test
+    public void callClassWithMatcher() throws Exception {
+        TestClass t = spy(TestClass.class);
+
+        when(t.concat(eq("a"), anyString())).thenReturn("match");
+
+        assertEquals("match", t.concat("a", "b"));
+        assertEquals("ba", t.concat("b", "a"));
+    }
+
+    @Test
+    public void callClassWithNullMatcher() throws Exception {
+        TestClass t = spy(TestClass.class);
+
+        when(t.concat(eq("a"), (String) isNull())).thenReturn("match");
+
+        assertEquals("match", t.concat("a", null));
+        assertEquals("ab", t.concat("a", "b"));
+    }
+
+    @Test
+    public void callSubClassWithoutMatcher() throws Exception {
+        TestSubClass t = spy(TestSubClass.class);
+
+        when(t.concat("a", "b")).thenReturn("match");
+
+        assertEquals("match", t.concat("a", "b"));
+        assertEquals("ba", t.concat("b", "a"));
+    }
+
+    @Test
+    public void callSubClassWithMatcher() throws Exception {
+        TestSubClass t = spy(TestSubClass.class);
+
+        when(t.concat(eq("a"), anyString())).thenReturn("match");
+
+        assertEquals("match", t.concat("a", "b"));
+        assertEquals("ba", t.concat("b", "a"));
+    }
+
+    @Test
+    public void callSubClassWithNullMatcher() throws Exception {
+        TestSubClass t = spy(TestSubClass.class);
+
+        when(t.concat(eq("a"), (String) isNull())).thenReturn("match");
+
+        assertEquals("match", t.concat("a", null));
+        assertEquals("ab", t.concat("a", "b"));
+    }
 }
diff --git a/dexmaker-mockito/build.gradle b/dexmaker-mockito/build.gradle
index e1f5133..96479c5 100644
--- a/dexmaker-mockito/build.gradle
+++ b/dexmaker-mockito/build.gradle
@@ -14,5 +14,5 @@ repositories {
 dependencies {
     compile project(":dexmaker")
 
-    compile 'org.mockito:mockito-core:2.12.0'
+    compile 'org.mockito:mockito-core:2.15.0', { exclude group: "net.bytebuddy" }
 }
diff --git a/dexmaker-mockito/src/main/java/com/android/dx/mockito/InterceptedInvocation.java b/dexmaker-mockito/src/main/java/com/android/dx/mockito/InterceptedInvocation.java
deleted file mode 100644
index b0b42e9..0000000
--- a/dexmaker-mockito/src/main/java/com/android/dx/mockito/InterceptedInvocation.java
+++ /dev/null
@@ -1,189 +0,0 @@
-/*
- * Copyright (c) 2016 Mockito contributors
- * This program is made available under the terms of the MIT License.
- */
-
-package com.android.dx.mockito;
-
-import org.mockito.internal.debugging.LocationImpl;
-import org.mockito.internal.exceptions.VerificationAwareInvocation;
-import org.mockito.internal.invocation.ArgumentsProcessor;
-import org.mockito.internal.invocation.MockitoMethod;
-import org.mockito.internal.reporting.PrintSettings;
-import org.mockito.invocation.Invocation;
-import org.mockito.invocation.Location;
-import org.mockito.invocation.StubInfo;
-
-import java.io.Serializable;
-import java.lang.reflect.Method;
-import java.util.Arrays;
-
-import static org.mockito.internal.exceptions.Reporter.cannotCallAbstractRealMethod;
-
-/**
- * {@link Invocation} used when intercepting methods from an method entry hook.
- */
-class InterceptedInvocation implements Invocation, VerificationAwareInvocation {
-    /** The mocked instance */
-    private final Object mock;
-
-    /** The method invoked */
-    private final MockitoMethod method;
-
-    /** expanded arguments to the method */
-    private final Object[] arguments;
-
-    /** raw arguments to the method */
-    private final Object[] rawArguments;
-
-    /** The super method */
-    private final SuperMethod superMethod;
-
-    /** sequence number of the invocation (different for each invocation) */
-    private final int sequenceNumber;
-
-    /** the location of the invocation (i.e. the stack trace) */
-    private final Location location;
-
-    /** Was this invocation {@link #markVerified() marked as verified} */
-    private boolean verified;
-
-    /** Should this be {@link #ignoreForVerification()} ignored for verification?} */
-    private boolean isIgnoredForVerification;
-
-    /** The stubinfo is this was {@link #markStubbed(StubInfo) markes as stubbed}*/
-    private StubInfo stubInfo;
-
-    /**
-     * Create a new invocation.
-     *
-     * @param mock mocked instance
-     * @param method method invoked
-     * @param arguments arguments to the method
-     * @param superMethod super method
-     * @param sequenceNumber sequence number of the invocation
-     */
-    InterceptedInvocation(Object mock, MockitoMethod method, Object[] arguments,
-                          SuperMethod superMethod, int sequenceNumber) {
-        this.mock = mock;
-        this.method = method;
-        this.arguments = ArgumentsProcessor.expandArgs(method, arguments);
-        this.rawArguments = arguments;
-        this.superMethod = superMethod;
-        this.sequenceNumber = sequenceNumber;
-        location = new LocationImpl();
-    }
-
-    @Override
-    public boolean isVerified() {
-        return verified || isIgnoredForVerification;
-    }
-
-    @Override
-    public int getSequenceNumber() {
-        return sequenceNumber;
-    }
-
-    @Override
-    public Location getLocation() {
-        return location;
-    }
-
-    @Override
-    public Object[] getRawArguments() {
-        return rawArguments;
-    }
-
-    @Override
-    public Class<?> getRawReturnType() {
-        return method.getReturnType();
-    }
-
-    @Override
-    public void markVerified() {
-        verified = true;
-    }
-
-    @Override
-    public StubInfo stubInfo() {
-        return stubInfo;
-    }
-
-    @Override
-    public void markStubbed(StubInfo stubInfo) {
-        this.stubInfo = stubInfo;
-    }
-
-    @Override
-    public boolean isIgnoredForVerification() {
-        return isIgnoredForVerification;
-    }
-
-    @Override
-    public void ignoreForVerification() {
-        isIgnoredForVerification = true;
-    }
-
-    @Override
-    public Object getMock() {
-        return mock;
-    }
-
-    @Override
-    public Method getMethod() {
-        return method.getJavaMethod();
-    }
-
-    @Override
-    public Object[] getArguments() {
-        return arguments;
-    }
-
-    @Override
-    @SuppressWarnings("unchecked")
-    public <T> T getArgument(int index) {
-        return (T) arguments[index];
-    }
-
-    @Override
-    public Object callRealMethod() throws Throwable {
-        if (!superMethod.isInvokable()) {
-            throw cannotCallAbstractRealMethod();
-        }
-        return superMethod.invoke();
-    }
-
-    @Override
-    public int hashCode() {
-        // TODO SF we need to provide hash code implementation so that there are no unexpected,
-        //         slight perf issues
-        return 1;
-    }
-
-    @Override
-    public boolean equals(Object o) {
-        if (o == null || !o.getClass().equals(this.getClass())) {
-            return false;
-        }
-        InterceptedInvocation other = (InterceptedInvocation) o;
-        return this.mock.equals(other.mock)
-                && this.method.equals(other.method)
-                && this.equalArguments(other.arguments);
-    }
-
-    private boolean equalArguments(Object[] arguments) {
-        return Arrays.equals(arguments, this.arguments);
-    }
-
-    @Override
-    public String toString() {
-        return new PrintSettings().print(ArgumentsProcessor.argumentsToMatchers(getArguments()),
-                this);
-    }
-
-    interface SuperMethod extends Serializable {
-        boolean isInvokable();
-
-        Object invoke() throws Throwable;
-    }
-}
diff --git a/dexmaker-mockito/src/main/java/com/android/dx/mockito/InvocationHandlerAdapter.java b/dexmaker-mockito/src/main/java/com/android/dx/mockito/InvocationHandlerAdapter.java
index bab9265..4a95e05 100644
--- a/dexmaker-mockito/src/main/java/com/android/dx/mockito/InvocationHandlerAdapter.java
+++ b/dexmaker-mockito/src/main/java/com/android/dx/mockito/InvocationHandlerAdapter.java
@@ -18,21 +18,14 @@ package com.android.dx.mockito;
 
 import com.android.dx.stock.ProxyBuilder;
 
-import org.mockito.internal.creation.DelegatingMethod;
-import org.mockito.internal.debugging.LocationImpl;
-import org.mockito.internal.exceptions.VerificationAwareInvocation;
-import org.mockito.internal.invocation.ArgumentsProcessor;
-import org.mockito.internal.progress.SequenceNumber;
-import org.mockito.invocation.Invocation;
-import org.mockito.invocation.Location;
+import org.mockito.Mockito;
+import org.mockito.invocation.InvocationFactory.RealMethodBehavior;
 import org.mockito.invocation.MockHandler;
-import org.mockito.invocation.StubInfo;
 
 import java.lang.reflect.InvocationHandler;
 import java.lang.reflect.Method;
-import java.lang.reflect.Modifier;
 
-import static org.mockito.internal.exceptions.Reporter.cannotCallAbstractRealMethod;
+import static org.mockito.Mockito.withSettings;
 
 /**
  * Handles proxy method invocations to dexmaker's InvocationHandler by calling
@@ -46,17 +39,24 @@ final class InvocationHandlerAdapter implements InvocationHandler {
     }
 
     @Override
-    public Object invoke(Object proxy, Method method, Object[] args) throws Throwable {
+    public Object invoke(final Object proxy, final Method method, final Object[] rawArgs)
+            throws Throwable {
         // args can be null if the method invoked has no arguments, but Mockito expects a non-null array
-        args = args != null ? args : new Object[0];
+        Object[] args = rawArgs != null ? rawArgs : new Object[0];
         if (isEqualsMethod(method)) {
             return proxy == args[0];
         } else if (isHashCodeMethod(method)) {
             return System.identityHashCode(proxy);
         }
 
-        return handler.handle(new ProxyInvocation(proxy, method, args, new DelegatingMethod
-                (method), SequenceNumber.next(), new LocationImpl()));
+        return handler.handle(Mockito.framework().getInvocationFactory().createInvocation(proxy,
+                withSettings().build(proxy.getClass().getSuperclass()), method,
+                new RealMethodBehavior() {
+            @Override
+            public Object call() throws Throwable {
+                return ProxyBuilder.callSuper(proxy, method, rawArgs);
+            }
+        }, args));
     }
 
     public MockHandler getHandler() {
@@ -77,108 +77,4 @@ final class InvocationHandlerAdapter implements InvocationHandler {
         return method.getName().equals("hashCode")
                 && method.getParameterTypes().length == 0;
     }
-
-    /**
-     * Invocation on a proxy
-     */
-    private class ProxyInvocation implements Invocation, VerificationAwareInvocation {
-        private final Object proxy;
-        private final Method method;
-        private final Object[] rawArgs;
-        private final int sequenceNumber;
-        private final Location location;
-        private final Object[] args;
-
-        private StubInfo stubInfo;
-        private boolean isIgnoredForVerification;
-        private boolean verified;
-
-        private ProxyInvocation(Object proxy, Method method, Object[] rawArgs, DelegatingMethod
-                mockitoMethod, int sequenceNumber, Location location) {
-            this.rawArgs = rawArgs;
-            this.proxy = proxy;
-            this.method = method;
-            this.sequenceNumber = sequenceNumber;
-            this.location = location;
-            args = ArgumentsProcessor.expandArgs(mockitoMethod, rawArgs);
-        }
-
-        @Override
-        public Object getMock() {
-            return proxy;
-        }
-
-        @Override
-        public Method getMethod() {
-            return method;
-        }
-
-        @Override
-        public Object[] getArguments() {
-            return args;
-        }
-
-        @Override
-        public <T> T getArgument(int index) {
-            return (T)args[index];
-        }
-
-        @Override
-        public Object callRealMethod() throws Throwable {
-            if (Modifier.isAbstract(method.getModifiers())) {
-                throw cannotCallAbstractRealMethod();
-            }
-            return ProxyBuilder.callSuper(proxy, method, rawArgs);
-        }
-
-        @Override
-        public boolean isVerified() {
-            return verified || isIgnoredForVerification;
-        }
-
-        @Override
-        public int getSequenceNumber() {
-            return sequenceNumber;
-        }
-
-        @Override
-        public Location getLocation() {
-            return location;
-        }
-
-        @Override
-        public Object[] getRawArguments() {
-            return rawArgs;
-        }
-
-        @Override
-        public Class<?> getRawReturnType() {
-            return method.getReturnType();
-        }
-
-        @Override
-        public void markVerified() {
-            verified = true;
-        }
-
-        @Override
-        public StubInfo stubInfo() {
-            return stubInfo;
-        }
-
-        @Override
-        public void markStubbed(StubInfo stubInfo) {
-            this.stubInfo = stubInfo;
-        }
-
-        @Override
-        public boolean isIgnoredForVerification() {
-            return isIgnoredForVerification;
-        }
-
-        @Override
-        public void ignoreForVerification() {
-            isIgnoredForVerification = true;
-        }
-    }
 }
diff --git a/dexmaker-tests/src/androidTest/java/com/android/dx/AnnotationIdTest.java b/dexmaker-tests/src/androidTest/java/com/android/dx/AnnotationIdTest.java
new file mode 100644
index 0000000..43731ee
--- /dev/null
+++ b/dexmaker-tests/src/androidTest/java/com/android/dx/AnnotationIdTest.java
@@ -0,0 +1,400 @@
+/*
+ * Copyright (C) 2017 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package com.android.dx;
+
+import android.support.test.InstrumentationRegistry;
+import org.junit.After;
+import org.junit.Before;
+import org.junit.Test;
+
+import java.io.File;
+import java.lang.annotation.*;
+import java.lang.reflect.Method;
+import java.util.HashMap;
+import java.util.List;
+import java.util.Map;
+
+import static com.android.dx.TypeId.*;
+import static java.lang.reflect.Modifier.PUBLIC;
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.fail;
+
+public final class AnnotationIdTest {
+
+    /**
+     *  Method Annotation definition for test
+     */
+    @Retention(RetentionPolicy.RUNTIME)
+    @Target({ElementType.METHOD})
+    @interface MethodAnnotation {
+        boolean elementBoolean() default false;
+        byte elementByte() default Byte.MIN_VALUE;
+        char elementChar() default 'a';
+        double elementDouble() default Double.MIN_NORMAL;
+        float elementFloat() default Float.MIN_NORMAL;
+        int elementInt() default Integer.MIN_VALUE;
+        long elementLong() default Long.MIN_VALUE;
+        short elementShort() default Short.MIN_VALUE;
+        String elementString() default "foo";
+        ElementEnum elementEnum() default ElementEnum.INSTANCE_0;
+        Class<?> elementClass() default Object.class;
+    }
+
+    enum ElementEnum {
+        INSTANCE_0,
+        INSTANCE_1,
+    }
+
+    private DexMaker dexMaker;
+    private static TypeId<?> GENERATED = TypeId.get("LGenerated;");
+    private static final Map<TypeId<?>, Class<?>> TYPE_TO_PRIMITIVE = new HashMap<>();
+    static {
+        TYPE_TO_PRIMITIVE.put(BOOLEAN, boolean.class);
+        TYPE_TO_PRIMITIVE.put(BYTE, byte.class);
+        TYPE_TO_PRIMITIVE.put(CHAR, char.class);
+        TYPE_TO_PRIMITIVE.put(DOUBLE, double.class);
+        TYPE_TO_PRIMITIVE.put(FLOAT, float.class);
+        TYPE_TO_PRIMITIVE.put(INT, int.class);
+        TYPE_TO_PRIMITIVE.put(LONG, long.class);
+        TYPE_TO_PRIMITIVE.put(SHORT, short.class);
+        TYPE_TO_PRIMITIVE.put(VOID, void.class);
+    }
+
+    @Before
+    public void setUp() {
+        init();
+    }
+
+    /**
+     *  Test adding a method annotation with new value of Boolean element.
+     */
+    @Test
+    public void addMethodAnnotationWithBooleanElement() throws Exception {
+        MethodId<?, Void> methodId = generateVoidMethod(TypeId.BOOLEAN);
+        AnnotationId.Element element = new AnnotationId.Element("elementBoolean", true);
+        addAnnotationToMethod(methodId, element);
+
+        Annotation[] methodAnnotations = getMethodAnnotations(methodId);
+        assertEquals(methodAnnotations.length, 1);
+
+        Boolean elementBoolean = ((MethodAnnotation)methodAnnotations[0]).elementBoolean();
+        assertEquals(true, elementBoolean);
+    }
+
+    /**
+     *  Test adding a method annotation with new value of Byte element.
+     */
+    @Test
+    public void addMethodAnnotationWithByteElement() throws Exception {
+        MethodId<?, Void> methodId = generateVoidMethod(TypeId.BYTE);
+        AnnotationId.Element element = new AnnotationId.Element("elementByte", Byte.MAX_VALUE);
+        addAnnotationToMethod(methodId, element);
+
+        Annotation[] methodAnnotations = getMethodAnnotations(methodId);
+        assertEquals(methodAnnotations.length, 1);
+
+        byte elementByte = ((MethodAnnotation)methodAnnotations[0]).elementByte();
+        assertEquals(Byte.MAX_VALUE, elementByte);
+    }
+
+    /**
+     *  Test adding a method annotation with new value of Char element.
+     */
+    @Test
+    public void addMethodAnnotationWithCharElement() throws Exception {
+        MethodId<?, Void> methodId = generateVoidMethod(TypeId.CHAR);
+        AnnotationId.Element element = new AnnotationId.Element("elementChar", 'X');
+        addAnnotationToMethod(methodId, element);
+
+        Annotation[] methodAnnotations = getMethodAnnotations(methodId);
+        assertEquals(methodAnnotations.length, 1);
+
+        char elementChar = ((MethodAnnotation)methodAnnotations[0]).elementChar();
+        assertEquals('X', elementChar);
+    }
+
+    /**
+     *  Test adding a method annotation with new value of Double element.
+     */
+    @Test
+    public void addMethodAnnotationWithDoubleElement() throws Exception {
+        MethodId<?, Void> methodId = generateVoidMethod(TypeId.DOUBLE);
+        AnnotationId.Element element = new AnnotationId.Element("elementDouble", Double.NaN);
+        addAnnotationToMethod(methodId, element);
+
+        Annotation[] methodAnnotations = getMethodAnnotations(methodId);
+        assertEquals(methodAnnotations.length, 1);
+
+        double elementDouble = ((MethodAnnotation)methodAnnotations[0]).elementDouble();
+        assertEquals(Double.NaN, elementDouble, 0);
+    }
+
+    /**
+     *  Test adding a method annotation with new value of Float element.
+     */
+    @Test
+    public void addMethodAnnotationWithFloatElement() throws Exception {
+        MethodId<?, Void> methodId = generateVoidMethod(TypeId.FLOAT);
+        AnnotationId.Element element = new AnnotationId.Element("elementFloat", Float.NaN);
+        addAnnotationToMethod(methodId, element);
+
+        Annotation[] methodAnnotations = getMethodAnnotations(methodId);
+        assertEquals(methodAnnotations.length, 1);
+
+        float elementFloat = ((MethodAnnotation)methodAnnotations[0]).elementFloat();
+        assertEquals(Float.NaN, elementFloat, 0);
+    }
+
+    /**
+     *  Test adding a method annotation with new value of Int element.
+     */
+    @Test
+    public void addMethodAnnotationWithIntElement() throws Exception {
+        MethodId<?, Void> methodId = generateVoidMethod(TypeId.INT);
+        AnnotationId.Element element = new AnnotationId.Element("elementInt", Integer.MAX_VALUE);
+        addAnnotationToMethod(methodId, element);
+
+        Annotation[] methodAnnotations = getMethodAnnotations(methodId);
+        assertEquals(methodAnnotations.length, 1);
+
+        int elementInt = ((MethodAnnotation)methodAnnotations[0]).elementInt();
+        assertEquals(Integer.MAX_VALUE, elementInt);
+    }
+
+    /**
+     *  Test adding a method annotation with new value of Long element.
+     */
+    @Test
+    public void addMethodAnnotationWithLongElement() throws Exception {
+        MethodId<?, Void> methodId = generateVoidMethod(TypeId.LONG);
+        AnnotationId.Element element = new AnnotationId.Element("elementLong", Long.MAX_VALUE);
+        addAnnotationToMethod(methodId, element);
+
+        Annotation[] methodAnnotations = getMethodAnnotations(methodId);
+        assertEquals(methodAnnotations.length, 1);
+
+        long elementLong = ((MethodAnnotation)methodAnnotations[0]).elementLong();
+        assertEquals(Long.MAX_VALUE, elementLong);
+    }
+
+    /**
+     *  Test adding a method annotation with new value of Short element.
+     */
+    @Test
+    public void addMethodAnnotationWithShortElement() throws Exception {
+        MethodId<?, Void> methodId = generateVoidMethod(TypeId.SHORT);
+        AnnotationId.Element element = new AnnotationId.Element("elementShort", Short.MAX_VALUE);
+        addAnnotationToMethod(methodId, element);
+
+        Annotation[] methodAnnotations = getMethodAnnotations(methodId);
+        assertEquals(methodAnnotations.length, 1);
+
+        short elementShort = ((MethodAnnotation)methodAnnotations[0]).elementShort();
+        assertEquals(Short.MAX_VALUE, elementShort);
+    }
+
+    /**
+     *  Test adding a method annotation with new value of String element.
+     */
+    @Test
+    public void addMethodAnnotationWithStingElement() throws Exception {
+        MethodId<?, Void> methodId = generateVoidMethod(TypeId.STRING);
+        AnnotationId.Element element = new AnnotationId.Element("elementString", "hello");
+        addAnnotationToMethod(methodId, element);
+
+        Annotation[] methodAnnotations = getMethodAnnotations(methodId);
+        assertEquals(methodAnnotations.length, 1);
+
+        String elementString = ((MethodAnnotation)methodAnnotations[0]).elementString();
+        assertEquals("hello", elementString);
+    }
+
+    /**
+     *  Test adding a method annotation with new value of Enum element.
+     */
+    @Test
+    public void addMethodAnnotationWithEnumElement() throws Exception {
+        MethodId<?, Void> methodId = generateVoidMethod(TypeId.get(Enum.class));
+        AnnotationId.Element element = new AnnotationId.Element("elementEnum", ElementEnum.INSTANCE_1);
+        addAnnotationToMethod(methodId, element);
+
+        Annotation[] methodAnnotations = getMethodAnnotations(methodId);
+        assertEquals(methodAnnotations.length, 1);
+
+        ElementEnum elementEnum = ((MethodAnnotation)methodAnnotations[0]).elementEnum();
+        assertEquals(ElementEnum.INSTANCE_1, elementEnum);
+    }
+
+    /**
+     *  Test adding a method annotation with new value of Class element.
+     */
+    @Test
+    public void addMethodAnnotationWithClassElement() throws Exception {
+        MethodId<?, Void> methodId = generateVoidMethod(TypeId.get(AnnotationId.class));
+        AnnotationId.Element element = new AnnotationId.Element("elementClass", AnnotationId.class);
+        addAnnotationToMethod(methodId, element);
+
+        Annotation[] methodAnnotations = getMethodAnnotations(methodId);
+        assertEquals(methodAnnotations.length, 1);
+
+        Class<?> elementClass = ((MethodAnnotation)methodAnnotations[0]).elementClass();
+        assertEquals(AnnotationId.class, elementClass);
+    }
+
+    /**
+     *  Test adding a method annotation with new multiple values of an element.
+     */
+    @Test
+    public void addMethodAnnotationWithMultiElements() throws Exception {
+        MethodId<?, Void> methodId = generateVoidMethod();
+        AnnotationId.Element element1 = new AnnotationId.Element("elementClass", AnnotationId.class);
+        AnnotationId.Element element2 = new AnnotationId.Element("elementEnum", ElementEnum.INSTANCE_1);
+        AnnotationId.Element[] elements = {element1, element2};
+        addAnnotationToMethod(methodId, elements);
+
+        Annotation[] methodAnnotations = getMethodAnnotations(methodId);
+        assertEquals(methodAnnotations.length, 1);
+
+        ElementEnum elementEnum = ((MethodAnnotation)methodAnnotations[0]).elementEnum();
+        assertEquals(ElementEnum.INSTANCE_1, elementEnum);
+        Class<?> elementClass = ((MethodAnnotation)methodAnnotations[0]).elementClass();
+        assertEquals(AnnotationId.class, elementClass);
+    }
+
+    /**
+     *  Test adding a method annotation with duplicate values of an element. The previous value will
+     *  be replaced by latter one.
+     */
+    @Test
+    public void addMethodAnnotationWithDuplicateElements() throws Exception {
+        MethodId<?, Void> methodId = generateVoidMethod();
+        AnnotationId.Element element1 = new AnnotationId.Element("elementEnum", ElementEnum.INSTANCE_1);
+        AnnotationId.Element element2 = new AnnotationId.Element("elementEnum", ElementEnum.INSTANCE_0);
+        addAnnotationToMethod(methodId, element1, element2);
+
+        Annotation[] methodAnnotations = getMethodAnnotations(methodId);
+        assertEquals(methodAnnotations.length, 1);
+
+        ElementEnum elementEnum = ((MethodAnnotation)methodAnnotations[0]).elementEnum();
+        assertEquals(ElementEnum.INSTANCE_0, elementEnum);
+    }
+
+
+    /**
+     *  Test adding a method annotation with new array value of an element. It's not supported yet.
+     */
+    @Test
+    public void addMethodAnnotationWithArrayElementValue() {
+        try {
+            MethodId<?, Void> methodId = generateVoidMethod();
+            int[] a = {1, 2};
+            AnnotationId.Element element = new AnnotationId.Element("elementInt", a);
+            addAnnotationToMethod(methodId, element);
+            fail();
+        } catch (UnsupportedOperationException e) {
+            System.out.println(e);
+        }
+    }
+
+    /**
+     *  Test adding a method annotation with new TypeId value of an element. It's not supported yet.
+     */
+    @Test
+    public void addMethodAnnotationWithTypeIdElementValue() {
+        try {
+            MethodId<?, Void> methodId = generateVoidMethod();
+            AnnotationId.Element element = new AnnotationId.Element("elementInt", INT);
+            addAnnotationToMethod(methodId, element);
+            fail();
+        } catch (UnsupportedOperationException e) {
+            System.out.println(e);
+        }
+    }
+
+    @After
+    public void tearDown() {
+    }
+
+    /**
+     *  Internal methods
+     */
+    private void init() {
+        clearDataDirectory();
+
+        dexMaker = new DexMaker();
+        dexMaker.declare(GENERATED, "Generated.java", PUBLIC, TypeId.OBJECT);
+    }
+
+    private void clearDataDirectory() {
+        for (File f : getDataDirectory().listFiles()) {
+            if (f.getName().endsWith(".jar") || f.getName().endsWith(".dex")) {
+                f.delete();
+            }
+        }
+    }
+
+    private static File getDataDirectory() {
+        String dataDir = InstrumentationRegistry.getTargetContext().getApplicationInfo().dataDir;
+        return new File(dataDir + "/cache" );
+    }
+
+    private MethodId<?, Void> generateVoidMethod(TypeId<?>... parameters) {
+        MethodId<?, Void> methodId = GENERATED.getMethod(VOID, "call", parameters);
+        Code code = dexMaker.declare(methodId, PUBLIC);
+        code.returnVoid();
+        return methodId;
+    }
+
+    private void addAnnotationToMethod(MethodId<?, Void> methodId, AnnotationId.Element... elements) {
+        TypeId<MethodAnnotation> annotationTypeId = TypeId.get(MethodAnnotation.class);
+        AnnotationId<?, MethodAnnotation> annotationId = AnnotationId.get(GENERATED, annotationTypeId, ElementType.METHOD);
+        for (AnnotationId.Element element : elements) {
+            annotationId.set(element);
+        }
+        annotationId.addToMethod(dexMaker, methodId);
+    }
+
+    private Annotation[] getMethodAnnotations(MethodId<?, Void> methodId) throws Exception {
+        Class<?> generatedClass = generateAndLoad();
+        Class<?>[] parameters = getMethodParameters(methodId);
+        Method method = generatedClass.getMethod(methodId.getName(), parameters);
+        return method.getAnnotations();
+    }
+
+    private Class<?>[] getMethodParameters(MethodId<?, Void> methodId) throws ClassNotFoundException {
+        List<TypeId<?>> paras = methodId.getParameters();
+        Class<?>[] p = null;
+        if (paras.size() > 0) {
+            p = new Class<?>[paras.size()];
+            for (int i = 0; i < paras.size(); i++) {
+                p[i] = TYPE_TO_PRIMITIVE.get(paras.get(i));
+                if (p[i] == null) {
+                    String name = paras.get(i).getName().replace('/', '.');
+                    if (name.charAt(0) == 'L') {
+                        name = name.substring(1, name.length()-1);
+                    }
+                    p[i] = Class.forName(name);
+                }
+            }
+        }
+        return p;
+    }
+
+    private Class<?> generateAndLoad() throws Exception {
+        return dexMaker.generateAndLoad(getClass().getClassLoader(), getDataDirectory())
+                .loadClass("Generated");
+    }
+}
diff --git a/dexmaker/src/main/java/com/android/dx/AnnotationId.java b/dexmaker/src/main/java/com/android/dx/AnnotationId.java
new file mode 100644
index 0000000..bcc201b
--- /dev/null
+++ b/dexmaker/src/main/java/com/android/dx/AnnotationId.java
@@ -0,0 +1,257 @@
+/*
+ * Copyright (C) 2017 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package com.android.dx;
+
+import com.android.dx.dex.file.ClassDefItem;
+import com.android.dx.rop.annotation.Annotation;
+import com.android.dx.rop.annotation.AnnotationVisibility;
+import com.android.dx.rop.annotation.Annotations;
+import com.android.dx.rop.annotation.NameValuePair;
+import com.android.dx.rop.cst.*;
+
+import java.lang.annotation.ElementType;
+import java.util.HashMap;
+
+/**
+ * Identifies an annotation on a program element, see {@link java.lang.annotation.ElementType}.
+ *
+ * Currently it is only targeting Class, Method, Field and Parameter because those are supported by
+ * {@link com.android.dx.dex.file.AnnotationsDirectoryItem} so far.
+ *
+ * <p><strong>NOTE:</strong>
+ * So far it only supports adding method annotation. The annotation of class, field and parameter
+ * will be implemented later.
+ *
+ * <p><strong>WARNING:</strong>
+ * The declared element of an annotation type should either have a default value or be set a value via
+ * {@code AnnotationId.set(Element)}. Otherwise it will incur
+ * {@link java.lang.annotation.IncompleteAnnotationException} when accessing the annotation element
+ * through reflection. The example is as follows:
+ * <pre>
+ *     {@code @Retention(RetentionPolicy.RUNTIME)}
+ *     {@code @Target({ElementType.METHOD})}
+ *     {@code @interface MethodAnnotation {
+ *                boolean elementBoolean();
+ *                // boolean elementBoolean() default false;
+ *            }
+ *
+ *            TypeId<?> GENERATED = TypeId.get("LGenerated;");
+ *            MethodId<?, Void> methodId = GENERATED.getMethod(VOID, "call");
+ *            Code code = dexMaker.declare(methodId, PUBLIC);
+ *            code.returnVoid();
+ *
+ *            TypeId<MethodAnnotation> annotationTypeId = TypeId.get(MethodAnnotation.class);
+ *            AnnotationId<?, MethodAnnotation> annotationId = AnnotationId.get(GENERATED,
+ *              annotationTypeId, ElementType.METHOD);
+ *
+ *            AnnotationId.Element element = new AnnotationId.Element("elementBoolean", true);
+ *            annotationId.set(element);
+ *            annotationId.addToMethod(dexMaker, methodId);
+ *     }
+ * </pre>
+ *
+ * @param <D> the type that declares the program element.
+ * @param <V> the annotation type. It should be a known type before compile.
+ */
+public final class AnnotationId<D, V> {
+    private final TypeId<D> declaringType;
+    private final TypeId<V> type;
+    /** The type of program element to be annotated */
+    private final ElementType annotatedElement;
+    /** The elements this annotation holds */
+    private final HashMap<String, NameValuePair> elements;
+
+    private AnnotationId(TypeId<D> declaringType, TypeId<V> type, ElementType annotatedElement) {
+        this.declaringType = declaringType;
+        this.type = type;
+        this.annotatedElement = annotatedElement;
+        this.elements = new HashMap<>();
+    }
+
+    /**
+     *  Construct an instance. It initially contains no elements.
+     *
+     * @param declaringType the type declaring the program element.
+     * @param type the annotation type.
+     * @param annotatedElement the program element type to be annotated.
+     * @return an annotation {@code AnnotationId<D,V>} instance.
+     */
+    public static <D, V> AnnotationId<D, V> get(TypeId<D> declaringType, TypeId<V> type,
+                                                ElementType annotatedElement) {
+        if (annotatedElement != ElementType.TYPE &&
+                annotatedElement != ElementType.METHOD &&
+                annotatedElement != ElementType.FIELD &&
+                annotatedElement != ElementType.PARAMETER) {
+            throw new IllegalArgumentException("element type is not supported to annotate yet.");
+        }
+
+        return new AnnotationId<>(declaringType, type, annotatedElement);
+    }
+
+    /**
+     * Set an annotation element of this instance.
+     * If there is a preexisting element with the same name, it will be
+     * replaced by this method.
+     *
+     * @param element {@code non-null;} the annotation element to be set.
+     */
+    public void set(Element element) {
+        if (element == null) {
+            throw new NullPointerException("element == null");
+        }
+
+        CstString pairName = new CstString(element.getName());
+        Constant pairValue = Element.toConstant(element.getValue());
+        NameValuePair nameValuePair = new NameValuePair(pairName, pairValue);
+        elements.put(element.getName(), nameValuePair);
+    }
+
+    /**
+     * Add this annotation to a method.
+     *
+     * @param dexMaker DexMaker instance.
+     * @param method Method to be added to.
+     */
+    public void addToMethod(DexMaker dexMaker, MethodId<?, ?> method) {
+        if (annotatedElement != ElementType.METHOD) {
+            throw new IllegalStateException("This annotation is not for method");
+        }
+
+        if (method.declaringType != declaringType) {
+            throw new IllegalArgumentException("Method" + method + "'s declaring type is inconsistent with" + this);
+        }
+
+        ClassDefItem classDefItem = dexMaker.getTypeDeclaration(declaringType).toClassDefItem();
+
+        if (classDefItem == null) {
+            throw new NullPointerException("No class defined item is found");
+        } else {
+            CstMethodRef cstMethodRef = method.constant;
+
+            if (cstMethodRef == null) {
+                throw new NullPointerException("Method reference is NULL");
+            } else {
+                // Generate CstType
+                CstType cstType = CstType.intern(type.ropType);
+
+                // Generate Annotation
+                Annotation annotation = new Annotation(cstType, AnnotationVisibility.RUNTIME);
+
+                // Add generated annotation
+                Annotations annotations = new Annotations();
+                for (NameValuePair nvp : elements.values()) {
+                    annotation.add(nvp);
+                }
+                annotations.add(annotation);
+                classDefItem.addMethodAnnotations(cstMethodRef, annotations, dexMaker.getDexFile());
+            }
+        }
+    }
+
+    /**
+     *  A wrapper of <code>NameValuePair</code> represents a (name, value) pair used as the contents
+     *  of an annotation.
+     *
+     *  An {@code Element} instance is stored in {@code AnnotationId.elements} by calling {@code
+     *  AnnotationId.set(Element)}.
+     *
+     *  <p><strong>WARNING: </strong></p>
+     *  the name should be exact same as the annotation element declared in the annotation type
+     *  which is referred by field {@code AnnotationId.type},otherwise the annotation will fail
+     *  to add and {@code java.lang.reflect.Method.getAnnotations()} will return nothing.
+     *
+     */
+    public static final class Element {
+        /** {@code non-null;} the name */
+        private final String name;
+        /** {@code non-null;} the value */
+        private final Object value;
+
+        /**
+         * Construct an instance.
+         *
+         * @param name {@code non-null;} the name
+         * @param value {@code non-null;} the value
+         */
+        public Element(String name, Object value) {
+            if (name == null) {
+                throw new NullPointerException("name == null");
+            }
+
+            if (value == null) {
+                throw new NullPointerException("value == null");
+            }
+            this.name = name;
+            this.value = value;
+        }
+
+        public String getName() {
+            return name;
+        }
+
+        public Object getValue() {
+            return value;
+        }
+
+        /** {@inheritDoc} */
+        @Override
+        public String toString() {
+            return "[" + name + ", " + value + "]";
+        }
+
+        /** {@inheritDoc} */
+        @Override
+        public int hashCode() {
+            return name.hashCode() * 31 + value.hashCode();
+        }
+
+        /** {@inheritDoc} */
+        @Override
+        public boolean equals(Object other) {
+            if (! (other instanceof Element)) {
+                return false;
+            }
+
+            Element otherElement = (Element) other;
+
+            return name.equals(otherElement.name)
+                    && value.equals(otherElement.value);
+        }
+
+        /**
+         *  Convert a value of an element to a {@code Constant}.
+         *  <p><strong>Warning:</strong> Array or TypeId value is not supported yet.
+         *
+         * @param value an annotation element value.
+         * @return a Constant
+         */
+        static Constant toConstant(Object value) {
+            Class clazz = value.getClass();
+            if (clazz.isEnum()) {
+                CstString descriptor = new CstString(TypeId.get(clazz).getName());
+                CstString name = new CstString(((Enum)value).name());
+                CstNat cstNat = new CstNat(name, descriptor);
+                return new CstEnumRef(cstNat);
+            } else if (clazz.isArray()) {
+                throw new UnsupportedOperationException("Array is not supported yet");
+            } else if (value instanceof TypeId) {
+                throw new UnsupportedOperationException("TypeId is not supported yet");
+            } else {
+                return  Constants.getConstant(value);
+            }
+        }
+    }
+}
diff --git a/dexmaker/src/main/java/com/android/dx/DexMaker.java b/dexmaker/src/main/java/com/android/dx/DexMaker.java
index f10ad8e..02baa9b 100644
--- a/dexmaker/src/main/java/com/android/dx/DexMaker.java
+++ b/dexmaker/src/main/java/com/android/dx/DexMaker.java
@@ -198,6 +198,7 @@ import static java.lang.reflect.Modifier.STATIC;
 public final class DexMaker {
     private final Map<TypeId<?>, TypeDeclaration> types = new LinkedHashMap<>();
     private ClassLoader sharedClassLoader;
+    private DexFile outputDex;
 
     /**
      * Creates a new {@code DexMaker} instance, which can be used to create a
@@ -206,7 +207,7 @@ public final class DexMaker {
     public DexMaker() {
     }
 
-    private TypeDeclaration getTypeDeclaration(TypeId<?> type) {
+    TypeDeclaration getTypeDeclaration(TypeId<?> type) {
         TypeDeclaration result = types.get(type);
         if (result == null) {
             result = new TypeDeclaration(type);
@@ -313,9 +314,11 @@ public final class DexMaker {
      * Generates a dex file and returns its bytes.
      */
     public byte[] generate() {
-        DexOptions options = new DexOptions();
-        options.targetApiLevel = DexFormat.API_NO_EXTENDED_OPCODES;
-        DexFile outputDex = new DexFile(options);
+        if (outputDex == null) {
+            DexOptions options = new DexOptions();
+            options.targetApiLevel = DexFormat.API_NO_EXTENDED_OPCODES;
+            outputDex = new DexFile(options);
+        }
 
         for (TypeDeclaration typeDeclaration : types.values()) {
             outputDex.add(typeDeclaration.toClassDefItem());
@@ -451,7 +454,16 @@ public final class DexMaker {
         return generateClassLoader(result, dexCache, parent);
     }
 
-    private static class TypeDeclaration {
+    DexFile getDexFile() {
+        if (outputDex == null) {
+            DexOptions options = new DexOptions();
+            options.targetApiLevel = DexFormat.API_NO_EXTENDED_OPCODES;
+            outputDex = new DexFile(options);
+        }
+        return outputDex;
+    }
+
+    static class TypeDeclaration {
         private final TypeId<?> type;
 
         /** declared state */
@@ -460,6 +472,7 @@ public final class DexMaker {
         private TypeId<?> supertype;
         private String sourceFile;
         private TypeList interfaces;
+        private ClassDefItem classDefItem;
 
         private final Map<FieldId, FieldDeclaration> fields = new LinkedHashMap<>();
         private final Map<MethodId, MethodDeclaration> methods = new LinkedHashMap<>();
@@ -479,27 +492,29 @@ public final class DexMaker {
 
             CstType thisType = type.constant;
 
-            ClassDefItem out = new ClassDefItem(thisType, flags, supertype.constant,
-                    interfaces.ropTypes, new CstString(sourceFile));
-
-            for (MethodDeclaration method : methods.values()) {
-                EncodedMethod encoded = method.toEncodedMethod(dexOptions);
-                if (method.isDirect()) {
-                    out.addDirectMethod(encoded);
-                } else {
-                    out.addVirtualMethod(encoded);
+            if (classDefItem == null) {
+                classDefItem = new ClassDefItem(thisType, flags, supertype.constant,
+                        interfaces.ropTypes, new CstString(sourceFile));
+
+                for (MethodDeclaration method : methods.values()) {
+                    EncodedMethod encoded = method.toEncodedMethod(dexOptions);
+                    if (method.isDirect()) {
+                        classDefItem.addDirectMethod(encoded);
+                    } else {
+                        classDefItem.addVirtualMethod(encoded);
+                    }
                 }
-            }
-            for (FieldDeclaration field : fields.values()) {
-                EncodedField encoded = field.toEncodedField();
-                if (field.isStatic()) {
-                    out.addStaticField(encoded, Constants.getConstant(field.staticValue));
-                } else {
-                    out.addInstanceField(encoded);
+                for (FieldDeclaration field : fields.values()) {
+                    EncodedField encoded = field.toEncodedField();
+                    if (field.isStatic()) {
+                        classDefItem.addStaticField(encoded, Constants.getConstant(field.staticValue));
+                    } else {
+                        classDefItem.addInstanceField(encoded);
+                    }
                 }
             }
 
-            return out;
+            return classDefItem;
         }
     }