[baksmali] Load full dex minimally for commenting class & method hierarchy of large no of dex files

Author: REAndroid

src/main/java/com/reandroid/apkeditor/smali/SmaliDecompiler.java

@@ -50,6 +50,7 @@ public class SmaliDecompiler implements DexDecoder {
     private SmaliWriterSetting smaliWriterSetting;
     private Opcodes mCurrentOpcodes;
     private APKLogger apkLogger;
+    private boolean mDexForCommentLoaded;
 
     public SmaliDecompiler(TableBlock tableBlock, DecompileOptions decompileOptions) {
         this.tableBlock = tableBlock;
@@ -74,11 +75,16 @@ public void decodeDex(ApkModule apkModule, File mainDirectory) throws IOExceptio
         }
         DexDirectory directory = (DexDirectory) apkModule.getTag(DexDirectory.class);
         if (directory == null) {
-            if (apkModule.listDexFiles().size() > decompileOptions.loadDex) {
+            int size = apkModule.listDexFiles().size();
+            logMessage("Dex files: " + size);
+            if (size > decompileOptions.loadDex) {
+                if (size < decompileOptions.loadDex * 5) {
+                    loadMinimalDexForComment(apkModule);
+                }
                 DexDecoder.super.decodeDex(apkModule, mainDirectory);
                 return;
             }
-            logMessage("Loading full dex files: " + apkModule.listDexFiles().size());
+            logMessage("Loading full dex files ...");
             Predicate<SectionType<?>> filter;
             if (decompileOptions.noDexDebug) {
                 filter = sectionType -> sectionType != SectionType.DEBUG_INFO;
@@ -106,6 +112,19 @@ public void decodeDex(ApkModule apkModule, File mainDirectory) throws IOExceptio
         }
     }
 
+    private void loadMinimalDexForComment(ApkModule apkModule) throws IOException {
+        String commentLevel = decompileOptions.commentLevel;
+        if (!DecompileOptions.COMMENT_LEVEL_DETAIL.equals(commentLevel) &&
+                !DecompileOptions.COMMENT_LEVEL_FULL.equals(commentLevel)) {
+            return;
+        }
+        logMessage("Loading basic structures of dex ...");
+        DexDirectory dexDirectory = DexDirectory.fromZip(
+                apkModule.getZipEntryMap(), SectionType.minimal());
+        mDexForCommentLoaded = false;
+        getSmaliWriterSetting(dexDirectory);
+        mDexForCommentLoaded = true;
+    }
     private void disassembleWithJesusFrekeLib(DexFileInputSource inputSource, File mainDir) throws IOException {
         File dir = toOutDir(inputSource, mainDir);
         BaksmaliOptions options = new BaksmaliOptions();
@@ -135,8 +154,10 @@ private void disassembleWithInternalDexLib(DexFileInputSource inputSource, File
         SmaliWriter smaliWriter = new SmaliWriter();
         smaliWriter.setWriterSetting(setting);
         dexFile.writeSmali(smaliWriter, toSmaliRoot(mainDir));
-        setting.clearClassComments();
-        setting.clearMethodComments();
+        if (!mDexForCommentLoaded) {
+            setting.clearClassComments();
+            setting.clearMethodComments();
+        }
         dexFile.close();
     }
     private void writeDexCache(DexFileInputSource inputSource, File mainDir) throws IOException {
@@ -191,12 +212,14 @@ public ResourceComment getComment() {
 
     private SmaliWriterSetting getSmaliWriterSetting(DexClassRepository classRepository) {
         SmaliWriterSetting setting = getSmaliWriterSetting();
-        setting.clearClassComments();
-        setting.clearMethodComments();
-        String commentLevel = decompileOptions.commentLevel;
-        if (DecompileOptions.COMMENT_LEVEL_DETAIL.equals(commentLevel) || DecompileOptions.COMMENT_LEVEL_FULL.equals(commentLevel)) {
-            setting.addClassComments(classRepository);
-            setting.addMethodComments(classRepository);
+        if (!mDexForCommentLoaded) {
+            setting.clearClassComments();
+            setting.clearMethodComments();
+            String commentLevel = decompileOptions.commentLevel;
+            if (DecompileOptions.COMMENT_LEVEL_DETAIL.equals(commentLevel) || DecompileOptions.COMMENT_LEVEL_FULL.equals(commentLevel)) {
+                setting.addClassComments(classRepository);
+                setting.addMethodComments(classRepository);
+            }
         }
         return setting;
     }