From 8c2348864666127b34f1d3bcbf2b0a069f733f90 Mon Sep 17 00:00:00 2001
From: yhs0602 <jourhyang123@gmail.com>
Date: Wed, 15 May 2024 23:32:42 +0900
Subject: [PATCH] :bug: Implement lazy loading

---
 src/main/kotlin/vm/classloader/DexClassLoader.kt | 13 ++++++++++---
 1 file changed, 10 insertions(+), 3 deletions(-)

diff --git a/src/main/kotlin/vm/classloader/DexClassLoader.kt b/src/main/kotlin/vm/classloader/DexClassLoader.kt
index b85b4da..d1c4ce2 100644
--- a/src/main/kotlin/vm/classloader/DexClassLoader.kt
+++ b/src/main/kotlin/vm/classloader/DexClassLoader.kt
@@ -16,6 +16,7 @@ class DexClassLoader(
         TypeId("Ljava/lang/Object;") to ObjectType
     )
 
+    private val typeIdToClassDef = mutableMapOf<TypeId, ParsedClass>()
     private val loadedMockedClasses = mutableSetOf<Class<*>>()
 
     private fun loadClass(clazz: Class<*>): Type {
@@ -87,16 +88,22 @@ class DexClassLoader(
         mockedClasses.forEach { (typeId, mockedClass) ->
             loadClass(typeId, mockedClass)
         }
-        // load dex classes
+        // lazy load dex classes
         dexFiles.forEach { dexFile ->
             dexFile.classDefs.forEach { classDef ->
                 val typeId = classDef.classDef.typeId
-                loadClass(typeId, classDef)
+                typeIdToClassDef[typeId] = classDef
+//                loadClass(typeId, classDef)
             }
         }
     }
 
     fun getClass(typeId: TypeId): Type {
-        return loadedTypes[typeId] ?: error("Class ${typeId.descriptor} not loaded")
+        val loadedClass = loadedTypes[typeId]
+        if (loadedClass != null) {
+            return loadedClass
+        }
+        val classDef = typeIdToClassDef[typeId] ?: error("Class ${typeId.descriptor} not found")
+        return loadClass(typeId, classDef)
     }
 }
\ No newline at end of file