MegEngine/imperative/python/src/pyext17.h

#pragma once

#include <Python.h>
#include <pybind11/pybind11.h>
#include <exception>
#include <stdexcept>
#include <utility>
#include <vector>

namespace pyext17 {

#ifdef METH_FASTCALL
constexpr bool has_fastcall = true;
#else
constexpr bool has_fastcall = false;
#endif

#ifdef _Py_TPFLAGS_HAVE_VECTORCALL
constexpr bool has_vectorcall = true;
#else
constexpr bool has_vectorcall = false;
#endif

template <typename... Args>
struct invocable_with {
    template <typename T>
    constexpr bool operator()(T&& lmb) {
        return std::is_invocable_v<T, Args...>;
    }
};

#define HAS_MEMBER_TYPE(T, U) \
    invocable_with<T>{}([](auto&& x) -> typename std::decay_t<decltype(x)>::U {})
#define HAS_MEMBER(T, m) \
    invocable_with<T>{}([](auto&& x) -> decltype(&std::decay_t<decltype(x)>::m) {})

inline PyObject* cvt_retval(PyObject* rv) {
    return rv;
}

#define CVT_RET_PYOBJ(...)                                       \
    if constexpr (std::is_same_v<decltype(__VA_ARGS__), void>) { \
        __VA_ARGS__;                                             \
        Py_RETURN_NONE;                                          \
    } else {                                                     \
        return cvt_retval(__VA_ARGS__);                          \
    }

inline int cvt_retint(int ret) {
    return ret;
}

#define CVT_RET_INT(...)                                         \
    if constexpr (std::is_same_v<decltype(__VA_ARGS__), void>) { \
        __VA_ARGS__;                                             \
        return 0;                                                \
    } else {                                                     \
        return cvt_retint(__VA_ARGS__);                          \
    }

struct py_err_set : std::exception {};

// refer to pybind11 for the following exception handling helper

inline void pybind11_translate_exception(std::exception_ptr last_exception) {
    auto& registered_exception_translators =
            pybind11::detail::get_internals().registered_exception_translators;
    for (auto& translator : registered_exception_translators) {
        try {
            translator(last_exception);
        } catch (...) {
            last_exception = std::current_exception();
            continue;
        }
        return;
    }
    PyErr_SetString(
            PyExc_SystemError, "Exception escaped from default exception translator!");
}

inline void pybind11_translate_exception() {
    pybind11_translate_exception(std::current_exception());
}

#if defined(__GNUG__) && !defined(__clang__)
#define PYEXT17_TRANSLATE_EXC_CATCH_FORCED_UNWIND \
    catch (::abi::__forced_unwind&) {             \
        throw;                                    \
    }
#else
#define PYEXT17_TRANSLATE_EXC_CATCH_FORCED_UNWIND
#endif

#define PYEXT17_TRANSLATE_EXC                      \
    catch (::pyext17::py_err_set&) {               \
    }                                              \
    catch (::pybind11::error_already_set & e) {    \
        e.restore();                               \
    }                                              \
    PYEXT17_TRANSLATE_EXC_CATCH_FORCED_UNWIND      \
    catch (...) {                                  \
        ::pyext17::pybind11_translate_exception(); \
    }

#define PYEXT17_TRANSLATE_EXC_RET(RET)             \
    catch (::pyext17::py_err_set&) {               \
        return RET;                                \
    }                                              \
    catch (::pybind11::error_already_set & e) {    \
        e.restore();                               \
        return RET;                                \
    }                                              \
    PYEXT17_TRANSLATE_EXC_CATCH_FORCED_UNWIND      \
    catch (...) {                                  \
        ::pyext17::pybind11_translate_exception(); \
        return RET;                                \
    };

template <typename T>
struct wrap {
private:
    typedef wrap<T> wrap_t;

public:
    PyObject_HEAD std::aligned_storage_t<sizeof(T), alignof(T)> storage;
#ifdef _Py_TPFLAGS_HAVE_VECTORCALL
    PyObject* (*vectorcall_slot)(PyObject*, PyObject* const*, size_t, PyObject*);
#endif

    inline T* inst() { return reinterpret_cast<T*>(&storage); }

    inline static PyObject* pycast(T* ptr) {
        return (PyObject*)((char*)ptr - offsetof(wrap_t, storage));
    }

private:
    // method wrapper

    enum struct meth_type { noarg, varkw, fastcall, singarg };

