Add ability for external executor

Author: DinoV

Include/cpython/pyframe.h

@@ -40,6 +40,7 @@ PyAPI_FUNC(int) PyUnstable_InterpreterFrame_GetLine(struct _PyInterpreterFrame *
 #define PyUnstable_EXECUTABLE_KIND_PY_FUNCTION 1
 #define PyUnstable_EXECUTABLE_KIND_BUILTIN_FUNCTION 3
 #define PyUnstable_EXECUTABLE_KIND_METHOD_DESCRIPTOR 4
-#define PyUnstable_EXECUTABLE_KINDS 5
+#define PyUnstable_EXECUTABLE_KIND_JIT 5
+#define PyUnstable_EXECUTABLE_KINDS 6
 
 PyAPI_DATA(const PyTypeObject *) const PyUnstable_ExecutableKinds[PyUnstable_EXECUTABLE_KINDS+1];

Include/internal/pycore_genobject.h

@@ -8,7 +8,7 @@ extern "C" {
 #  error "this header requires Py_BUILD_CORE define"
 #endif
 
-#include "pycore_interpframe_structs.h" // _PyGenObject
+#include "pycore_interpframe_structs.h" // _PyInterpreterFrame
 
 #include <stddef.h>               // offsetof()
 

Include/internal/pycore_interp_structs.h

@@ -528,7 +528,7 @@ struct _py_func_state {
 
 /* For now we hard-code this to a value for which we are confident
    all the static builtin types will fit (for all builds). */
-#define _Py_MAX_MANAGED_STATIC_BUILTIN_TYPES 202
+#define _Py_MAX_MANAGED_STATIC_BUILTIN_TYPES 210
 #define _Py_MAX_MANAGED_STATIC_EXT_TYPES 10
 #define _Py_MAX_MANAGED_STATIC_TYPES \
     (_Py_MAX_MANAGED_STATIC_BUILTIN_TYPES + _Py_MAX_MANAGED_STATIC_EXT_TYPES)

Include/internal/pycore_interpframe.h

@@ -17,9 +17,39 @@ extern "C" {
 #define _PyInterpreterFrame_LASTI(IF) \
     ((int)((IF)->instr_ptr - _PyFrame_GetBytecode((IF))))
 
+PyAPI_DATA(PyTypeObject) PyUnstable_ExternalExecutable_Type;
+
+#define PyUnstable_ExternalExecutable_Check(op) Py_IS_TYPE((op), &PyUnstable_ExternalExecutable_Type)
+
+// Initialize a potentially external frame and make it safe to access the
+// all of the members of the returned _PyInterpreterFrame. The returned
+// value will be the same address as the passed in pointer.
+PyAPI_FUNC(void) _PyFrame_InitializeExternalFrame(_PyInterpreterFrame *frame);
+
+PyAPI_FUNC(PyObject *) PyUnstable_MakeExternalExecutable(_PyFrame_Reifier reifier, PyCodeObject *code, PyObject *state);
+
+static inline bool _PyFrame_IsExternalFrame(_PyInterpreterFrame *frame)
+{
+    if (PyStackRef_IsNull(frame->f_executable)) {
+        return false;
+    }
+    return PyUnstable_ExternalExecutable_Check(PyStackRef_AsPyObjectBorrow(frame->f_executable));
+}
+
+static inline void
+_PyFrame_EnsureFrameFullyInitialized(_PyInterpreterFrame *frame)
+{
+    if (_PyFrame_IsExternalFrame(frame)) {
+        _PyFrame_InitializeExternalFrame(frame);
+    }
+}
+
 static inline PyCodeObject *_PyFrame_GetCode(_PyInterpreterFrame *f) {
     assert(!PyStackRef_IsNull(f->f_executable));
     PyObject *executable = PyStackRef_AsPyObjectBorrow(f->f_executable);
+    if (PyUnstable_ExternalExecutable_Check(executable)) {
+        return ((PyUnstable_PyExternalExecutable *)executable)->ef_code;
+    }
     assert(PyCode_Check(executable));
     return (PyCodeObject *)executable;
 }
@@ -30,12 +60,6 @@ static inline PyCodeObject *_PyFrame_GetCode(_PyInterpreterFrame *f) {
 static inline PyCodeObject* _Py_NO_SANITIZE_THREAD
 _PyFrame_SafeGetCode(_PyInterpreterFrame *f)
 {
-    // globals and builtins may be NULL on a legit frame, but it's unlikely.
-    // It's more likely that it's a sign of an invalid frame.
-    if (f->f_globals == NULL || f->f_builtins == NULL) {
-        return NULL;
-    }
-
     if (PyStackRef_IsNull(f->f_executable)) {
         return NULL;
     }
@@ -48,6 +72,18 @@ _PyFrame_SafeGetCode(_PyInterpreterFrame *f)
     if (_PyObject_IsFreed(executable)) {
         return NULL;
     }
+    if (_PyFrame_IsExternalFrame(f)) {
+        executable = (PyObject *)((PyUnstable_PyExternalExecutable *)executable)->ef_code;
+        if (_PyObject_IsFreed(executable)) {
+            return NULL;
+        }
+    } else {
+        // globals and builtins may be NULL on a legit frame, but it's unlikely.
+        // It's more likely that it's a sign of an invalid frame.
+        if (f->f_globals == NULL || f->f_builtins == NULL) {
+            return NULL;
+        }
+    }
     if (!PyCode_Check(executable)) {
         return NULL;
     }
@@ -81,6 +117,7 @@ _PyFrame_SafeGetLasti(struct _PyInterpreterFrame *f)
     }
 
     _Py_CODEUNIT *bytecode;
+    _PyFrame_EnsureFrameFullyInitialized(f);
 #ifdef Py_GIL_DISABLED
     _PyCodeArray *tlbc = _PyCode_GetTLBCArray(co);
     assert(f->tlbc_index >= 0 && f->tlbc_index < tlbc->size);
@@ -262,6 +299,9 @@ _PyFrame_IsIncomplete(_PyInterpreterFrame *frame)
     if (frame->owner >= FRAME_OWNED_BY_INTERPRETER) {
         return true;
     }
+    if (frame->instr_ptr == NULL) {
+        return false;
+    }
     return frame->owner != FRAME_OWNED_BY_GENERATOR &&
            frame->instr_ptr < _PyFrame_GetBytecode(frame) +
                                   _PyFrame_GetCode(frame)->_co_firsttraceable;
@@ -276,12 +316,69 @@ _PyFrame_GetFirstComplete(_PyInterpreterFrame *frame)
     return frame;
 }
 
+#if Py_DEBUG
+
+static inline bool _Py_NO_SANITIZE_THREAD
+_PyFrame_IsIncompleteOrUninitialized(_PyInterpreterFrame *frame)
+{
+    if (frame->owner >= FRAME_OWNED_BY_INTERPRETER || _PyFrame_IsExternalFrame(frame)) {
+        return true;
+    }
+    if (frame->instr_ptr == NULL) {
+        return true;
+    }
+    return frame->owner != FRAME_OWNED_BY_GENERATOR &&
+           frame->instr_ptr < _PyFrame_GetBytecode(frame) +
+                             _PyFrame_GetCode(frame)->_co_firsttraceable;
+}
+
+static inline _PyInterpreterFrame *
+_PyFrame_GetFirstCompleteInitialized(_PyInterpreterFrame *frame)
+{
+    while (frame && _PyFrame_IsIncompleteOrUninitialized(frame)) {
+        frame = frame->previous;
+    }
+    return frame;
+}
+
+#endif
+
+static inline bool
+_PyFrame_StackpointerSaved(void)
+{
+#if Py_DEBUG
+    PyThreadState *tstate = PyThreadState_GET();
+    return _PyFrame_GetFirstCompleteInitialized(tstate->current_frame) == NULL ||
+           _PyFrame_GetFirstCompleteInitialized(tstate->current_frame)->stackpointer != NULL;
+#else
+    return true;
+#endif
+}
+
+
+
 static inline _PyInterpreterFrame *
 _PyThreadState_GetFrame(PyThreadState *tstate)
 {
     return _PyFrame_GetFirstComplete(tstate->current_frame);
 }
 
+static inline PyObject *
+_PyFrame_GetGlobals(_PyInterpreterFrame *frame) {
+    if (frame->f_globals == NULL) {
+        frame->f_globals = _PyFrame_GetFunction(frame)->func_globals;
+    }
+    return frame->f_globals;
+}
+
+static inline PyObject *
+_PyFrame_GetBuiltins(_PyInterpreterFrame *frame) {
+    if (frame->f_builtins == NULL) {
+        frame->f_builtins = _PyFrame_GetFunction(frame)->func_builtins;
+    }
+    return frame->f_builtins;
+}
+
 /* For use by _PyFrame_GetFrameObject
   Do not call directly. */
 PyAPI_FUNC(PyFrameObject *)
@@ -293,9 +390,8 @@ _PyFrame_MakeAndSetFrameObject(_PyInterpreterFrame *frame);
 static inline PyFrameObject *
 _PyFrame_GetFrameObject(_PyInterpreterFrame *frame)
 {
-
     assert(!_PyFrame_IsIncomplete(frame));
-    PyFrameObject *res =  frame->frame_obj;
+    PyFrameObject *res = frame->frame_obj;
     if (res != NULL) {
         return res;
     }
@@ -314,7 +410,7 @@ _PyFrame_ClearLocals(_PyInterpreterFrame *frame);
  * take should  be set to 1 for heap allocated
  * frames like the ones in generators and coroutines.
  */
-void
+PyAPI_FUNC(void)
 _PyFrame_ClearExceptCode(_PyInterpreterFrame * frame);
 
 int
@@ -363,7 +459,7 @@ _PyFrame_PushUnchecked(PyThreadState *tstate, _PyStackRef func, int null_locals_
 /* Pushes a trampoline frame without checking for space.
  * Must be guarded by _PyThreadState_HasStackSpace() */
 static inline _PyInterpreterFrame *
-_PyFrame_PushTrampolineUnchecked(PyThreadState *tstate, PyCodeObject *code, int stackdepth, _PyInterpreterFrame * previous)
+_PyFrame_PushTrampolineUnchecked(PyThreadState *tstate, PyCodeObject *code, int stackdepth, _PyInterpreterFrame *previous)
 {
     CALL_STAT_INC(frames_pushed);
     _PyInterpreterFrame *frame = (_PyInterpreterFrame *)tstate->datastack_top;

Include/internal/pycore_interpframe_structs.h

@@ -85,6 +85,15 @@ struct _PyAsyncGenObject {
     _PyGenObject_HEAD(ag)
 };
 
+typedef void (*_PyFrame_Reifier)(struct _PyInterpreterFrame *, PyObject *reifier);
+
+typedef struct {
+    PyObject_HEAD
+    PyCodeObject *ef_code;
+    PyObject *ef_state;
+    _PyFrame_Reifier ef_reifier;
+} PyUnstable_PyExternalExecutable;
+
 #undef _PyGenObject_HEAD
 
 

Lib/test/test_capi/test_misc.py

@@ -4,6 +4,7 @@
 import _thread
 from collections import deque
 import contextlib
+import dis
 import importlib.machinery
 import importlib.util
 import json
@@ -2869,6 +2870,61 @@ def func():
         names = ["func", "outer", "outer", "inner", "inner", "outer", "inner"]
         self.do_test(func, names)
 
+    def test_jit_frame(self):
+        def fakefunc():
+            pass
+
+        def f():
+            return sys._getframe(1)
+
+        res = _testinternalcapi.call_with_jit_frame(fakefunc, f, ())
+
+    def test_jit_frame_instr_ptr(self):
+        """jit executable can fill in the instr ptr each time the frame is queried"""
+        def fakefunc():
+            pass
+            pass
+            pass
+            pass
+
+        offset = 0
+        linenos = []
+        def test():
+            for op in dis.get_instructions(fakefunc):
+                if op.opname in ("RESUME", "NOP", "RETURN_VALUE"):
+                    nonlocal offset
+                    offset = op.offset//2
+                    linenos.append(sys._getframe(1).f_lineno)
+
+        def callback():
+            return {"instr_ptr": offset}
+
+        _testinternalcapi.call_with_jit_frame(fakefunc, test, (), callback)
+        base = fakefunc.__code__.co_firstlineno
+        self.assertEqual(linenos, [base, base + 1, base + 2, base + 3, base + 4])
+
+    def test_jit_frame_code(self):
+        """internal C api checks the for a code executor"""
+        def fakefunc():
+            pass
+
+        def callback():
+            return _testinternalcapi.iframe_getcode(sys._getframe(1))
+
+        res = _testinternalcapi.call_with_jit_frame(fakefunc, callback, ())
+        self.assertEqual(res, fakefunc.__code__)
+
+    def test_jit_frame_line(self):
+        """internal C api checks the for a code executor"""
+        def fakefunc():
+            pass
+
+        def callback():
+            return _testinternalcapi.iframe_getline(sys._getframe(1))
+
+        res = _testinternalcapi.call_with_jit_frame(fakefunc, callback, ())
+        self.assertEqual(res, fakefunc.__code__.co_firstlineno)
+
 
 class Test_Pep523AllowSpecialization(unittest.TestCase):
     """Tests for _PyInterpreterState_SetEvalFrameFunc with