Merge "Refactor JVMTI ClassFileLoadHook handling"

6 files changed, 276 insertions, 169 deletions

diff --git a/openjdkjvmti/ti_class.cc b/openjdkjvmti/ti_class.cc
index 1868631020..4d54d756d5 100644
--- a/openjdkjvmti/ti_class.cc
+++ b/openjdkjvmti/ti_class.cc
@@ -184,74 +184,27 @@ struct ClassCallback : public art::ClassLoadCallback {
       return;
     }
 
-    // Strip the 'L' and ';' from the descriptor
-    std::string name(std::string(descriptor).substr(1, strlen(descriptor) - 2));
-
     art::Thread* self = art::Thread::Current();
-    art::JNIEnvExt* env = self->GetJniEnv();
-    ScopedLocalRef<jobject> loader(
-        env, class_loader.IsNull() ? nullptr : env->AddLocalReference<jobject>(class_loader.Get()));
-    std::unique_ptr<FixedUpDexFile> dex_file_copy(FixedUpDexFile::Create(initial_dex_file,
-                                                                         descriptor));
-
-    // Go back to native.
-    art::ScopedThreadSuspension sts(self, art::ThreadState::kNative);
-    // Call all Non-retransformable agents.
-    jint post_no_redefine_len = 0;
-    unsigned char* post_no_redefine_dex_data = nullptr;
-    std::unique_ptr<const unsigned char, FakeJvmtiDeleter<const unsigned char>>
-        post_no_redefine_unique_ptr(nullptr, FakeJvmtiDeleter<const unsigned char>());
-    event_handler->DispatchEvent<ArtJvmtiEvent::kClassFileLoadHookNonRetransformable>(
-        self,
-        static_cast<JNIEnv*>(env),
-        static_cast<jclass>(nullptr),  // The class doesn't really exist yet so send null.
-        loader.get(),
-        name.c_str(),
-        static_cast<jobject>(nullptr),  // Android doesn't seem to have protection domains
-        static_cast<jint>(dex_file_copy->Size()),
-        static_cast<const unsigned char*>(dex_file_copy->Begin()),
-        static_cast<jint*>(&post_no_redefine_len),
-        static_cast<unsigned char**>(&post_no_redefine_dex_data));
-    if (post_no_redefine_dex_data == nullptr) {
-      DCHECK_EQ(post_no_redefine_len, 0);
-      post_no_redefine_dex_data = const_cast<unsigned char*>(dex_file_copy->Begin());
-      post_no_redefine_len = dex_file_copy->Size();
-    } else {
-      post_no_redefine_unique_ptr =
-          std::unique_ptr<const unsigned char, FakeJvmtiDeleter<const unsigned char>>(
-              post_no_redefine_dex_data, FakeJvmtiDeleter<const unsigned char>());
-      DCHECK_GT(post_no_redefine_len, 0);
+    ArtClassDefinition def;
+    def.InitFirstLoad(descriptor, class_loader, initial_dex_file);
+
+    // Call all non-retransformable agents.
+    Transformer::TransformSingleClassDirect<ArtJvmtiEvent::kClassFileLoadHookNonRetransformable>(
+        event_handler, self, &def);
+
+    std::vector<unsigned char> post_non_retransform;
+    if (def.IsModified()) {
+      // Copy the dex data after the non-retransformable events.
+      post_non_retransform.resize(def.GetDexData().size());
+      memcpy(post_non_retransform.data(), def.GetDexData().data(), post_non_retransform.size());
     }
+
     // Call all retransformable agents.
-    jint final_len = 0;
-    unsigned char* final_dex_data = nullptr;
-    std::unique_ptr<const unsigned char, FakeJvmtiDeleter<const unsigned char>>
-        final_dex_unique_ptr(nullptr, FakeJvmtiDeleter<const unsigned char>());
-    event_handler->DispatchEvent<ArtJvmtiEvent::kClassFileLoadHookRetransformable>(
-        self,
-        static_cast<JNIEnv*>(env),
-        static_cast<jclass>(nullptr),  // The class doesn't really exist yet so send null.
-        loader.get(),
-        name.c_str(),
-        static_cast<jobject>(nullptr),  // Android doesn't seem to have protection domains
-        static_cast<jint>(post_no_redefine_len),
-        static_cast<const unsigned char*>(post_no_redefine_dex_data),
-        static_cast<jint*>(&final_len),
-        static_cast<unsigned char**>(&final_dex_data));
-    if (final_dex_data == nullptr) {
-      DCHECK_EQ(final_len, 0);
-      final_dex_data = post_no_redefine_dex_data;
-      final_len = post_no_redefine_len;
-    } else {
-      final_dex_unique_ptr =
-          std::unique_ptr<const unsigned char, FakeJvmtiDeleter<const unsigned char>>(
-              final_dex_data, FakeJvmtiDeleter<const unsigned char>());
-      DCHECK_GT(final_len, 0);
-    }
+    Transformer::TransformSingleClassDirect<ArtJvmtiEvent::kClassFileLoadHookRetransformable>(
+        event_handler, self, &def);
 
-    if (final_dex_data != dex_file_copy->Begin()) {
+    if (def.IsModified()) {
       LOG(WARNING) << "Changing class " << descriptor;
-      art::ScopedObjectAccess soa(self);
       art::StackHandleScope<2> hs(self);
       // Save the results of all the non-retransformable agents.
       // First allocate the ClassExt
@@ -268,7 +221,7 @@ struct ClassCallback : public art::ClassLoadCallback {
 
       // Allocate the byte array to store the dex file bytes in.
       art::MutableHandle<art::mirror::Object> arr(hs.NewHandle<art::mirror::Object>(nullptr));
-      if (post_no_redefine_dex_data == dex_file_copy->Begin() && name != "java/lang/Long") {
+      if (post_non_retransform.empty() && strcmp(descriptor, "Ljava/lang/Long;") != 0) {
         // we didn't have any non-retransformable agents. We can just cache a pointer to the
         // initial_dex_file. It will be kept live by the class_loader.
         jlong dex_ptr = reinterpret_cast<uintptr_t>(&initial_dex_file);
@@ -278,8 +231,8 @@ struct ClassCallback : public art::ClassLoadCallback {
       } else {
         arr.Assign(art::mirror::ByteArray::AllocateAndFill(
             self,
-            reinterpret_cast<const signed char*>(post_no_redefine_dex_data),
-            post_no_redefine_len));
+            reinterpret_cast<const signed char*>(post_non_retransform.data()),
+            post_non_retransform.size()));
       }
       if (arr.IsNull()) {
         LOG(WARNING) << "Unable to allocate memory for initial dex-file. Aborting transformation";
@@ -290,8 +243,8 @@ struct ClassCallback : public art::ClassLoadCallback {
       std::unique_ptr<const art::DexFile> dex_file(MakeSingleDexFile(self,
                                                                      descriptor,
                                                                      initial_dex_file.GetLocation(),
-                                                                     final_len,
-                                                                     final_dex_data));
+                                                                     def.GetDexData().size(),
+                                                                     def.GetDexData().data()));
       if (dex_file.get() == nullptr) {
         return;
       }
diff --git a/openjdkjvmti/ti_class_definition.cc b/openjdkjvmti/ti_class_definition.cc
index c8a3047d9a..7f2f80009a 100644
--- a/openjdkjvmti/ti_class_definition.cc
+++ b/openjdkjvmti/ti_class_definition.cc
@@ -46,29 +46,33 @@
 namespace openjdkjvmti {
 
 bool ArtClassDefinition::IsModified() const {
-  // RedefineClasses calls always are 'modified' since they need to change the original_dex_file of
+  // RedefineClasses calls always are 'modified' since they need to change the current_dex_file of
   // the class.
   if (redefined_) {
     return true;
   }
+
+  // Check to see if any change has taken place.
+  if (current_dex_file_.data() == dex_data_.data()) {
+    // no change at all.
+    return false;
+  }
+
   // Check if the dex file we want to set is the same as the current one.
   // Unfortunately we need to do this check even if no modifications have been done since it could
   // be that agents were removed in the mean-time so we still have a different dex file. The dex
   // checksum means this is likely to be fairly fast.
-  return static_cast<jint>(original_dex_file_.size()) != dex_len_ ||
-      memcmp(original_dex_file_.data(), dex_data_.get(), dex_len_) != 0;
+  return current_dex_file_.size() != dex_data_.size() ||
+      memcmp(current_dex_file_.data(), dex_data_.data(), current_dex_file_.size()) != 0;
 }
 
-jvmtiError ArtClassDefinition::InitCommon(ArtJvmTiEnv* env, jclass klass) {
-  JNIEnv* jni_env = GetJniEnv(env);
-  if (jni_env == nullptr) {
-    return ERR(INTERNAL);
-  }
-  art::ScopedObjectAccess soa(jni_env);
+jvmtiError ArtClassDefinition::InitCommon(art::Thread* self, jclass klass) {
+  art::ScopedObjectAccess soa(self);
   art::ObjPtr<art::mirror::Class> m_klass(soa.Decode<art::mirror::Class>(klass));
   if (m_klass.IsNull()) {
     return ERR(INVALID_CLASS);
   }
+  initialized_ = true;
   klass_ = klass;
   loader_ = soa.AddLocalReference<jobject>(m_klass->GetClassLoader());
   std::string descriptor_store;
@@ -79,11 +83,18 @@ jvmtiError ArtClassDefinition::InitCommon(ArtJvmTiEnv* env, jclass klass) {
   return OK;
 }
 
+static void DequickenDexFile(const art::DexFile* dex_file,
+                             const char* descriptor,
+                             /*out*/std::vector<unsigned char>* dex_data)
+    REQUIRES_SHARED(art::Locks::mutator_lock_) {
+  std::unique_ptr<FixedUpDexFile> fixed_dex_file(FixedUpDexFile::Create(*dex_file, descriptor));
+  dex_data->resize(fixed_dex_file->Size());
+  memcpy(dex_data->data(), fixed_dex_file->Begin(), fixed_dex_file->Size());
+}
+
 // Gets the data surrounding the given class.
-static jvmtiError GetDexDataForRetransformation(ArtJvmTiEnv* env,
-                                                art::Handle<art::mirror::Class> klass,
-                                                /*out*/jint* dex_data_len,
-                                                /*out*/unsigned char** dex_data)
+static void GetDexDataForRetransformation(art::Handle<art::mirror::Class> klass,
+                                          /*out*/std::vector<unsigned char>* dex_data)
     REQUIRES_SHARED(art::Locks::mutator_lock_) {
   art::StackHandleScope<3> hs(art::Thread::Current());
   art::Handle<art::mirror::ClassExt> ext(hs.NewHandle(klass->GetExtData()));
@@ -95,12 +106,9 @@ static jvmtiError GetDexDataForRetransformation(ArtJvmTiEnv* env,
         DCHECK(orig_dex->GetClass()->GetComponentType()->IsPrimitiveByte());
         art::Handle<art::mirror::ByteArray> orig_dex_bytes(
             hs.NewHandle(art::down_cast<art::mirror::ByteArray*>(orig_dex->AsArray())));
-        *dex_data_len = static_cast<jint>(orig_dex_bytes->GetLength());
-        return CopyDataIntoJvmtiBuffer(
-            env,
-            reinterpret_cast<const unsigned char*>(orig_dex_bytes->GetData()),
-            *dex_data_len,
-            /*out*/dex_data);
+        dex_data->resize(orig_dex_bytes->GetLength());
+        memcpy(dex_data->data(), orig_dex_bytes->GetData(), dex_data->size());
+        return;
       } else if (orig_dex->IsDexCache()) {
         dex_file = orig_dex->AsDexCache()->GetDexFile();
       } else {
@@ -113,7 +121,7 @@ static jvmtiError GetDexDataForRetransformation(ArtJvmTiEnv* env,
         art::JValue val;
         if (!art::UnboxPrimitiveForResult(orig_dex.Get(), prim_long_class, &val)) {
           // This should never happen.
-          return ERR(INTERNAL);
+          LOG(FATAL) << "Unable to unbox a primitive long value!";
         }
         dex_file = reinterpret_cast<const art::DexFile*>(static_cast<uintptr_t>(val.GetJ()));
       }
@@ -123,59 +131,153 @@ static jvmtiError GetDexDataForRetransformation(ArtJvmTiEnv* env,
     dex_file = &klass->GetDexFile();
   }
   std::string storage;
-  std::unique_ptr<FixedUpDexFile> fixed_dex_file(
-      FixedUpDexFile::Create(*dex_file, klass->GetDescriptor(&storage)));
-  *dex_data_len = static_cast<jint>(fixed_dex_file->Size());
-  return CopyDataIntoJvmtiBuffer(env,
-                                 fixed_dex_file->Begin(),
-                                 fixed_dex_file->Size(),
-                                 /*out*/dex_data);
+  DequickenDexFile(dex_file, klass->GetDescriptor(&storage), dex_data);
 }
 
-jvmtiError ArtClassDefinition::Init(ArtJvmTiEnv* env, jclass klass) {
-  jvmtiError res = InitCommon(env, klass);
+static bool DexNeedsDequickening(art::Handle<art::mirror::Class> klass,
+                                 /*out*/ bool* from_class_ext)
+    REQUIRES_SHARED(art::Locks::mutator_lock_) {
+  art::ObjPtr<art::mirror::ClassExt> ext(klass->GetExtData());
+  if (ext.IsNull()) {
+    // We don't seem to have ever been redefined so be conservative and say we need de-quickening.
+    *from_class_ext = false;
+    return true;
+  }
+  art::ObjPtr<art::mirror::Object> orig_dex(ext->GetOriginalDexFile());
+  if (orig_dex.IsNull()) {
+    // We don't seem to have ever been redefined so be conservative and say we need de-quickening.
+    *from_class_ext = false;
+    return true;
+  } else if (!orig_dex->IsArrayInstance()) {
+    // We were redefined but the original is held in a dex-cache or dex file. This means that the
+    // original dex file is the one from the disk, which might be quickened.
+    DCHECK(orig_dex->IsDexCache() || orig_dex->GetClass()->DescriptorEquals("Ljava/lang/Long;"));
+    *from_class_ext = true;
+    return true;
+  } else {
+    // An array instance means the original-dex-file is from a redefineClasses which cannot have any
+    // quickening, so it's fine to use directly.
+    DCHECK(orig_dex->GetClass()->GetComponentType()->IsPrimitiveByte());
+    *from_class_ext = true;
+    return false;
+  }
+}
+
+static const art::DexFile* GetQuickenedDexFile(art::Handle<art::mirror::Class> klass)
+    REQUIRES_SHARED(art::Locks::mutator_lock_) {
+  art::ObjPtr<art::mirror::ClassExt> ext(klass->GetExtData());
+  if (ext.IsNull() || ext->GetOriginalDexFile() == nullptr) {
+    return &klass->GetDexFile();
+  }
+
+  art::ObjPtr<art::mirror::Object> orig_dex(ext->GetOriginalDexFile());
+  DCHECK(!orig_dex->IsArrayInstance());
+  if (orig_dex->IsDexCache()) {
+    return orig_dex->AsDexCache()->GetDexFile();
+  }
+
+  DCHECK(orig_dex->GetClass()->DescriptorEquals("Ljava/lang/Long;"))
+      << "Expected java/lang/Long but found object of type "
+      << orig_dex->GetClass()->PrettyClass();
+  art::ObjPtr<art::mirror::Class> prim_long_class(
+      art::Runtime::Current()->GetClassLinker()->GetClassRoot(
+          art::ClassLinker::kPrimitiveLong));
+  art::JValue val;
+  if (!art::UnboxPrimitiveForResult(orig_dex.Ptr(), prim_long_class, &val)) {
+    LOG(FATAL) << "Unable to unwrap a long value!";
+  }
+  return reinterpret_cast<const art::DexFile*>(static_cast<uintptr_t>(val.GetJ()));
+}
+
+template<typename GetOriginalDexFile>
+void ArtClassDefinition::InitWithDex(GetOriginalDexFile get_original,
+                                     const art::DexFile* quick_dex) {
+  art::Thread* self = art::Thread::Current();
+  DCHECK(quick_dex != nullptr);
+  get_original(/*out*/&dex_data_memory_);
+  dex_data_ = art::ArrayRef<const unsigned char>(dex_data_memory_);
+  if (from_class_ext_) {
+    // We got initial from class_ext so the current one must have undergone redefinition so no
+    // cdex or quickening stuff.
+    // We can only do this if it's not a first load.
+    DCHECK(klass_ != nullptr);
+    const art::DexFile& cur_dex = self->DecodeJObject(klass_)->AsClass()->GetDexFile();
+    current_dex_file_ = art::ArrayRef<const unsigned char>(cur_dex.Begin(), cur_dex.Size());
+  } else {
+    // No redefinition must have ever happened so the (dequickened) cur_dex is the same as the
+    // initial dex_data. We need to copy it into another buffer to keep it around if we have a
+    // real redefinition.
+    current_dex_memory_.resize(dex_data_.size());
+    memcpy(current_dex_memory_.data(), dex_data_.data(), current_dex_memory_.size());
+    current_dex_file_ = art::ArrayRef<const unsigned char>(current_dex_memory_);
+  }
+}
+
+jvmtiError ArtClassDefinition::Init(art::Thread* self, jclass klass) {
+  jvmtiError res = InitCommon(self, klass);
   if (res != OK) {
     return res;
   }
-  unsigned char* new_data = nullptr;
-  art::Thread* self = art::Thread::Current();
   art::ScopedObjectAccess soa(self);
   art::StackHandleScope<1> hs(self);
   art::Handle<art::mirror::Class> m_klass(hs.NewHandle(self->DecodeJObject(klass)->AsClass()));
-  res = GetDexDataForRetransformation(env, m_klass, &dex_len_, &new_data);
-  if (res != OK) {
-    return res;
-  }
-  dex_data_ = MakeJvmtiUniquePtr(env, new_data);
-  if (m_klass->GetExtData() == nullptr || m_klass->GetExtData()->GetOriginalDexFile() == nullptr) {
-    // We have never redefined class this yet. Keep track of what the (de-quickened) dex file looks
-    // like so we can tell if anything has changed. Really we would like to just always do the
-    // 'else' block but the fact that we de-quickened stuff screws us over.
-    unsigned char* original_data_memory = nullptr;
-    res = CopyDataIntoJvmtiBuffer(env, dex_data_.get(), dex_len_, &original_data_memory);
-    original_dex_file_memory_ = MakeJvmtiUniquePtr(env, original_data_memory);
-    original_dex_file_ = art::ArrayRef<const unsigned char>(original_data_memory, dex_len_);
-  } else {
-    // We know that we have been redefined at least once (there is an original_dex_file set in
-    // the class) so we can just use the current dex file directly.
-    const art::DexFile& dex_file = m_klass->GetDexFile();
-    original_dex_file_ = art::ArrayRef<const unsigned char>(dex_file.Begin(), dex_file.Size());
+  if (!DexNeedsDequickening(m_klass, &from_class_ext_)) {
+    // We don't need to do any dequickening. We want to copy the data just so we don't need to deal
+    // with the GC moving it around.
+    art::ObjPtr<art::mirror::ByteArray> orig_dex(
+        m_klass->GetExtData()->GetOriginalDexFile()->AsByteArray());
+    dex_data_memory_.resize(orig_dex->GetLength());
+    memcpy(dex_data_memory_.data(), orig_dex->GetData(), dex_data_memory_.size());
+    dex_data_ = art::ArrayRef<const unsigned char>(dex_data_memory_);
+
+    // Since we are here we must not have any quickened instructions since we were redefined.
+    const art::DexFile& cur_dex = m_klass->GetDexFile();
+    DCHECK(from_class_ext_);
+    current_dex_file_ = art::ArrayRef<const unsigned char>(cur_dex.Begin(), cur_dex.Size());
+    return OK;
   }
-  return res;
+
+  // We need to dequicken stuff. This is often super slow (10's of ms). Instead we will do it
+  // dynamically.
+  const art::DexFile* quick_dex = GetQuickenedDexFile(m_klass);
+  auto get_original = [&](/*out*/std::vector<unsigned char>* dex_data)
+      REQUIRES_SHARED(art::Locks::mutator_lock_) {
+    GetDexDataForRetransformation(m_klass, dex_data);
+  };
+  InitWithDex(get_original, quick_dex);
+  return OK;
 }
 
-jvmtiError ArtClassDefinition::Init(ArtJvmTiEnv* env, const jvmtiClassDefinition& def) {
-  jvmtiError res = InitCommon(env, def.klass);
+jvmtiError ArtClassDefinition::Init(art::Thread* self, const jvmtiClassDefinition& def) {
+  jvmtiError res = InitCommon(self, def.klass);
   if (res != OK) {
     return res;
   }
-  unsigned char* new_data = nullptr;
-  original_dex_file_ = art::ArrayRef<const unsigned char>(def.class_bytes, def.class_byte_count);
+  // We are being directly redefined.
   redefined_ = true;
-  dex_len_ = def.class_byte_count;
-  res = CopyDataIntoJvmtiBuffer(env, def.class_bytes, def.class_byte_count, /*out*/ &new_data);
-  dex_data_ = MakeJvmtiUniquePtr(env, new_data);
-  return res;
+  current_dex_file_ = art::ArrayRef<const unsigned char>(def.class_bytes, def.class_byte_count);
+  dex_data_ = art::ArrayRef<const unsigned char>(def.class_bytes, def.class_byte_count);
+  return OK;
+}
+
+void ArtClassDefinition::InitFirstLoad(const char* descriptor,
+                                       art::Handle<art::mirror::ClassLoader> klass_loader,
+                                       const art::DexFile& dex_file) {
+  art::Thread* self = art::Thread::Current();
+  art::ScopedObjectAccess soa(self);
+  initialized_ = true;
+  // No Class
+  klass_ = nullptr;
+  loader_ = soa.AddLocalReference<jobject>(klass_loader.Get());
+  std::string descriptor_str(descriptor);
+  name_ = descriptor_str.substr(1, descriptor_str.size() - 2);
+  // Android doesn't really have protection domains.
+  protection_domain_ = nullptr;
+  auto get_original = [&](/*out*/std::vector<unsigned char>* dex_data)
+      REQUIRES_SHARED(art::Locks::mutator_lock_) {
+    DequickenDexFile(&dex_file, descriptor, dex_data);
+  };
+  InitWithDex(get_original, &dex_file);
 }
 
 }  // namespace openjdkjvmti
diff --git a/openjdkjvmti/ti_class_definition.h b/openjdkjvmti/ti_class_definition.h
index accc456957..00847394e7 100644
--- a/openjdkjvmti/ti_class_definition.h
+++ b/openjdkjvmti/ti_class_definition.h
@@ -48,41 +48,49 @@ class ArtClassDefinition {
         loader_(nullptr),
         name_(),
         protection_domain_(nullptr),
-        dex_len_(0),
-        dex_data_(nullptr),
-        original_dex_file_memory_(nullptr),
-        original_dex_file_(),
-        redefined_(false) {}
-
-  jvmtiError Init(ArtJvmTiEnv* env, jclass klass);
-  jvmtiError Init(ArtJvmTiEnv* env, const jvmtiClassDefinition& def);
+        dex_data_memory_(),
+        dex_data_(),
+        current_dex_file_(),
+        redefined_(false),
+        from_class_ext_(false),
+        initialized_(false) {}
+
+  void InitFirstLoad(const char* descriptor,
+                     art::Handle<art::mirror::ClassLoader> klass_loader,
+                     const art::DexFile& dex_file);
+  jvmtiError Init(art::Thread* self, jclass klass);
+  jvmtiError Init(art::Thread* self, const jvmtiClassDefinition& def);
 
   ArtClassDefinition(ArtClassDefinition&& o) = default;
   ArtClassDefinition& operator=(ArtClassDefinition&& o) = default;
 
-  void SetNewDexData(ArtJvmTiEnv* env, jint new_dex_len, unsigned char* new_dex_data) {
+  void SetNewDexData(jint new_dex_len, unsigned char* new_dex_data) {
     DCHECK(IsInitialized());
     if (new_dex_data == nullptr) {
       return;
-    } else if (new_dex_data != dex_data_.get() || new_dex_len != dex_len_) {
-      dex_len_ = new_dex_len;
-      dex_data_ = MakeJvmtiUniquePtr(env, new_dex_data);
+    } else {
+      art::ArrayRef<const unsigned char> new_data(new_dex_data, new_dex_len);
+      if (new_data != dex_data_) {
+        dex_data_memory_.resize(new_dex_len);
+        memcpy(dex_data_memory_.data(), new_dex_data, new_dex_len);
+        dex_data_ = art::ArrayRef<const unsigned char>(dex_data_memory_);
+      }
     }
   }
 
   art::ArrayRef<const unsigned char> GetNewOriginalDexFile() const {
     DCHECK(IsInitialized());
     if (redefined_) {
-      return original_dex_file_;
+      return current_dex_file_;
     } else {
       return art::ArrayRef<const unsigned char>();
     }
   }
 
-  bool IsModified() const;
+  bool IsModified() const REQUIRES_SHARED(art::Locks::mutator_lock_);
 
   bool IsInitialized() const {
-    return klass_ != nullptr;
+    return initialized_;
   }
 
   jclass GetClass() const {
@@ -107,22 +115,43 @@ class ArtClassDefinition {
 
   art::ArrayRef<const unsigned char> GetDexData() const {
     DCHECK(IsInitialized());
-    return art::ArrayRef<const unsigned char>(dex_data_.get(), dex_len_);
+    return dex_data_;
   }
 
  private:
-  jvmtiError InitCommon(ArtJvmTiEnv* env, jclass klass);
+  jvmtiError InitCommon(art::Thread* self, jclass klass);
+
+  template<typename GetOriginalDexFile>
+  void InitWithDex(GetOriginalDexFile get_original, const art::DexFile* quick_dex)
+      REQUIRES_SHARED(art::Locks::mutator_lock_);
 
   jclass klass_;
   jobject loader_;
   std::string name_;
   jobject protection_domain_;
-  jint dex_len_;
-  JvmtiUniquePtr<unsigned char> dex_data_;
-  JvmtiUniquePtr<unsigned char> original_dex_file_memory_;
-  art::ArrayRef<const unsigned char> original_dex_file_;
+
+  // A unique_ptr to the current dex_data if it needs to be cleaned up.
+  std::vector<unsigned char> dex_data_memory_;
+
+  // A ref to the current dex data. This is either dex_data_memory_, or current_dex_file_. This is
+  // what the dex file will be turned into.
+  art::ArrayRef<const unsigned char> dex_data_;
+
+  // This is only used if we failed to create a mmap to store the dequickened data
+  std::vector<unsigned char> current_dex_memory_;
+
+  // This is a dequickened version of what is loaded right now. It is either current_dex_memory_ (if
+  // no other redefinition has ever happened to this) or the current dex_file_ directly (if this
+  // class has been redefined, thus it cannot have any quickened stuff).
+  art::ArrayRef<const unsigned char> current_dex_file_;
+
   bool redefined_;
 
+  // If we got the initial dex_data_ from a class_ext
+  bool from_class_ext_;
+
+  bool initialized_;
+
   DISALLOW_COPY_AND_ASSIGN(ArtClassDefinition);
 };
 
diff --git a/openjdkjvmti/ti_redefine.cc b/openjdkjvmti/ti_redefine.cc
index 717b2ba669..c3fb946b9a 100644
--- a/openjdkjvmti/ti_redefine.cc
+++ b/openjdkjvmti/ti_redefine.cc
@@ -351,15 +351,14 @@ jvmtiError Redefiner::RedefineClasses(ArtJvmTiEnv* env,
     memcpy(class_bytes_copy, definitions[i].class_bytes, definitions[i].class_byte_count);
 
     ArtClassDefinition def;
-    res = def.Init(env, definitions[i]);
+    res = def.Init(self, definitions[i]);
     if (res != OK) {
       return res;
     }
     def_vector.push_back(std::move(def));
   }
   // Call all the transformation events.
-  jvmtiError res = Transformer::RetransformClassesDirect(env,
-                                                         event_handler,
+  jvmtiError res = Transformer::RetransformClassesDirect(event_handler,
                                                          self,
                                                          &def_vector);
   if (res != OK) {
diff --git a/openjdkjvmti/transform.cc b/openjdkjvmti/transform.cc
index 3a5fcccf35..8445ecab16 100644
--- a/openjdkjvmti/transform.cc
+++ b/openjdkjvmti/transform.cc
@@ -63,27 +63,46 @@
 
 namespace openjdkjvmti {
 
+// Initialize templates.
+template
+void Transformer::TransformSingleClassDirect<ArtJvmtiEvent::kClassFileLoadHookNonRetransformable>(
+    EventHandler* event_handler, art::Thread* self, /*in-out*/ArtClassDefinition* def);
+template
+void Transformer::TransformSingleClassDirect<ArtJvmtiEvent::kClassFileLoadHookRetransformable>(
+    EventHandler* event_handler, art::Thread* self, /*in-out*/ArtClassDefinition* def);
+
+template<ArtJvmtiEvent kEvent>
+void Transformer::TransformSingleClassDirect(EventHandler* event_handler,
+                                             art::Thread* self,
+                                             /*in-out*/ArtClassDefinition* def) {
+  static_assert(kEvent == ArtJvmtiEvent::kClassFileLoadHookNonRetransformable ||
+                kEvent == ArtJvmtiEvent::kClassFileLoadHookRetransformable,
+                "bad event type");
+  jint new_len = -1;
+  unsigned char* new_data = nullptr;
+  art::ArrayRef<const unsigned char> dex_data = def->GetDexData();
+  event_handler->DispatchEvent<kEvent>(
+      self,
+      static_cast<JNIEnv*>(self->GetJniEnv()),
+      def->GetClass(),
+      def->GetLoader(),
+      def->GetName().c_str(),
+      def->GetProtectionDomain(),
+      static_cast<jint>(dex_data.size()),
+      dex_data.data(),
+      /*out*/&new_len,
+      /*out*/&new_data);
+  def->SetNewDexData(new_len, new_data);
+}
+
 jvmtiError Transformer::RetransformClassesDirect(
-      ArtJvmTiEnv* env,
       EventHandler* event_handler,
       art::Thread* self,
       /*in-out*/std::vector<ArtClassDefinition>* definitions) {
   for (ArtClassDefinition& def : *definitions) {
-    jint new_len = -1;
-    unsigned char* new_data = nullptr;
-    art::ArrayRef<const unsigned char> dex_data = def.GetDexData();
-    event_handler->DispatchEvent<ArtJvmtiEvent::kClassFileLoadHookRetransformable>(
-        self,
-        GetJniEnv(env),
-        def.GetClass(),
-        def.GetLoader(),
-        def.GetName().c_str(),
-        def.GetProtectionDomain(),
-        static_cast<jint>(dex_data.size()),
-        dex_data.data(),
-        /*out*/&new_len,
-        /*out*/&new_data);
-    def.SetNewDexData(env, new_len, new_data);
+    TransformSingleClassDirect<ArtJvmtiEvent::kClassFileLoadHookRetransformable>(event_handler,
+                                                                                 self,
+                                                                                 &def);
   }
   return OK;
 }
@@ -120,13 +139,13 @@ jvmtiError Transformer::RetransformClasses(ArtJvmTiEnv* env,
       return ERR(UNMODIFIABLE_CLASS);
     }
     ArtClassDefinition def;
-    res = def.Init(env, classes[i]);
+    res = def.Init(self, classes[i]);
     if (res != OK) {
       return res;
     }
     definitions.push_back(std::move(def));
   }
-  res = RetransformClassesDirect(env, event_handler, self, &definitions);
+  res = RetransformClassesDirect(event_handler, self, &definitions);
   if (res != OK) {
     return res;
   }
diff --git a/openjdkjvmti/transform.h b/openjdkjvmti/transform.h
index 6bbe60a91f..f43af174f0 100644
--- a/openjdkjvmti/transform.h
+++ b/openjdkjvmti/transform.h
@@ -48,8 +48,13 @@ jvmtiError GetClassLocation(ArtJvmTiEnv* env, jclass klass, /*out*/std::string*
 
 class Transformer {
  public:
+  template<ArtJvmtiEvent kEvent>
+  static void TransformSingleClassDirect(
+      EventHandler* event_handler,
+      art::Thread* self,
+      /*in-out*/ArtClassDefinition* def);
+
   static jvmtiError RetransformClassesDirect(
-      ArtJvmTiEnv* env,
       EventHandler* event_handler,
       art::Thread* self,
       /*in-out*/std::vector<ArtClassDefinition>* definitions);