    template <auto f>
    struct detect_meth_type {
        static constexpr meth_type value = []() {
            using F = decltype(f);
            static_assert(std::is_member_function_pointer_v<F>);
            if constexpr (std::is_invocable_v<F, T>) {
                return meth_type::noarg;
            } else if constexpr (std::is_invocable_v<F, T, PyObject*, PyObject*>) {
                return meth_type::varkw;
            } else if constexpr (std::is_invocable_v<
                                         F, T, PyObject* const*, Py_ssize_t>) {
                return meth_type::fastcall;
            } else if constexpr (std::is_invocable_v<F, T, PyObject*>) {
                return meth_type::singarg;
            } else {
                static_assert(!std::is_same_v<F, F>);
            }
        }();
    };

    template <meth_type, auto f>
    struct meth {};

    template <auto f>
    struct meth<meth_type::noarg, f> {
        static constexpr int flags = METH_NOARGS;

        static PyObject* impl(PyObject* self, PyObject*) {
            auto* inst = reinterpret_cast<wrap_t*>(self)->inst();
            try {
                CVT_RET_PYOBJ((inst->*f)());
            }
            PYEXT17_TRANSLATE_EXC_RET(nullptr)
        }
    };

    template <auto f>
    struct meth<meth_type::varkw, f> {
        static constexpr int flags = METH_VARARGS | METH_KEYWORDS;

        static PyObject* impl(PyObject* self, PyObject* args, PyObject* kwargs) {
            auto* inst = reinterpret_cast<wrap_t*>(self)->inst();
            try {
                CVT_RET_PYOBJ((inst->*f)(args, kwargs));
            }
            PYEXT17_TRANSLATE_EXC_RET(nullptr)
        }
    };

    template <auto f>
    struct meth<meth_type::fastcall, f> {
#ifdef METH_FASTCALL
        static constexpr int flags = METH_FASTCALL;

        static PyObject* impl(PyObject* self, PyObject* const* args, Py_ssize_t nargs) {
            auto* inst = reinterpret_cast<wrap_t*>(self)->inst();
            try {
                CVT_RET_PYOBJ((inst->*f)(args, nargs));
            }
            PYEXT17_TRANSLATE_EXC_RET(nullptr)
        }
#else
        static constexpr int flags = METH_VARARGS;

        static PyObject* impl(PyObject* self, PyObject* args) {
            auto* inst = reinterpret_cast<wrap_t*>(self)->inst();
            auto* arr = &PyTuple_GET_ITEM(args, 0);
            auto size = PyTuple_GET_SIZE(args);
            try {
                CVT_RET_PYOBJ((inst->*f)(arr, size));
            }
            PYEXT17_TRANSLATE_EXC_RET(nullptr)
        }
#endif
    };

    template <auto f>
    struct meth<meth_type::singarg, f> {
        static constexpr int flags = METH_O;

        static PyObject* impl(PyObject* self, PyObject* obj) {
            auto* inst = reinterpret_cast<wrap_t*>(self)->inst();
            try {
                CVT_RET_PYOBJ((inst->*f)(obj));
            }
            PYEXT17_TRANSLATE_EXC_RET(nullptr)
        }
    };

    template <auto f>
    static constexpr PyMethodDef make_meth_def(
            const char* name, const char* doc = nullptr) {
        using M = meth<detect_meth_type<f>::value, f>;
        return {name, (PyCFunction)M::impl, M::flags, doc};
    }

    template <auto f>
    struct getter {
        using F = decltype(f);

        static PyObject* impl(PyObject* self, void* closure) {
            auto* inst = reinterpret_cast<wrap_t*>(self)->inst();
            try {
                if constexpr (std::is_invocable_v<F, PyObject*, void*>) {
                    CVT_RET_PYOBJ(f(self, closure));
                } else if constexpr (std::is_invocable_v<F, T, void*>) {
                    CVT_RET_PYOBJ((inst->*f)(closure));
                } else if constexpr (std::is_invocable_v<F, T>) {
                    CVT_RET_PYOBJ((inst->*f)());
                } else {
                    static_assert(!std::is_same_v<F, F>);
                }
            }
            PYEXT17_TRANSLATE_EXC_RET(nullptr)
        }
    };

    template <auto f>
    struct setter {
        using F = decltype(f);

        template <typename = void>
        static int impl_(PyObject* self, PyObject* val, void* closure) {
            auto* inst = reinterpret_cast<wrap_t*>(self)->inst();
            try {
                if constexpr (std::is_invocable_v<F, PyObject*, PyObject*, void*>) {
                    CVT_RET_INT(f(self, val, closure));
                } else if constexpr (std::is_invocable_v<F, T, PyObject*, void*>) {
                    CVT_RET_INT((inst->*f)(val, closure));
                } else if constexpr (std::is_invocable_v<F, T, PyObject*>) {
                    CVT_RET_INT((inst->*f)(val));
                } else {
                    static_assert(!std::is_same_v<F, F>);
                }
            }
            PYEXT17_TRANSLATE_EXC_RET(-1)
        }

