From 71027c48cbeb4df9b5dea148c27598303f6642b7 Mon Sep 17 00:00:00 2001
From: Dominik Peters <github@d-peters.com>
Date: Sun, 1 Jun 2025 16:15:57 +0200
Subject: [PATCH] feat: allow score script to generate treemap file for decomp
 progress (#598)

---
 requirements.txt      |   3 +-
 tools/python/score.py | 116 ++++++++++++++++++++++++++++++++++++++++--
 2 files changed, 115 insertions(+), 4 deletions(-)

diff --git a/requirements.txt b/requirements.txt
index 7f04f0ea..6e8916ef 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -7,4 +7,5 @@ Levenshtein>=0.26.1
 cxxfilt>=0.3.0
 mapfile-parser>=2.7.5
 PyYAML>=6.0.2
-colour>=0.1.5
\ No newline at end of file
+colour>=0.1.5
+plotly>=6.1.0
\ No newline at end of file
diff --git a/tools/python/score.py b/tools/python/score.py
index 25408e40..1933e016 100644
--- a/tools/python/score.py
+++ b/tools/python/score.py
@@ -1,10 +1,10 @@
 import re
 import sys
 import argparse
-import time
 import os
 from file_util import FileUtil
 from score_display import ScoreDisplay
+import plotly.graph_objects as go
 
 ASM_FOLDERS = [
     './asm',
@@ -62,7 +62,7 @@ FUNCTION_REGEX = r'^(?<!static\s)(?:(\/[*][*!][*]*\n(?:[^\/]*\n)+?\s*[*]\/\n)(?:
 GLOBAL_ASM_REGEX = r'\#pragma\sGLOBAL_ASM[(]".*(?=\/)\/([^.]+).s"[)]'
 WIP_REGEX = r'ifdef\s+(?:NON_MATCHING|NON_EQUIVALENT)(?:.|\n)*?\#else\s*(\#pragma\sGLOBAL_ASM[(][^)]*[)])(.|\n)*?'
 NON_MATCHING_REGEX = re.compile(r'^#ifdef[ ]+NON_MATCHING(?:.|\n)*?(?:\s*UNUSED\s+)?(?:[^\s]+)\s(?:\s|[*])*?([0-9A-Za-z_]+)\s*[(][^)]*[)]\s*{', re.MULTILINE)
-NON_EQUVIALENT_REGEX = re.compile(r'^#ifdef[ ]+NON_EQUIVALENT', re.MULTILINE)
+NON_EQUVIALENT_REGEX = re.compile(r'^#ifdef[ ]+NON_EQUIVALENT(?:.|\n)*?(?:\s*UNUSED\s+)?(?:[^\s]+)\s(?:\s|[*])*?([0-9A-Za-z_]+)\s*[(][^)]*[)]\s*{', re.MULTILINE)
 
 CODE_START = 0x80000400
 CODE_END = 0x800D75F4
@@ -142,12 +142,19 @@ class ScoreFile:
     def read_file(self):
         with open(self.path, "r") as inFile:
             self.text = inFile.read()
+            
             self.nonMatchings = re.findall(NON_MATCHING_REGEX, self.text)
             self.nonMatchingsSizes = 0
             for nonMatching in self.nonMatchings:
                 self.nonMatchingsSizes += MAP_FILE.functionSizes[nonMatching]
             self.numNonMatchings = len(self.nonMatchings)
-            self.numNonEquivalents = len(re.findall(NON_EQUVIALENT_REGEX, self.text))
+
+            self.nonEquivalents = re.findall(NON_EQUVIALENT_REGEX, self.text)
+            self.nonEquivalentsSizes = 0
+            for nonEquivalent in self.nonEquivalents:
+                self.nonEquivalentsSizes += MAP_FILE.functionSizes[nonEquivalent]
+            self.numNonEquivalents = len(self.nonEquivalents)
+            
             self.text = re.sub(WIP_REGEX, r"GLOBAL_ASM(\1)", self.text)
             
     def get_matches(self):
@@ -198,6 +205,7 @@ def main():
     parser.add_argument("-t", "--top", help="(Optional) Shows the top N files remaining.")
     parser.add_argument("-a", "--adventure", help="(Optional) Only shows adventure 1 or 2 based on passed in value.", choices=['1', '2'])
     parser.add_argument("-s", "--summary", help="(Optional) Only prints the percentages for adventure 1 and 2", action='store_true')
+    parser.add_argument("--treemap", help="(Optional) Generates a treemap .html file", default="progress-treemap.html", metavar="path/to/treemap-file.html")
     args = parser.parse_args()
     adventureSelect = 3 # Show both adventures by default
     if args.adventure != None:
@@ -267,6 +275,108 @@ def main():
         print(f"Decomp progress: {adventureOnePercentage:5.2f}%")
         print(f"Documentation progress: {adventureTwoPercentage:5.2f}%")
         sys.exit(0)
+
+    if args.treemap:
+        print(f"Generating progress treemap, outputting file to {args.treemap}")
+
+        # Prepare data
+        labels = ["Decomp"]
+        parents = [None]
+        values = [0]
+        colours = [None]
+        legends = [None]
+
+        for scoreFile in scoreFiles:
+            parentName = scoreFile.path.replace('./src/', '')
+            
+            for func in scoreFile.functions:
+                color = '#92ac68'
+                legend = 'matched'
+                if func.isDocumented:
+                    color = 'green'
+                    legend = 'documented'
+                
+                labels.append(func.functionName)
+                parents.append(parentName)
+                values.append(func.size)
+                colours.append(color)
+                legends.append(legend)
+            
+            matches = re.finditer(GLOBAL_ASM_REGEX, scoreFile.text, re.MULTILINE)
+            for match in matches:
+                funcName = match.groups()[0]
+                size = MAP_FILE.functionSizes[funcName]
+                
+                color = "grey"
+                legend = "N/A"
+                if funcName in scoreFile.nonMatchings:
+                    color = 'orange'
+                    legend = "non matching"
+                elif funcName in scoreFile.nonEquivalents:
+                    color = 'red'
+                    legend = "non equivalent"
+                
+                labels.append(funcName)
+                parents.append(parentName)
+                values.append(size)
+                colours.append(color)
+                legends.append(legend)
+
+            labels.append(parentName)
+            parents.append("Decomp")
+            values.append(0)
+            colours.append(None)
+            legends.append(None)
+
+        matchPercents = []
+        for value in values:
+            if value > 0:
+                matchPercents.append(f"Size: {value} bytes<br>Total: {value / CODE_SIZE * 100:.2f}%")
+            else:
+                matchPercents.append(None)
+
+        fig = go.Figure(go.Treemap(
+            labels = labels,
+            parents = parents,
+            values = values,
+            text=matchPercents,
+            hovertemplate='<b>%{label}</b><br>Size: %{value} bytes<extra></extra>',
+            marker = dict(colors = colours),
+            root_color="lightgrey"
+        ))
+
+        # Add custom legend using annotations
+        legend_items = [
+            ("green", "documented"),
+            ("#92ac68", "matched"),
+            ("orange", "non matching"),
+            ("red", "non equivalent"), 
+            ("grey", "N/A")
+        ]
+
+        annotations = []
+        for i, (color, label) in enumerate(legend_items):
+            annotations.append(dict(
+                x=1.05,
+                y=1 - (i * 0.05),
+                xref="paper",
+                yref="paper",
+                showarrow=False,
+                text=f"<span style='color:{color}'>■</span> {label}",
+                font=dict(size=12)
+            ))
+
+
+        fig.update_layout(
+            margin=dict(t=50, l=25, r=150, b=25),  # Extra right margin for legend
+            title="Decomp Progress",
+            annotations=annotations
+        )
+        output_path = "treemap.html"
+        # This will raise an error if writing fails
+        fig.write_html(output_path)
+        sys.exit(0)
+
     scoreDisplay = ScoreDisplay()
     print(scoreDisplay.getDisplay(adventureOnePercentage, adventureOnePercentageWithNonMatching, adventureTwoPercentage, adventureSelect, totalNumberOfDecompiledFunctions, totalNumberOfGlobalAsms, totalNumberOfNonMatching, totalNumberOfNonEquivalent, totalNumberOfDocumentedFunctions, (totalNumberOfFunctions - ignoreNumberDocumentedFunctions) - totalNumberOfDocumentedFunctions))