From fbe51b960b685db455a558729956bdaf982a2af9 Mon Sep 17 00:00:00 2001
From: Romain Guy <romainguy@curious-creature.com>
Date: Mon, 14 Oct 2024 11:15:58 -0700
Subject: [PATCH] Add code size in bytes to method dumps in OAT format

---
 .../dev/romainguy/kotlin/explorer/code/CodeBuilder.kt | 11 ++++++++---
 .../dev/romainguy/kotlin/explorer/code/DataModels.kt  |  2 +-
 .../romainguy/kotlin/explorer/oat/OatDumpParser.kt    |  8 ++++++--
 3 files changed, 15 insertions(+), 6 deletions(-)

diff --git a/src/jvmMain/kotlin/dev/romainguy/kotlin/explorer/code/CodeBuilder.kt b/src/jvmMain/kotlin/dev/romainguy/kotlin/explorer/code/CodeBuilder.kt
index 97ea2726..6db13e92 100644
--- a/src/jvmMain/kotlin/dev/romainguy/kotlin/explorer/code/CodeBuilder.kt
+++ b/src/jvmMain/kotlin/dev/romainguy/kotlin/explorer/code/CodeBuilder.kt
@@ -20,6 +20,8 @@ import androidx.collection.IntIntPair
 import androidx.collection.IntObjectMap
 import androidx.collection.mutableIntIntMapOf
 import androidx.collection.mutableIntObjectMapOf
+import androidx.compose.ui.util.fastForEach
+import androidx.compose.ui.util.fastSumBy
 import kotlin.math.max
 
 fun buildCode(codeStyle: CodeStyle = CodeStyle(), builderAction: CodeBuilder.() -> Unit): CodeBuilder {
@@ -85,14 +87,17 @@ class CodeBuilder(private val codeStyle: CodeStyle) {
         sb.append(" ".repeat(codeStyle.indent))
         writeLine(method.header)
 
-        sb.append("  ".repeat(codeStyle.indent))
+        val indent = "  ".repeat(codeStyle.indent)
+
+        sb.append(indent)
+        val codeSize = method.codeSize
         val instructionCount = method.instructionSet.instructions.size
-        writeLine("-- $instructionCount instruction${if (instructionCount > 1) "s" else ""}")
+        writeLine("-- $instructionCount instruction${if (instructionCount > 1) "s" else ""} ${if (codeSize >= 0) "($codeSize bytes)" else ""}")
 
         val (pre, post) = countBranches(method.instructionSet)
         val branches = pre + post
         if (branches > 0) {
-            sb.append("  ".repeat(codeStyle.indent))
+            sb.append(indent)
             writeLine("-- $branches branch${if (branches > 1) "es" else ""} ($pre + $post)")
         }
     }
diff --git a/src/jvmMain/kotlin/dev/romainguy/kotlin/explorer/code/DataModels.kt b/src/jvmMain/kotlin/dev/romainguy/kotlin/explorer/code/DataModels.kt
index ca61e583..069bacff 100644
--- a/src/jvmMain/kotlin/dev/romainguy/kotlin/explorer/code/DataModels.kt
+++ b/src/jvmMain/kotlin/dev/romainguy/kotlin/explorer/code/DataModels.kt
@@ -81,7 +81,7 @@ enum class ISA(val branchInstructions: ScatterSet<String>, val returnInstruction
 
 data class Class(val header: String, val methods: List<Method>, val builtIn: Boolean)
 
-data class Method(val header: String, val instructionSet: InstructionSet, val index: Int = -1)
+data class Method(val header: String, val instructionSet: InstructionSet, val index: Int = -1, val codeSize: Int = -1)
 
 data class InstructionSet(
     val isa: ISA,
diff --git a/src/jvmMain/kotlin/dev/romainguy/kotlin/explorer/oat/OatDumpParser.kt b/src/jvmMain/kotlin/dev/romainguy/kotlin/explorer/oat/OatDumpParser.kt
index 8d935e02..85052c97 100644
--- a/src/jvmMain/kotlin/dev/romainguy/kotlin/explorer/oat/OatDumpParser.kt
+++ b/src/jvmMain/kotlin/dev/romainguy/kotlin/explorer/oat/OatDumpParser.kt
@@ -23,6 +23,8 @@ import dev.romainguy.kotlin.explorer.code.*
 import dev.romainguy.kotlin.explorer.code.CodeContent.Error
 import dev.romainguy.kotlin.explorer.code.CodeContent.Success
 
+private val CodeStartRegex = Regex("^\\s+CODE: \\(code_offset=0x[a-fA-F0-9]+ size=(?<codeSize>\\d+)\\)[.]{3}")
+
 private val ClassNameRegex = Regex("^\\d+: L(?<class>[^;]+); \\(offset=0x$HexDigit+\\) \\(type_idx=\\d+\\).+")
 private val MethodRegex = Regex("^\\s+\\d+:\\s+(?<method>.+)\\s+\\(dex_method_idx=(?<methodIndex>\\d+)\\)")
 private val CodeRegex = Regex("^\\s+0x(?<address>$HexDigit+):\\s+$HexDigit+\\s+(?<code>.+)")
@@ -119,13 +121,15 @@ internal class OatDumpParser {
         consumeUntil("DEX CODE:")
         val methodReferences = readMethodReferences()
 
-        consumeUntil("CODE:")
+        val codeStart = consumeUntil(CodeStartRegex)
+        val codeSize = codeStart?.getValue("codeSize")?.toInt() ?: -1
+
         val instructions = readNativeInstructions(jumpRegex, methodCallRegex, builtIn)
 
         val method = match.getValue("method")
         val index = match.getValue("methodIndex").toInt()
 
-        return Method(method, InstructionSet(isa, instructions, methodReferences), index)
+        return Method(method, InstructionSet(isa, instructions, methodReferences), index, codeSize)
     }
 
     private fun PeekingIterator<String>.readMethodReferences(): IntObjectMap<MethodReference> {