        static constexpr auto impl = []() {
            if constexpr (std::is_same_v<F, std::nullptr_t>)
                return nullptr;
            else
                return impl_<>;
        }();
    };

    template <auto get, auto set = nullptr>
    static constexpr PyGetSetDef make_getset_def(
            const char* name, const char* doc = nullptr, void* closure = nullptr) {
        return {const_cast<char*>(name), getter<get>::impl, setter<set>::impl,
                const_cast<char*>(doc), closure};
    }

    // polyfills

    struct tp_vectorcall {
        static constexpr bool valid = HAS_MEMBER(T, tp_vectorcall);
        static constexpr bool haskw = []() {
            if constexpr (valid)
                if constexpr (std::is_invocable_v<
                                      decltype(&T::tp_vectorcall), T, PyObject* const*,
                                      size_t, PyObject*>)
                    return true;
            return false;
        }();

        template <typename = void>
        static PyObject* impl(
                PyObject* self, PyObject* const* args, size_t nargsf,
                PyObject* kwnames) {
            auto* inst = reinterpret_cast<wrap_t*>(self)->inst();
            if constexpr (haskw) {
                CVT_RET_PYOBJ(inst->tp_vectorcall(args, nargsf, kwnames));
            } else {
                if (kwnames && PyTuple_GET_SIZE(kwnames)) {
                    PyErr_SetString(PyExc_TypeError, "expect no keyword argument");
                    return nullptr;
                }
                CVT_RET_PYOBJ(inst->tp_vectorcall(args, nargsf));
            }
        }

        static constexpr Py_ssize_t offset = []() {
            if constexpr (valid)
                return offsetof(wrap_t, vectorcall_slot);
            else
                return 0;
        }();
    };

    struct tp_call {
        static constexpr bool provided = HAS_MEMBER(T, tp_call);
        static constexpr bool static_form =
                invocable_with<T, PyObject*, PyObject*, PyObject*>{}(
                        [](auto&& t, auto... args)
                                -> decltype(std::decay_t<decltype(t)>::tp_call(
                                        args...)) {});
        static constexpr bool valid = provided || tp_vectorcall::valid;

        template <typename = void>
        static PyObject* impl(PyObject* self, PyObject* args, PyObject* kwargs) {
            auto* inst = reinterpret_cast<wrap_t*>(self)->inst();
            CVT_RET_PYOBJ(inst->tp_call(args, kwargs));
        }

        static constexpr ternaryfunc value = []() {
            if constexpr (static_form)
                return T::tp_call;
            else if constexpr (provided)
                return impl<>;
#ifdef _Py_TPFLAGS_HAVE_VECTORCALL
            else if constexpr (valid)
                return PyVectorcall_Call;
#endif
            else
                return nullptr;
        }();
    };

    struct tp_new {
        static constexpr bool provided = HAS_MEMBER(T, tp_new);
        static constexpr bool varkw = std::is_constructible_v<T, PyObject*, PyObject*>;
        static constexpr bool noarg = std::is_default_constructible_v<T>;

        template <typename = void>
        static PyObject* impl(PyTypeObject* type, PyObject* args, PyObject* kwargs) {
            struct FreeGuard {
                PyObject* self;
                PyTypeObject* type;
                ~FreeGuard() {
                    if (self)
                        type->tp_free(self);
                }
            };

            auto* self = type->tp_alloc(type, 0);
            FreeGuard free_guard{self, type};
            auto* inst = reinterpret_cast<wrap_t*>(self)->inst();
            if constexpr (has_vectorcall && tp_vectorcall::valid) {
                reinterpret_cast<wrap_t*>(self)->vectorcall_slot =
                        &tp_vectorcall::template impl<>;
            }
            try {
                if constexpr (varkw) {
                    new (inst) T(args, kwargs);
                } else {
                    new (inst) T();
                }
            }
            PYEXT17_TRANSLATE_EXC_RET(nullptr)
            free_guard.self = nullptr;
            return self;
        }

        static constexpr newfunc value = []() {
            if constexpr (provided)
                return T::tp_new;
            else if constexpr (varkw || noarg)
                return impl<>;
            else
                return nullptr;
        }();
    };

    struct tp_dealloc {
        static constexpr bool provided = HAS_MEMBER(T, tp_dealloc);

        template <typename = void>
        static void impl(PyObject* self) {
            reinterpret_cast<wrap_t*>(self)->inst()->~T();
            Py_TYPE(self)->tp_free(self);
        }

        static constexpr destructor value = []() {
            if constexpr (provided)
                return T::tp_dealloc;
            else
                return impl<>;
        }();
    };

public:
    class TypeBuilder {
        std::vector<PyMethodDef> m_methods;
        std::vector<PyGetSetDef> m_getsets;
        PyTypeObject m_type;
        bool m_finalized = false;
        bool m_ready = false;

