Speed up dex mapping lookups

src/launchpad/cli.py

@@ -8,7 +8,7 @@
 from . import __version__
 from .distribution.cli import distribution_command
 from .service import run_service
-from .size.cli import app_icon_command, size_command
+from .size.cli import app_icon_command, profile_dex_parsing_command, size_command
 from .utils.console import console
 from .utils.logging import setup_logging
 
@@ -85,6 +85,7 @@ def serve(host: str, port: int, mode: str | None, verbose: bool) -> None:
 cli.add_command(size_command)
 cli.add_command(app_icon_command)
 cli.add_command(distribution_command)
+cli.add_command(profile_dex_parsing_command)
 
 
 def main() -> None:

src/launchpad/parsers/android/dex/dex_mapping.py

@@ -21,6 +21,8 @@ class DexMappingClass:
 class DexMapping:
     def __init__(self, bytes: bytes) -> None:
         self._classes: dict[str, DexMappingClass] = {}
+        self._classes_by_deobfuscated: dict[str, DexMappingClass] = {}
+        self._classes_by_deobfuscated_fqn: dict[str, DexMappingClass] = {}
 
         classes: list[DexMappingClass] = []
         current_class: DexMappingClass | None = None
@@ -42,6 +44,9 @@ def __init__(self, bytes: bytes) -> None:
 
         for clazz in classes:
             self._classes[clazz.obfuscated_fqn] = clazz
+            self._classes_by_deobfuscated[clazz.deobfuscated_signature] = clazz
+            self._classes_by_deobfuscated_fqn[clazz.deobfuscated_fqn] = clazz
+        return None
 
     @staticmethod
     def _parse_comment(line: str, pending_file_name: str | None) -> str | None:
@@ -178,10 +183,7 @@ def lookup_obfuscated_class(self, obfuscated_class_name: str) -> DexMappingClass
         return self._classes.get(obfuscated_class_name)
 
     def lookup_deobfuscated_signature(self, deobfuscated_class_signature: str) -> DexMappingClass | None:
-        for clazz in self._classes.values():
-            if clazz.deobfuscated_signature == deobfuscated_class_signature:
-                return clazz
-        return None
+        return self._classes_by_deobfuscated.get(deobfuscated_class_signature)
 
     def deobfuscate_method(self, class_name: str, obfuscated_method_name: str) -> str | None:
         """Deobfuscate a method name for a given class.
@@ -195,19 +197,16 @@ def deobfuscate_method(self, class_name: str, obfuscated_method_name: str) -> st
         """
         # First try to find the class by obfuscated name
         clazz = self.lookup_obfuscated_class(class_name)
+
+        # Try to find by deobfuscated signature
         if clazz is None:
-            # Try to find by deobfuscated signature
             clazz = self.lookup_deobfuscated_signature(class_name)
-            if clazz is None:
-                # Try to match by deobfuscated FQN
-                for c in self._classes.values():
-                    if c.deobfuscated_fqn == class_name:
-                        clazz = c
-                        break
-                if clazz is None:
-                    return None
-
-        if clazz.methods is None:
+
+        # Try to match by deobfuscated FQN
+        if clazz is None:
+            clazz = self._classes_by_deobfuscated_fqn.get(class_name)
+
+        if clazz is None or clazz.methods is None:
             return None
 
         return clazz.methods.get(obfuscated_method_name)
@@ -224,19 +223,16 @@ def deobfuscate_field(self, class_name: str, obfuscated_field_name: str) -> str
         """
         # First try to find the class by obfuscated name
         clazz = self.lookup_obfuscated_class(class_name)
+
+        # Try to find by deobfuscated signature
         if clazz is None:
-            # Try to find by deobfuscated signature
             clazz = self.lookup_deobfuscated_signature(class_name)
-            if clazz is None:
-                # Try to match by deobfuscated FQN
-                for c in self._classes.values():
-                    if c.deobfuscated_fqn == class_name:
-                        clazz = c
-                        break
-                if clazz is None:
-                    return None
-
-        if clazz.fields is None:
+
+        # Try to match by deobfuscated FQN
+        if clazz is None:
+            clazz = self._classes_by_deobfuscated_fqn.get(class_name)
+
+        if clazz is None or clazz.fields is None:
             return None
 
         return clazz.fields.get(obfuscated_field_name)

src/launchpad/size/cli.py

@@ -6,6 +6,8 @@
 from rich.table import Table
 
 from launchpad.artifacts.artifact_factory import ArtifactFactory
+from launchpad.parsers.android.dex.dex_file_parser import DexFileParser
+from launchpad.parsers.android.dex.dex_mapping import DexMapping
 from launchpad.size.models.android import AndroidAnalysisResults
 from launchpad.size.models.apple import AppleAnalysisResults
 from launchpad.size.models.common import BaseAnalysisResults, FileAnalysis
@@ -141,6 +143,20 @@ def app_icon_command(
         raise click.Abort()
 
 
+@click.command(name="profile-dex-parsing")
+@click.argument("input_path", type=click.Path(exists=True, path_type=Path), metavar="DEX_PATH")
+@click.argument("mapping_path", required=False, type=click.Path(exists=True, path_type=Path), metavar="MAPPING_PATH")
+def profile_dex_parsing_command(input_path: Path, mapping_path: Path | None = None):
+    mapping = None
+    if mapping_path:
+        with open(mapping_path, "rb") as f:
+            mapping = DexMapping(f.read())
+    with open(input_path, "rb") as f:
+        parser = DexFileParser(f.read(), mapping)
+
+    parser.get_class_definitions()
+
+
 def _print_results_as_table(results: BaseAnalysisResults) -> None:
     if isinstance(results, AndroidAnalysisResults):
         _print_android_table_output(results)