From 3c137dc613c860f605d3520d7fd722cd8ed79da6 Mon Sep 17 00:00:00 2001
From: Itamar Ostricher <itamarost@gmail.com>
Date: Fri, 2 Dec 2022 09:28:27 -0800
Subject: GH-91054: Add code object watchers API (GH-99859)

* Add API to allow extensions to set callback function on creation and destruction of PyCodeObject

Co-authored-by: Ye11ow-Flash <janshah@cs.stonybrook.edu>
---
 Objects/codeobject.c | 63 ++++++++++++++++++++++++++++++++++++++++++++++++++++
 1 file changed, 63 insertions(+)

(limited to 'Objects/codeobject.c')

diff --git a/Objects/codeobject.c b/Objects/codeobject.c
index f5d90cf65fc..0c197d767b0 100644
--- a/Objects/codeobject.c
+++ b/Objects/codeobject.c
@@ -12,6 +12,66 @@
 #include "clinic/codeobject.c.h"
 
 
+static void
+notify_code_watchers(PyCodeEvent event, PyCodeObject *co)
+{
+    PyInterpreterState *interp = _PyInterpreterState_GET();
+    if (interp->active_code_watchers) {
+        assert(interp->_initialized);
+        for (int i = 0; i < CODE_MAX_WATCHERS; i++) {
+            PyCode_WatchCallback cb = interp->code_watchers[i];
+            if ((cb != NULL) && (cb(event, co) < 0)) {
+                PyErr_WriteUnraisable((PyObject *) co);
+            }
+        }
+    }
+}
+
+int
+PyCode_AddWatcher(PyCode_WatchCallback callback)
+{
+    PyInterpreterState *interp = _PyInterpreterState_GET();
+    assert(interp->_initialized);
+
+    for (int i = 0; i < CODE_MAX_WATCHERS; i++) {
+        if (!interp->code_watchers[i]) {
+            interp->code_watchers[i] = callback;
+            interp->active_code_watchers |= (1 << i);
+            return i;
+        }
+    }
+
+    PyErr_SetString(PyExc_RuntimeError, "no more code watcher IDs available");
+    return -1;
+}
+
+static inline int
+validate_watcher_id(PyInterpreterState *interp, int watcher_id)
+{
+    if (watcher_id < 0 || watcher_id >= CODE_MAX_WATCHERS) {
+        PyErr_Format(PyExc_ValueError, "Invalid code watcher ID %d", watcher_id);
+        return -1;
+    }
+    if (!interp->code_watchers[watcher_id]) {
+        PyErr_Format(PyExc_ValueError, "No code watcher set for ID %d", watcher_id);
+        return -1;
+    }
+    return 0;
+}
+
+int
+PyCode_ClearWatcher(int watcher_id)
+{
+    PyInterpreterState *interp = _PyInterpreterState_GET();
+    assert(interp->_initialized);
+    if (validate_watcher_id(interp, watcher_id) < 0) {
+        return -1;
+    }
+    interp->code_watchers[watcher_id] = NULL;
+    interp->active_code_watchers &= ~(1 << watcher_id);
+    return 0;
+}
+
 /******************
  * generic helpers
  ******************/
@@ -355,6 +415,7 @@ init_code(PyCodeObject *co, struct _PyCodeConstructor *con)
     }
     co->_co_firsttraceable = entry_point;
     _PyCode_Quicken(co);
+    notify_code_watchers(PY_CODE_EVENT_CREATE, co);
 }
 
 static int
@@ -1615,6 +1676,8 @@ code_new_impl(PyTypeObject *type, int argcount, int posonlyargcount,
 static void
 code_dealloc(PyCodeObject *co)
 {
+    notify_code_watchers(PY_CODE_EVENT_DESTROY, co);
+
     if (co->co_extra != NULL) {
         PyInterpreterState *interp = _PyInterpreterState_GET();
         _PyCodeObjectExtra *co_extra = co->co_extra;
-- 
cgit v1.2.3