Explicit rooting of references for Python

python_explicit_rooting_mock.c

#include <stdbool.h>
#include <stddef.h>
#include <stdint.h>
#include <limits.h>
#include <assert.h>

typedef size_t Py_ssize_t;
#define PY_SSIZE_T_MAX ULONG_MAX

/* dummy PyObject structure */
typedef struct {
    int data;
} PyObject;

/* dummy thread state */
typedef struct {
    struct _PyRoot_Stack_Cell_Struct *stack_roots;
} PyThreadState;


/* used for a object return value that must go into a rooted location. */
typedef PyObject *PyRoot;

/* used for a object parameter that is already rooted, maybe on a
 * higher stack frame, as a global or as an object property */
typedef PyObject *PyHandle;

/* stack allocated structure to hold local roots.  Linked into thread state
 * in LIFO fashion */
typedef struct _PyRoot_Stack_Cell_Struct {
    struct _PyRoot_Stack_Cell_Struct *prev;
    PyRoot obj;
} _PyRoot_Stack_Cell;

#define _PyRoot_Decl(name) \
    _PyRoot_Stack_Cell _PyRoot_stack_var_ ## name = {0};


static void
_PyRoot_Link(PyThreadState *tstate, _PyRoot_Stack_Cell *c)
{
    /* push this cell on top of local variable roots stack */
    c->prev = tstate->stack_roots;
    tstate->stack_roots = c;
}

static PyRoot
_PyRoot_Unlink(PyThreadState *tstate, _PyRoot_Stack_Cell *c)
{
    /* require LIFO, this cell must be top of stack */
    assert(tstate->stack_roots == c);
    tstate->stack_roots = c->prev;
    PyRoot rv = (PyRoot)c->obj;
    c->obj = NULL;
    return rv;
}

/* Declare and init a new C stack root 'name' */
#define PyRoot_New(name) _PyRoot_Decl(name); _PyRoot_Link(tstate, &(_PyRoot_stack_var_ ## name))

/* remove a C stack root, must be LIFO with PyRoot_New() */
#define PyRoot_Del(name) (_PyRoot_Unlink(tstate, &(_PyRoot_stack_var_ ## name)))

/* need this to hook in pre/post methods for write barriers */
#define PyRoot_Set(name, val) ((_PyRoot_stack_var_ ## name).obj = val)


/* Convert a PyHandle to a PyRoot.  This is used for return values, to enforce
 that the caller has a properly "rooted" location to store the value.
 */
PyRoot
PyHandle_AsRoot(PyHandle h)
{
    return (PyRoot)h;
}

/* Convert a PyRoot value to a PyHandle, typically to pass as argument */
#define PyRoot_AsHandle(name) ((PyHandle)((_PyRoot_stack_var_ ## name).obj))

static inline bool
PyHandle_IsNull(PyHandle h)
{
    return h == NULL;
}

static inline bool
_PyRoot_IsNull(_PyRoot_Stack_Cell *c)
{
    return c->obj == NULL;
}

#define PyRoot_IsNull(name) _PyRoot_IsNull(&(_PyRoot_stack_var_ ## name))

PyThreadState *tstate;


/* Dummy functions follow to flesh out example code */
static int
PyTuple_CheckExact(PyHandle v)
{
    return (uintptr_t)v == 1;
}

static int
PyList_CheckExact(PyHandle v)
{
    return (uintptr_t)v == 2;
}

static PyRoot
PyList_AsTuple(PyHandle v)
{
    return PyHandle_AsRoot(v);
}

static PyRoot
PyObject_GetIter(PyHandle v)
{
    return PyHandle_AsRoot(v);
}

static Py_ssize_t
PyObject_LengthHint(PyHandle v, Py_ssize_t hint)
{
    return hint;
}

static PyRoot
PyTuple_New(Py_ssize_t n)
{
    return (PyRoot)(n);
}

static PyRoot
PyIter_Next(PyHandle v)
{
    return (PyRoot)(v);
}

static bool
PyErr_Occurred(void)
{
    return tstate == NULL;
}

static PyRoot
PyErr_NoMemory(void)
{
    return NULL;
}

static void
PyTuple_SET_ITEM(PyHandle obj, Py_ssize_t i, PyHandle val)
{
}

PyRoot
PySequence_Tuple(PyHandle v)
{
    if (PyHandle_IsNull(v)) {
        return NULL;
    }

    /* Special-case the common tuple and list cases, for efficiency. */
    if (PyTuple_CheckExact(v)) {
        /* Note that we can't know whether it's safe to return
           a tuple *subclass* instance as-is, hence the restriction
           to exact tuples here.  In contrast, lists always make
           a copy, so there's no need for exactness below. */
        return PyHandle_AsRoot(v);
    }
    if (PyList_CheckExact(v)) {
        return PyList_AsTuple(v);
    }

    /* create new stack roots */
    PyRoot_New(result);
    PyRoot_New(item); /* hoisted from loop for performance */
    PyRoot_New(it);  /* iter(v) */

    /* Get iterator. */
    PyRoot_Set(it, PyObject_GetIter(v));
    if (PyRoot_IsNull(it))
        goto Fail;

    /* Guess result size and allocate space. */
    Py_ssize_t n = PyObject_LengthHint(v, 10);
    if (n == -1)
        goto Fail;
    PyRoot_Set(result, PyTuple_New(n));
    if (PyRoot_IsNull(result))
        goto Fail;

    /* Fill the tuple. */
    for (Py_ssize_t j = 0; ; ++j) {
        PyRoot_Set(item, PyIter_Next(PyRoot_AsHandle(it)));
        if (PyRoot_IsNull(item)) {
            if (PyErr_Occurred())
                goto Fail;
            break;
        }
        if (j >= n) {
            size_t newn = (size_t)n;
            /* The over-allocation strategy can grow a bit faster
               than for lists because unlike lists the
               over-allocation isn't permanent -- we reclaim
               the excess before the end of this routine.
               So, grow by ten and then add 25%.
            */
            newn += 10u;
            newn += newn >> 2;
            if (newn > PY_SSIZE_T_MAX) {
                /* Check for overflow */
                PyErr_NoMemory();
                goto Fail;
            }
            n = (Py_ssize_t)newn;
        }
        PyTuple_SET_ITEM(PyRoot_AsHandle(result), j, PyRoot_AsHandle(item));
    }

    PyRoot_Del(it);
    PyRoot_Del(item);
    return PyRoot_Del(result);

Fail:
    PyRoot_Del(it);
    PyRoot_Del(item);
    PyRoot_Del(result);
    return NULL;
}