#include "parts.h"

#include "util.h"

static Py_ssize_t

get_code_extra_index(PyInterpreterState* interp) {

Py_ssize_t result = -1;

static const char *key = "_testcapi.frame_evaluation.code_index";

PyObject *interp_dict = PyInterpreterState_GetDict(interp); // borrowed

assert(interp_dict); // real users would handle missing dict... somehow

PyObject *index_obj;

if (PyDict_GetItemStringRef(interp_dict, key, &index_obj) < 0) {

goto finally;

}

Py_ssize_t index = 0;

if (!index_obj) {

index = PyUnstable_Eval_RequestCodeExtraIndex(NULL);

if (index < 0 || PyErr_Occurred()) {

goto finally;

}

index_obj = PyLong_FromSsize_t(index); // strong ref

if (!index_obj) {

goto finally;

}

int res = PyDict_SetItemString(interp_dict, key, index_obj);

Py_DECREF(index_obj);

if (res < 0) {

goto finally;

}

}

else {

index = PyLong_AsSsize_t(index_obj);

Py_DECREF(index_obj);

if (index == -1 && PyErr_Occurred()) {

goto finally;

}

}

result = index;

finally:

return result;

}

static PyObject *

test_code_extra(PyObject* self, PyObject *Py_UNUSED(callable))

{

PyObject *result = NULL;

PyObject *test_func = NULL;

// Get or initialize interpreter-specific code object storage index

PyInterpreterState *interp = PyInterpreterState_Get();

if (!interp) {

return NULL;

}

Py_ssize_t code_extra_index = get_code_extra_index(interp);

if (PyErr_Occurred()) {

goto finally;

}

// Get a function to test with

// This can be any Python function. Use `test.test_misc.testfunction`.

test_func = PyImport_ImportModuleAttrString("test.test_capi.test_misc",

"testfunction");

if (!test_func) {

goto finally;

}

PyObject *test_func_code = PyFunction_GetCode(test_func); // borrowed

if (!test_func_code) {

goto finally;

}

// Check the value is initially NULL

void *extra = UNINITIALIZED_PTR;

int res = PyUnstable_Code_GetExtra(test_func_code, code_extra_index, &extra);

if (res < 0) {

goto finally;

}

assert (extra == NULL);

// Set another code extra value

res = PyUnstable_Code_SetExtra(test_func_code, code_extra_index, (void*)(uintptr_t)77);

if (res < 0) {

goto finally;

}

// Assert it was set correctly

extra = UNINITIALIZED_PTR;

res = PyUnstable_Code_GetExtra(test_func_code, code_extra_index, &extra);

if (res < 0) {

goto finally;

}

assert ((uintptr_t)extra == 77);

// Revert to initial code extra value.

res = PyUnstable_Code_SetExtra(test_func_code, code_extra_index, NULL);

if (res < 0) {

goto finally;

}

result = Py_NewRef(Py_None);

finally:

Py_XDECREF(test_func);

return result;

}

static PyMethodDef TestMethods[] = {

{"test_code_extra", test_code_extra, METH_NOARGS},

{NULL},

};

_PyTestCapi_Init_Code(PyObject *m) {

if (PyModule_AddFunctions(m, TestMethods) < 0) {

return -1;

}

return 0;

}