Add hotspot analysis

3 files changed, 54 insertions, 6 deletions

diff --git a/atk16_emu/cli.py b/atk16_emu/cli.py
index 9cb268a..613af6c 100644
--- a/atk16_emu/cli.py
+++ b/atk16_emu/cli.py
@@ -57,6 +57,7 @@ def main():
     print("===============")
 
     machine.print_state_summary()
+    machine.print_hotspot_summary()
 
   else:
     debugger = Debugger(options.peripherals_enabled)
diff --git a/atk16_emu/emu.py b/atk16_emu/emu.py
index b3a3500..cc211ec 100644
--- a/atk16_emu/emu.py
+++ b/atk16_emu/emu.py
@@ -3,6 +3,7 @@ os.environ['PYGAME_HIDE_SUPPORT_PROMPT'] = "hide"
 
 import sys
 from dataclasses import dataclass
+import heapq
 
 from .opcodes import *
 from .memory import *
@@ -103,6 +104,12 @@ class ALU:
 
     raise ValueError(f"Invalid ALU S: {S}")
 
+@dataclass
+class HotspotSpan:
+  from_addr: int
+  to_addr: int
+  count: int
+
 class Machine:
   def __init__(self, peripherals_enabled: bool = False):
     self.rom = ROM(15, 16)
@@ -136,10 +143,22 @@ class Machine:
 
     self.peripherals_enabled = peripherals_enabled
     if peripherals_enabled:
-      self.peripherals = Peripherals(set_irq_line=self.set_irq_line)
+      self.peripherals = Peripherals(set_irq_line=self.set_irq_line,
+                                     request_quit=self.shutdown)
     else:
       self.peripherals = DummyPeripherals()
 
+    self.hotspot_table: dict[int, int] = {}
+
+  def shutdown(self):
+    self.running = False
+
+  def record_hotspot(self, addr: int):
+    if addr not in self.hotspot_table:
+      self.hotspot_table[addr] = 0
+
+    self.hotspot_table[addr] += 1
+
   def make_copy(self):
     new_machine = Machine()
 
@@ -304,6 +323,7 @@ class Machine:
       self.is_critical_section = True
 
     pc_addr = self.pc.value
+    self.record_hotspot(pc_addr)
     self.pc.step()
 
     instr = self.mem_read(pc_addr)
@@ -467,4 +487,30 @@ class Machine:
           f"Z: {as_num(self.fr.zero)}, "
           f"S: {as_num(self.fr.sign)}")
 
-    print()
\ No newline at end of file
+    print()
+
+  def print_hotspot_summary(self):
+    # show the top 10 most executed instructions using self.hotspot_table
+    N = 100
+    counts = heapq.nlargest(N, self.hotspot_table.items(), key=lambda item: item[1])
+
+    print("Hotspot analysis summary:")
+    current_span: HotspotSpan | None = None
+    for addr, count in counts:
+      if current_span is None:
+        current_span = HotspotSpan(from_addr=addr,
+                                   to_addr=addr,
+                                   count=count)
+      elif addr == current_span.to_addr + 1:
+        current_span.to_addr = addr
+        current_span.count += count
+      else:
+        print(C.OKGREEN + f"0x{current_span.from_addr:>04x}..0x{current_span.to_addr:>04x}" + C.ENDC + f"   {count}")
+        current_span = HotspotSpan(from_addr=addr,
+                                   to_addr=addr,
+                                   count=count)
+
+    if current_span is not None:
+      print(C.OKGREEN + f"0x{current_span.from_addr:>04x}..0x{current_span.to_addr:>04x}" + C.ENDC + f"   {count}")
+
+    print()
diff --git a/atk16_emu/peripherals.py b/atk16_emu/peripherals.py
index 3bc0443..9b65410 100644
--- a/atk16_emu/peripherals.py
+++ b/atk16_emu/peripherals.py
@@ -1,5 +1,5 @@
 import pygame
-import sys
+from sys import stdout
 import os
 import random
 from .memory import RAM, ROM
@@ -8,18 +8,19 @@ from typing import Callable
 IRQ_LINE_KEYBOARD = 0
 
 class Peripherals:
-  def __init__(self, set_irq_line: Callable[[int], None]):
+  def __init__(self, set_irq_line: Callable[[int], None], request_quit: Callable[[], None]):
     pygame.init()
     self.graphics = Graphics()
     self.keyboard = Keyboard(set_irq_line=set_irq_line)
     self.terminal = Terminal()
 
+    self.request_quit = request_quit
     self.set_irq_line = set_irq_line
 
   def step(self):
     for event in pygame.event.get():
       if event.type == pygame.QUIT:
-          sys.exit(0)
+        self.request_quit()
       elif event.type == pygame.KEYDOWN:
         self.keyboard.set_key_down(event.key)
       elif event.type == pygame.KEYUP:
@@ -242,4 +243,4 @@ class Terminal:
 
   def write(self, s: int):
     print(chr(s), end="")
-    sys.stdout.flush()
+    stdout.flush()