        void check_finalized() {
            if (m_finalized) {
                throw std::runtime_error("type is already finalized");
            }
        }

        static const char* to_c_str(const char* s) { return s; }

        template <size_t N, typename... Ts>
        static const char* to_c_str(const pybind11::detail::descr<N, Ts...>& desc) {
            return desc.text;
        }

    public:
        TypeBuilder(const TypeBuilder&) = delete;
        TypeBuilder& operator=(const TypeBuilder&) = delete;

        TypeBuilder() : m_type{PyVarObject_HEAD_INIT(nullptr, 0)} {
            constexpr auto has_tp_name = HAS_MEMBER(T, tp_name);
            if constexpr (has_tp_name) {
                m_type.tp_name = to_c_str(T::tp_name);
            }
            m_type.tp_dealloc = tp_dealloc::value;
#ifdef _Py_TPFLAGS_HAVE_VECTORCALL
            m_type.tp_vectorcall_offset = tp_vectorcall::offset;
#endif
            m_type.tp_call = tp_call::value;
            m_type.tp_basicsize = sizeof(wrap_t);
            m_type.tp_flags = Py_TPFLAGS_DEFAULT | Py_TPFLAGS_BASETYPE;
#ifdef _Py_TPFLAGS_HAVE_VECTORCALL
            if constexpr (tp_vectorcall::valid) {
                m_type.tp_flags |= _Py_TPFLAGS_HAVE_VECTORCALL;
            }
#endif
            m_type.tp_new = tp_new::value;
        }

        PyTypeObject* operator->() { return &m_type; }

        bool ready() const { return m_ready; }

        bool isinstance(PyObject* op) { return PyObject_TypeCheck(op, &m_type); }

        bool isexact(PyObject* op) { return Py_TYPE(op) == &m_type; }

        bool same_pytype(PyTypeObject* pt) { return pt == &m_type; }

        PyObject* finalize() {
            if (!m_finalized) {
                m_finalized = true;
                if (m_methods.size()) {
                    m_methods.push_back({0});
                    if (m_type.tp_methods) {
                        PyErr_SetString(PyExc_SystemError, "tp_method is already set");
                        return nullptr;
                    }
                    m_type.tp_methods = &m_methods[0];
                }
                if (m_getsets.size()) {
                    m_getsets.push_back({0});
                    if (m_type.tp_getset) {
                        PyErr_SetString(PyExc_SystemError, "tp_getset is already set");
                        return nullptr;
                    }
                    m_type.tp_getset = &m_getsets[0];
                }
                if (PyType_Ready(&m_type)) {
                    return nullptr;
                }
                m_ready = true;
            }
            return (PyObject*)&m_type;
        }

        template <auto f>
        TypeBuilder& def(const char* name, const char* doc = nullptr) {
            check_finalized();
            m_methods.push_back(make_meth_def<f>(name, doc));
            return *this;
        }

        template <auto get, auto set = nullptr>
        TypeBuilder& def_getset(
                const char* name, const char* doc = nullptr, void* closure = nullptr) {
            check_finalized();
            m_getsets.push_back(make_getset_def<get, set>(name, doc, closure));
            return *this;
        }
    };

    static TypeBuilder& type() {
        static TypeBuilder type_helper;
        return type_helper;
    }

    template <typename... Args>
    static PyObject* cnew(Args&&... args) {
        auto* pytype = type().operator->();
        return cnew_with_type(pytype, std::forward<Args>(args)...);
    }

    template <typename... Args>
    static PyObject* cnew_with_type(PyTypeObject* pytype, Args&&... args) {
        auto* self = pytype->tp_alloc(pytype, 0);
        auto* inst = reinterpret_cast<wrap_t*>(self)->inst();
        if constexpr (has_vectorcall && tp_vectorcall::valid) {
            reinterpret_cast<wrap_t*>(self)->vectorcall_slot =
                    &tp_vectorcall::template impl<>;
        }
        new (inst) T(std::forward<Args>(args)...);
        return self;
    }

    struct caster {
        static constexpr auto name = T::tp_name;

        T* value;

        bool load(pybind11::handle src, bool convert) {
            if (wrap_t::type().isinstance(src.ptr())) {
                value = reinterpret_cast<wrap_t*>(src.ptr())->inst();
                return true;
            }
            return false;
        }

        template <typename U>
        using cast_op_type = pybind11::detail::cast_op_type<U>;
        operator T*() { return value; }
        operator T&() { return *value; }
    };
};

}  // namespace pyext17

#undef HAS_MEMBER_TYPE
#undef HAS_MEMBER
#undef CVT_RET_PYOBJ
#undef CVT_RET_INT