ART: Correctly handle temporary classes in class-load events (4/3)

When a temporary class is given out in a ClassLoad event, all stored
references need to be fixed up before publishing a ClassPrepare event.

This CL handles objects stored in the heap as referents.

Bug: 31684920
Test: m test-art-host-run-test-912-classes
Change-Id: If140ecae675cd7bc648f622eaf200f8ad8b15438
diff --git a/runtime/openjdkjvmti/ti_class.cc b/runtime/openjdkjvmti/ti_class.cc
index fc4b6fe..7ca233f 100644
--- a/runtime/openjdkjvmti/ti_class.cc
+++ b/runtime/openjdkjvmti/ti_class.cc
@@ -52,6 +52,7 @@
 #include "mirror/class_ext.h"
 #include "mirror/object_reference.h"
 #include "mirror/object-inl.h"
+#include "mirror/reference.h"
 #include "runtime.h"
 #include "runtime_callbacks.h"
 #include "ScopedLocalRef.h"
@@ -463,8 +464,17 @@
         }
       }
 
+      void operator()(art::ObjPtr<art::mirror::Class> klass ATTRIBUTE_UNUSED,
+                      art::ObjPtr<art::mirror::Reference> reference) const
+          REQUIRES_SHARED(art::Locks::mutator_lock_) {
+        art::mirror::Object* val = reference->GetReferent();
+        if (val == input_) {
+          reference->SetReferent<false>(output_);
+        }
+      }
+
       void VisitRoot(art::mirror::CompressedReference<art::mirror::Object>* root ATTRIBUTE_UNUSED)
-      const {
+          const {
         LOG(FATAL) << "Unreachable";
       }
 
@@ -478,7 +488,7 @@
         HeapFixupVisitor* hfv = reinterpret_cast<HeapFixupVisitor*>(arg);
 
         // Visit references, not native roots.
-        obj->VisitReferences<false>(*hfv, art::VoidFunctor());
+        obj->VisitReferences<false>(*hfv, *hfv);
       }
 
      private:
diff --git a/test/912-classes/classes.cc b/test/912-classes/classes.cc
index c92e49f..3ccfe86 100644
--- a/test/912-classes/classes.cc
+++ b/test/912-classes/classes.cc
@@ -433,9 +433,13 @@
 class ClassLoadPrepareEquality {
  public:
   static constexpr const char* kClassName = "LMain$ClassE;";
-  static constexpr const char* kStorageClassName = "Main$ClassF";
   static constexpr const char* kStorageFieldName = "STATIC";
   static constexpr const char* kStorageFieldSig = "Ljava/lang/Object;";
+  static constexpr const char* kStorageWeakFieldName = "WEAK";
+  static constexpr const char* kStorageWeakFieldSig = "Ljava/lang/ref/Reference;";
+  static constexpr const char* kWeakClassName = "java/lang/ref/WeakReference";
+  static constexpr const char* kWeakInitSig = "(Ljava/lang/Object;)V";
+  static constexpr const char* kWeakGetSig = "()Ljava/lang/Object;";
 
   static void JNICALL ClassLoadCallback(jvmtiEnv* jenv,
                                         JNIEnv* jni_env,
@@ -472,6 +476,8 @@
 
   static void SetOrCompare(JNIEnv* jni_env, jobject value, bool set) {
     CHECK(storage_class_ != nullptr);
+
+    // Simple direct storage.
     jfieldID field = jni_env->GetStaticFieldID(storage_class_, kStorageFieldName, kStorageFieldSig);
     CHECK(field != nullptr);
 
@@ -482,6 +488,36 @@
       ScopedLocalRef<jobject> stored(jni_env, jni_env->GetStaticObjectField(storage_class_, field));
       CHECK(jni_env->IsSameObject(value, stored.get()));
     }
+
+    // Storage as a reference.
+    ScopedLocalRef<jclass> weak_ref_class(jni_env, jni_env->FindClass(kWeakClassName));
+    CHECK(weak_ref_class.get() != nullptr);
+    jfieldID weak_field = jni_env->GetStaticFieldID(storage_class_,
+                                                    kStorageWeakFieldName,
+                                                    kStorageWeakFieldSig);
+    CHECK(weak_field != nullptr);
+    if (set) {
+      // Create a WeakReference.
+      jmethodID weak_init = jni_env->GetMethodID(weak_ref_class.get(), "<init>", kWeakInitSig);
+      CHECK(weak_init != nullptr);
+      ScopedLocalRef<jobject> weak_obj(jni_env, jni_env->NewObject(weak_ref_class.get(),
+                                                                   weak_init,
+                                                                   value));
+      CHECK(weak_obj.get() != nullptr);
+      jni_env->SetStaticObjectField(storage_class_, weak_field, weak_obj.get());
+      CHECK(!jni_env->ExceptionCheck());
+    } else {
+      // Check the reference value.
+      jmethodID get_referent = jni_env->GetMethodID(weak_ref_class.get(), "get", kWeakGetSig);
+      CHECK(get_referent != nullptr);
+      ScopedLocalRef<jobject> weak_obj(jni_env, jni_env->GetStaticObjectField(storage_class_,
+                                                                              weak_field));
+      CHECK(weak_obj.get() != nullptr);
+      ScopedLocalRef<jobject> weak_referent(jni_env, jni_env->CallObjectMethod(weak_obj.get(),
+                                                                               get_referent));
+      CHECK(weak_referent.get() != nullptr);
+      CHECK(jni_env->IsSameObject(value, weak_referent.get()));
+    }
   }
 
   static void CheckFound() {
diff --git a/test/912-classes/src/Main.java b/test/912-classes/src/Main.java
index 52a5194..005074f 100644
--- a/test/912-classes/src/Main.java
+++ b/test/912-classes/src/Main.java
@@ -14,6 +14,7 @@
  * limitations under the License.
  */
 
+import java.lang.ref.Reference;
 import java.lang.reflect.Constructor;
 import java.lang.reflect.Proxy;
 import java.util.Arrays;
@@ -433,6 +434,7 @@
 
   public static class ClassF {
     public static Object STATIC = null;
+    public static Reference<Object> WEAK = null;
   }
 
   private static final String DEX1 = System.getenv("DEX_LOCATION") + "/912-classes.jar";