From f323cf4847a55740e11a0c7a3f54abca26801440 Mon Sep 17 00:00:00 2001
From: Michael Droettboom <mdboom@gmail.com>
Date: Wed, 10 Jul 2024 09:45:03 -0400
Subject: [PATCH] Fix hook check

---
 pyperf/_hooks.py | 74 +++++++++++++++++++++++++++++++++++++++++++++++-
 pyproject.toml   |  1 +
 2 files changed, 74 insertions(+), 1 deletion(-)

diff --git a/pyperf/_hooks.py b/pyperf/_hooks.py
index 36dd19e6..487355a8 100644
--- a/pyperf/_hooks.py
+++ b/pyperf/_hooks.py
@@ -4,8 +4,12 @@
 
 
 import abc
+from collections import defaultdict
 import importlib.metadata
+import math
+import os
 import sys
+from sys import monitoring
 
 
 def get_hooks():
@@ -80,7 +84,7 @@ def __exit__(self, _exc_type, _exc_value, _traceback):
 
 class pystats(HookBase):
     def __init__(self):
-        if not hasattr(sys, "_pystats_on"):
+        if not hasattr(sys, "_stats_on"):
             raise HookError(
                 "Can not collect pystats because python was not built with --enable-pystats"
             )
@@ -92,6 +96,74 @@ def teardown(self, metadata):
 
     def __enter__(self):
         sys._stats_on()
+        # os.environ["PYTHONSTATS"] = "1"
 
     def __exit__(self, _exc_type, _exc_value, _traceback):
+        # del os.environ["PYTHONSTATS"]
         sys._stats_off()
+
+
+class instruction_count:
+    def __init__(self):
+        self.unique_instructions = 0
+        self.total_instructions = 0
+        monitoring.register_callback(
+            1, monitoring.events.INSTRUCTION, self.instruction_handler
+        )
+        monitoring.register_callback(
+            2, monitoring.events.INSTRUCTION, self.instruction_counter
+        )
+        monitoring.use_tool_id(1, "instruction_handler")
+        monitoring.use_tool_id(2, "instruction_counter")
+
+    def teardown(self, metadata):
+        value = self.total_instructions // self.unique_instructions
+        metadata["mean_instr_count"] = value
+        metadata["instr_count"] = self.unique_instructions
+        monitoring.free_tool_id(1)
+        monitoring.free_tool_id(2)
+
+    def instruction_handler(self, code, instruction_offset):
+        self.unique_instructions += 1
+        return sys.monitoring.DISABLE
+
+    def instruction_counter(self, code, instruction_offset):
+        self.total_instructions += 1
+
+    def __enter__(self):
+        monitoring.set_events(1, monitoring.events.INSTRUCTION)
+        monitoring.set_events(2, monitoring.events.INSTRUCTION)
+
+    def __exit__(self, _exc_type, _exc_value, _traceback):
+        monitoring.set_events(1, 0)
+        monitoring.set_events(2, 0)
+
+
+class instruction_count:
+    def __init__(self):
+        self.counts = defaultdict(int)
+        monitoring.register_callback(
+            1, monitoring.events.INSTRUCTION, self.instruction_handler
+        )
+        monitoring.use_tool_id(1, "instruction_handler")
+
+    def teardown(self, metadata):
+        histogram = [0] * 32
+        for value in self.counts.values():
+            value = math.log2(value)
+            histogram[int(value)] += 1
+        metadata["histogram"] = str(histogram)
+
+        value = sum(self.counts.values()) / len(self.counts)
+        metadata["mean_instr_count"] = value
+        metadata["instr_count"] = len(self.counts)
+        monitoring.free_tool_id(1)
+
+    def instruction_handler(self, code, instruction_offset):
+        self.counts[(code, instruction_offset)] += 1
+
+    def __enter__(self):
+        monitoring.set_events(1, monitoring.events.INSTRUCTION)
+
+    def __exit__(self, _exc_type, _exc_value, _traceback):
+        monitoring.set_events(1, 0)
\ No newline at end of file
diff --git a/pyproject.toml b/pyproject.toml
index 892b4260..0d0bb3ea 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -58,6 +58,7 @@ pyperf = "pyperf.__main__:main"
 
 [project.entry-points."pyperf.hook"]
 pystats = "pyperf._hooks:pystats"
+instruction_count = "pyperf._hooks:instruction_count"
 _test_hook = "pyperf._hooks:_test_hook"
 
 [tool.setuptools]