Add co_branches iterator

Author: markshannon

Include/internal/pycore_code.h

@@ -590,6 +590,8 @@ extern int _Py_GetBaseOpcode(PyCodeObject *code, int offset);
 
 extern int _PyInstruction_GetLength(PyCodeObject *code, int offset);
 
+extern PyObject *_PyInstrumentation_BranchesIterator(PyCodeObject *code);
+
 #ifdef __cplusplus
 }
 #endif

Objects/codeobject.c

@@ -1277,6 +1277,7 @@ typedef struct {
 } lineiterator;
 
 
+
 static void
 lineiter_dealloc(lineiterator *li)
 {
@@ -2162,6 +2163,12 @@ code_linesiterator(PyCodeObject *code, PyObject *Py_UNUSED(args))
     return (PyObject *)new_linesiterator(code);
 }
 
+static PyObject *
+code_branchesiterator(PyCodeObject *code, PyObject *Py_UNUSED(args))
+{
+    return _PyInstrumentation_BranchesIterator(code);
+}
+
 /*[clinic input]
 @text_signature "($self, /, **changes)"
 code.replace
@@ -2302,6 +2309,7 @@ code__varname_from_oparg_impl(PyCodeObject *self, int oparg)
 static struct PyMethodDef code_methods[] = {
     {"__sizeof__", (PyCFunction)code_sizeof, METH_NOARGS},
     {"co_lines", (PyCFunction)code_linesiterator, METH_NOARGS},
+    {"co_branches", (PyCFunction)code_branchesiterator, METH_NOARGS},
     {"co_positions", (PyCFunction)code_positionsiterator, METH_NOARGS},
     CODE_REPLACE_METHODDEF
     CODE__VARNAME_FROM_OPARG_METHODDEF

Python/instrumentation.c

@@ -2877,3 +2877,87 @@ _PyMonitoring_RegisterCallback(int tool_id, int event_id, PyObject *obj)
     }
     return res;
 }
+
+/* Branch Iterator */
+
+typedef struct {
+    PyObject_HEAD
+    PyCodeObject *bi_code;
+    int bi_offset;
+} branchesiterator;
+
+PyObject *int_triple(int a, int b, int c) {
+    PyObject *obja = PyLong_FromLong(a);
+    PyObject *objb = NULL;
+    PyObject *objc = NULL;
+    if (obja == NULL) {
+        goto error;
+    }
+    objb = PyLong_FromLong(b);
+    if (objb == NULL) {
+        goto error;
+    }
+    objc = PyLong_FromLong(c);
+    if (objc == NULL) {
+        goto error;
+    }
+    PyObject *array[3] = { obja, objb, objc };
+    return _PyTuple_FromArraySteal(array, 3);
+error:
+    Py_XDECREF(obja);
+    Py_XDECREF(objb);
+    Py_XDECREF(objc);
+    return NULL;
+}
+
+static PyObject *
+branchesiter_next(branchesiterator *bi)
+{
+    int offset = bi->bi_offset;
+    while (offset < Py_SIZE(bi->bi_code)) {
+        _Py_CODEUNIT inst = _PyCode_CODE(bi->bi_code)[offset];
+        int next_offset = offset + _PyInstruction_GetLength(bi->bi_code, offset);
+        int event = EVENT_FOR_OPCODE[inst.op.code];
+        if (event == PY_MONITORING_EVENT_BRANCH_TAKEN) {
+            /* Skip NOT_TAKEN */
+            int not_taken = next_offset + 1;
+            bi->bi_offset = not_taken;
+            return int_triple(offset*2, not_taken*2, (next_offset + inst.op.arg)*2);
+        }
+        offset = next_offset;
+    }
+    return NULL;
+}
+
+static void
+branchesiter_dealloc(branchesiterator *bi)
+{
+    Py_DECREF(bi->bi_code);
+    PyObject_Free(bi);
+}
+
+PyTypeObject _PyBranchesIterator = {
+    PyVarObject_HEAD_INIT(&PyType_Type, 0)
+    "line_iterator",                    /* tp_name */
+    sizeof(branchesiterator),               /* tp_basicsize */
+    0,                                  /* tp_itemsize */
+    /* methods */
+    .tp_dealloc = (destructor)branchesiter_dealloc,
+    .tp_flags = Py_TPFLAGS_DEFAULT | Py_TPFLAGS_BASETYPE,
+    .tp_iter = PyObject_SelfIter,
+    .tp_iternext = (iternextfunc)branchesiter_next,
+    .tp_free = PyObject_Del,
+};
+
+PyObject *
+_PyInstrumentation_BranchesIterator(PyCodeObject *code)
+{
+
+    branchesiterator *bi = (branchesiterator *)PyType_GenericAlloc(&_PyBranchesIterator, 0);
+    if (bi == NULL) {
+        return NULL;
+    }
+    bi->bi_code = (PyCodeObject*)Py_NewRef(code);
+    bi->bi_offset = 0;
+    return (PyObject *)bi;
+}