#include "Python.h"
#include <stdlib.h>

static int installed = 0;

struct {
    PyMemAllocator raw;
    PyMemAllocator mem;
    PyMemAllocator obj;
} hook;

static void
allocation_failed(void)
{
    _exit(1);
}

static void*
hook_malloc(void *ctx, size_t size)
{
    PyMemAllocator *alloc = (PyMemAllocator *)ctx;
    void *ptr;
    ptr = alloc->malloc(alloc->ctx, size);
    if (ptr == NULL)
        allocation_failed();
    return ptr;
}

static void*
hook_realloc(void *ctx, void *ptr, size_t new_size)
{
    PyMemAllocator *alloc = (PyMemAllocator *)ctx;
    void *ptr2;
    ptr2 = alloc->realloc(alloc->ctx, ptr, new_size);
    if (ptr2 == NULL)
        allocation_failed();
    return ptr2;
}

static void
hook_free(void *ctx, void *ptr)
{
    PyMemAllocator *alloc = (PyMemAllocator *)ctx;
    alloc->free(alloc->ctx, ptr);
}

static void
setup_hooks(void)
{
    PyMemAllocator alloc;

    alloc.malloc = hook_malloc;
    alloc.realloc = hook_realloc;
    alloc.free = hook_free;
    PyMem_GetAllocator(PYMEM_DOMAIN_RAW, &hook.raw);
    PyMem_GetAllocator(PYMEM_DOMAIN_MEM, &hook.mem);
    PyMem_GetAllocator(PYMEM_DOMAIN_OBJ, &hook.obj);

    alloc.ctx = &hook.raw;
    PyMem_SetAllocator(PYMEM_DOMAIN_RAW, &alloc);

    alloc.ctx = &hook.mem;
    PyMem_SetAllocator(PYMEM_DOMAIN_MEM, &alloc);

    alloc.ctx = &hook.obj;
    PyMem_SetAllocator(PYMEM_DOMAIN_OBJ, &alloc);
}

static void
remove_hooks(void)
{
    PyMem_SetAllocator(PYMEM_DOMAIN_RAW, &hook.raw);
    PyMem_SetAllocator(PYMEM_DOMAIN_MEM, &hook.mem);
    PyMem_SetAllocator(PYMEM_DOMAIN_OBJ, &hook.obj);
}

static PyObject*
fatalmalloc_enable(PyObject *module)
{
    if (!installed) {
        installed = 1;
        setup_hooks();
    }
    Py_RETURN_NONE;
}

static PyObject*
fatalmalloc_disable(PyObject *module)
{
    if (installed) {
        installed = 0;
        remove_hooks();
    }
    Py_RETURN_NONE;
}

PyDoc_STRVAR(module_doc,
"fatalmalloc module.");

static PyMethodDef module_methods[] = {
    {"enable",
     (PyCFunction)fatalmalloc_enable, METH_NOARGS,
     PyDoc_STR("enable()")},
    {"disable",
     (PyCFunction)fatalmalloc_disable, METH_NOARGS,
     PyDoc_STR("disable()")},
    {NULL, NULL}  /* sentinel */
};

static struct PyModuleDef module_def = {
    PyModuleDef_HEAD_INIT,
    "fatalmalloc",
    module_doc,
    0,
    module_methods,
    NULL,
    NULL,
    NULL,
    NULL
};

PyMODINIT_FUNC
PyInit_fatalmalloc(void)
{
    return PyModule_Create(&module_def);
}