Make InitializeStaticStorageFromCode handle requests for uninitialized storage from <clinit>

Change-Id: I8562ad4fdb33c02c575ddc8986e49ee37c566cfd
diff --git a/src/class_linker.cc b/src/class_linker.cc
index 8ddcc5f..7cd320e 100644
--- a/src/class_linker.cc
+++ b/src/class_linker.cc
@@ -1390,9 +1390,16 @@
   if (klass == NULL) {
     UNIMPLEMENTED(FATAL) << "throw exception due to unresolved class";
   }
+  // If we are the <clinit> of this class, just return our storage.
+  //
+  // Do not set the DexCache InitializedStaticStorage, since that
+  // implies <clinit> has finished running.
+  if (klass == referrer->GetDeclaringClass() && referrer->GetName()->Equals("<clinit>")) {
+    return klass;
+  }
   if (!class_linker->EnsureInitialized(klass)) {
     CHECK(Thread::Current()->IsExceptionPending());
-    UNIMPLEMENTED(FATAL) << "throw exception due to class initializtion problem";
+    UNIMPLEMENTED(FATAL) << "throw exception due to class initialization problem";
   }
   referrer->GetDexCacheInitializedStaticStorage()->Set(type_idx, klass);
   return klass;
diff --git a/src/class_linker_test.cc b/src/class_linker_test.cc
index 92b3fca..d30f4bf 100644
--- a/src/class_linker_test.cc
+++ b/src/class_linker_test.cc
@@ -566,19 +566,27 @@
 }
 
 TEST_F(ClassLinkerTest, InitializeStaticStorageFromCode) {
-  // pretend we are trying to ensure we have initilized storage for Static from Statics.<clinit>
+  // pretend we are trying to get the static storage for the Statics class.
+
+  // case 1, get the uninitialized storage from Statics.<clinit>
+  // case 2, get the initialized storage from Statics.getS8
+
   const ClassLoader* class_loader = LoadDex("Statics");
   const DexFile* dex_file = ClassLoader::GetClassPath(class_loader)[0];
   CHECK(dex_file != NULL);
 
   Class* Statics = class_linker_->FindClass("LStatics;", class_loader);
   Method* clinit = Statics->FindDirectMethod("<clinit>", "()V");
+  Method* getS8 = Statics->FindDirectMethod("getS8", "()Ljava/lang/Object;");
   uint32_t type_idx = FindTypeIdxByDescriptor(*dex_file, "LStatics;");
 
   EXPECT_TRUE(clinit->GetDexCacheInitializedStaticStorage()->Get(type_idx) == NULL);
-  StaticStorageBase* storage = class_linker_->InitializeStaticStorageFromCode(type_idx, clinit);
-  EXPECT_TRUE(storage != NULL);
-  EXPECT_EQ(storage, clinit->GetDexCacheInitializedStaticStorage()->Get(type_idx));
+  StaticStorageBase* uninit = class_linker_->InitializeStaticStorageFromCode(type_idx, clinit);
+  EXPECT_TRUE(uninit != NULL);
+  EXPECT_TRUE(clinit->GetDexCacheInitializedStaticStorage()->Get(type_idx) == NULL);
+  StaticStorageBase* init = class_linker_->InitializeStaticStorageFromCode(type_idx, getS8);
+  EXPECT_TRUE(init != NULL);
+  EXPECT_EQ(init, clinit->GetDexCacheInitializedStaticStorage()->Get(type_idx));
 }
 
 }  // namespace art