/* PickleBuffer object implementation */

#define PY_SSIZE_T_CLEAN
#include "Python.h"
#include <stddef.h>

typedef struct {
 PyObject_HEAD
 /* The view exported by the original object */
 Py_buffer view;
 PyObject *weakreflist;
} PyPickleBufferObject;

/* C API */

PyObject *
PyPickleBuffer_FromObject(PyObject *base)
{
 PyTypeObject *type = &PyPickleBuffer_Type;
 PyPickleBufferObject *self;

 self = (PyPickleBufferObject *) type->tp_alloc(type, 0);
 if (self == NULL) {
 return NULL;
 }
 self->view.obj = NULL;
 self->weakreflist = NULL;
 if (PyObject_GetBuffer(base, &self->view, PyBUF_FULL_RO) < 0) {
 Py_DECREF(self);
 return NULL;
 }
 return (PyObject *) self;
}

const Py_buffer *
PyPickleBuffer_GetBuffer(PyObject *obj)
{
 PyPickleBufferObject *self = (PyPickleBufferObject *) obj;

 if (!PyPickleBuffer_Check(obj)) {
 PyErr_Format(PyExc_TypeError,
 "expected PickleBuffer, %.200s found",
 Py_TYPE(obj)->tp_name);
 return NULL;
 }
 if (self->view.obj == NULL) {
 PyErr_SetString(PyExc_ValueError,
 "operation forbidden on released PickleBuffer object");
 return NULL;
 }
 return &self->view;
}

int
PyPickleBuffer_Release(PyObject *obj)
{
 PyPickleBufferObject *self = (PyPickleBufferObject *) obj;

 if (!PyPickleBuffer_Check(obj)) {
 PyErr_Format(PyExc_TypeError,
 "expected PickleBuffer, %.200s found",
 Py_TYPE(obj)->tp_name);
 return -1;
 }
 PyBuffer_Release(&self->view);
 return 0;
}

static PyObject *
picklebuf_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
{
 PyPickleBufferObject *self;
 PyObject *base;
 char *keywords[] = {"", NULL};

 if (!PyArg_ParseTupleAndKeywords(args, kwds, "O:PickleBuffer",
 keywords, &base)) {
 return NULL;
 }

 self = (PyPickleBufferObject *) type->tp_alloc(type, 0);
 if (self == NULL) {
 return NULL;
 }
 self->view.obj = NULL;
 self->weakreflist = NULL;
 if (PyObject_GetBuffer(base, &self->view, PyBUF_FULL_RO) < 0) {
 Py_DECREF(self);
 return NULL;
 }
 return (PyObject *) self;
}

static int
picklebuf_traverse(PyPickleBufferObject *self, visitproc visit, void *arg)
{
 Py_VISIT(self->view.obj);
 return 0;
}

static int
picklebuf_clear(PyPickleBufferObject *self)
{
 PyBuffer_Release(&self->view);
 return 0;
}

static void
picklebuf_dealloc(PyPickleBufferObject *self)
{
 PyObject_GC_UnTrack(self);
 if (self->weakreflist != NULL)
 PyObject_ClearWeakRefs((PyObject *) self);
 PyBuffer_Release(&self->view);
 Py_TYPE(self)->tp_free((PyObject *) self);
}

/* Buffer API */

static int
picklebuf_getbuf(PyPickleBufferObject *self, Py_buffer *view, int flags)
{
 if (self->view.obj == NULL) {
 PyErr_SetString(PyExc_ValueError,
 "operation forbidden on released PickleBuffer object");
 return -1;
 }
 return PyObject_GetBuffer(self->view.obj, view, flags);
}

static void
picklebuf_releasebuf(PyPickleBufferObject *self, Py_buffer *view)
{
 /* Since our bf_getbuffer redirects to the original object, this
 * implementation is never called. It only exists to signal that
 * buffers exported by PickleBuffer have non-trivial releasing
 * behaviour (see check in Python/getargs.c).
 */
}

static PyBufferProcs picklebuf_as_buffer = {
 .bf_getbuffer = (getbufferproc) picklebuf_getbuf,
 .bf_releasebuffer = (releasebufferproc) picklebuf_releasebuf,
};

/* Methods */

static PyObject *
picklebuf_raw(PyPickleBufferObject *self, PyObject *Py_UNUSED(ignored))
{
 if (self->view.obj == NULL) {
 PyErr_SetString(PyExc_ValueError,
 "operation forbidden on released PickleBuffer object");
 return NULL;
 }
 if (self->view.suboffsets != NULL
 || !PyBuffer_IsContiguous(&self->view, 'A')) {
 PyErr_SetString(PyExc_BufferError,
 "cannot extract raw buffer from non-contiguous buffer");
 return NULL;
 }
 PyObject *m = PyMemoryView_FromObject((PyObject *) self);
 if (m == NULL) {
 return NULL;
 }
 PyMemoryViewObject *mv = (PyMemoryViewObject *) m;
 assert(mv->view.suboffsets == NULL);
 /* Mutate memoryview instance to make it a "raw" memoryview */
 mv->view.format = "B";
 mv->view.ndim = 1;
 mv->view.itemsize = 1;
 /* shape = (length,) */
 mv->view.shape = &mv->view.len;
 /* strides = (1,) */
 mv->view.strides = &mv->view.itemsize;
 /* Fix memoryview state flags */
 /* XXX Expose memoryobject.c's init_flags() instead? */
 mv->flags = _Py_MEMORYVIEW_C | _Py_MEMORYVIEW_FORTRAN;
 return m;
}

PyDoc_STRVAR(picklebuf_raw_doc,
"raw($self, /)\n--\n\
\n\
Return a memoryview of the raw memory underlying this buffer.\n\
Will raise BufferError is the buffer isn't contiguous.");

static PyObject *
picklebuf_release(PyPickleBufferObject *self, PyObject *Py_UNUSED(ignored))
{
 PyBuffer_Release(&self->view);
 Py_RETURN_NONE;
}

PyDoc_STRVAR(picklebuf_release_doc,
"release($self, /)\n--\n\
\n\
Release the underlying buffer exposed by the PickleBuffer object.");

static PyMethodDef picklebuf_methods[] = {
 {"raw", (PyCFunction) picklebuf_raw, METH_NOARGS, picklebuf_raw_doc},
 {"release", (PyCFunction) picklebuf_release, METH_NOARGS, picklebuf_release_doc},
 {NULL, NULL}
};

PyTypeObject PyPickleBuffer_Type = {
 PyVarObject_HEAD_INIT(NULL, 0)
 .tp_name = "pickle.PickleBuffer",
 .tp_doc = "Wrapper for potentially out-of-band buffers",
 .tp_basicsize = sizeof(PyPickleBufferObject),
 .tp_flags = Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC,
 .tp_new = picklebuf_new,
 .tp_dealloc = (destructor) picklebuf_dealloc,
 .tp_traverse = (traverseproc) picklebuf_traverse,
 .tp_clear = (inquiry) picklebuf_clear,
 .tp_weaklistoffset = offsetof(PyPickleBufferObject, weakreflist),
 .tp_as_buffer = &picklebuf_as_buffer,
 .tp_methods = picklebuf_methods,
};