4 files changed, 42 insertions, 4 deletions

diff --git a/runtime/java_vm_ext.cc b/runtime/java_vm_ext.cc
index b795d72..3142089 100644
--- a/runtime/java_vm_ext.cc
+++ b/runtime/java_vm_ext.cc
@@ -562,6 +562,11 @@ mirror::Object* JavaVMExt::DecodeGlobal(Thread* self, IndirectRef ref) {
   return globals_.SynchronizedGet(self, &globals_lock_, ref);
 }
 
+void JavaVMExt::UpdateGlobal(Thread* self, IndirectRef ref, mirror::Object* result) {
+  WriterMutexLock mu(self, globals_lock_);
+  globals_.Update(ref, result);
+}
+
 mirror::Object* JavaVMExt::DecodeWeakGlobal(Thread* self, IndirectRef ref) {
   MutexLock mu(self, weak_globals_lock_);
   while (UNLIKELY(!allow_new_weak_globals_)) {
@@ -570,6 +575,11 @@ mirror::Object* JavaVMExt::DecodeWeakGlobal(Thread* self, IndirectRef ref) {
   return weak_globals_.Get(ref);
 }
 
+void JavaVMExt::UpdateWeakGlobal(Thread* self, IndirectRef ref, mirror::Object* result) {
+  MutexLock mu(self, weak_globals_lock_);
+  weak_globals_.Update(ref, result);
+}
+
 void JavaVMExt::DumpReferenceTables(std::ostream& os) {
   Thread* self = Thread::Current();
   {
diff --git a/runtime/java_vm_ext.h b/runtime/java_vm_ext.h
index deec6a9..594027c 100644
--- a/runtime/java_vm_ext.h
+++ b/runtime/java_vm_ext.h
@@ -125,9 +125,17 @@ class JavaVMExt : public JavaVM {
   mirror::Object* DecodeGlobal(Thread* self, IndirectRef ref)
       SHARED_LOCKS_REQUIRED(Locks::mutator_lock_);
 
+  void UpdateGlobal(Thread* self, IndirectRef ref, mirror::Object* result)
+      SHARED_LOCKS_REQUIRED(Locks::mutator_lock_)
+      LOCKS_EXCLUDED(globals_lock_);
+
   mirror::Object* DecodeWeakGlobal(Thread* self, IndirectRef ref)
       SHARED_LOCKS_REQUIRED(Locks::mutator_lock_);
 
+  void UpdateWeakGlobal(Thread* self, IndirectRef ref, mirror::Object* result)
+      SHARED_LOCKS_REQUIRED(Locks::mutator_lock_)
+      LOCKS_EXCLUDED(weak_globals_lock_);
+
   const JNIInvokeInterface* GetUncheckedFunctions() const {
     return unchecked_functions_;
   }
diff --git a/runtime/reflection.cc b/runtime/reflection.cc
index d321d27..f8c7081 100644
--- a/runtime/reflection.cc
+++ b/runtime/reflection.cc
@@ -464,7 +464,7 @@ JValue InvokeWithVarArgs(const ScopedObjectAccessAlreadyRunnable& soa, jobject o
   InvokeWithArgArray(soa, method, &arg_array, &result, shorty);
   if (is_string_init) {
     // For string init, remap original receiver to StringFactory result.
-    soa.Self()->GetJniEnv()->locals.Update(obj, result.GetL());
+    UpdateReference(soa.Self(), obj, result.GetL());
   }
   return result;
 }
@@ -494,7 +494,7 @@ JValue InvokeWithJValues(const ScopedObjectAccessAlreadyRunnable& soa, jobject o
   InvokeWithArgArray(soa, method, &arg_array, &result, shorty);
   if (is_string_init) {
     // For string init, remap original receiver to StringFactory result.
-    soa.Self()->GetJniEnv()->locals.Update(obj, result.GetL());
+    UpdateReference(soa.Self(), obj, result.GetL());
   }
   return result;
 }
@@ -525,7 +525,7 @@ JValue InvokeVirtualOrInterfaceWithJValues(const ScopedObjectAccessAlreadyRunnab
   InvokeWithArgArray(soa, method, &arg_array, &result, shorty);
   if (is_string_init) {
     // For string init, remap original receiver to StringFactory result.
-    soa.Self()->GetJniEnv()->locals.Update(obj, result.GetL());
+    UpdateReference(soa.Self(), obj, result.GetL());
   }
   return result;
 }
@@ -556,7 +556,7 @@ JValue InvokeVirtualOrInterfaceWithVarArgs(const ScopedObjectAccessAlreadyRunnab
   InvokeWithArgArray(soa, method, &arg_array, &result, shorty);
   if (is_string_init) {
     // For string init, remap original receiver to StringFactory result.
-    soa.Self()->GetJniEnv()->locals.Update(obj, result.GetL());
+    UpdateReference(soa.Self(), obj, result.GetL());
   }
   return result;
 }
@@ -882,4 +882,21 @@ void InvalidReceiverError(mirror::Object* o, mirror::Class* c) {
                                              actual_class_name.c_str()).c_str());
 }
 
+// This only works if there's one reference which points to the object in obj.
+// Will need to be fixed if there's cases where it's not.
+void UpdateReference(Thread* self, jobject obj, mirror::Object* result) {
+  IndirectRef ref = reinterpret_cast<IndirectRef>(obj);
+  IndirectRefKind kind = GetIndirectRefKind(ref);
+  if (kind == kLocal) {
+    self->GetJniEnv()->locals.Update(obj, result);
+  } else if (kind == kHandleScopeOrInvalid) {
+    LOG(FATAL) << "Unsupported UpdateReference for kind kHandleScopeOrInvalid";
+  } else if (kind == kGlobal) {
+    self->GetJniEnv()->vm->UpdateGlobal(self, ref, result);
+  } else {
+    DCHECK_EQ(kind, kWeakGlobal);
+    self->GetJniEnv()->vm->UpdateWeakGlobal(self, ref, result);
+  }
+}
+
 }  // namespace art
diff --git a/runtime/reflection.h b/runtime/reflection.h
index 6b5ffc7..df3b9d3 100644
--- a/runtime/reflection.h
+++ b/runtime/reflection.h
@@ -85,6 +85,9 @@ mirror::Class* GetCallingClass(Thread* self, size_t num_frames)
 void InvalidReceiverError(mirror::Object* o, mirror::Class* c)
     SHARED_LOCKS_REQUIRED(Locks::mutator_lock_);
 
+void UpdateReference(Thread* self, jobject obj, mirror::Object* result)
+    SHARED_LOCKS_REQUIRED(Locks::mutator_lock_);
+
 }  // namespace art
 
 #endif  // ART_RUNTIME_REFLECTION_H_