pocketpy/include/pybind11/internal/cpp_function.h

#pragma once

#include "cast.h"
#include <bitset>

namespace pybind11 {

template <std::size_t Nurse, std::size_t... Patients>
struct keep_alive {};

template <typename T>
struct call_guard {
    static_assert(std::is_default_constructible_v<T>, "call_guard must be default constructible");
};

// append the overload to the beginning of the overload list
struct prepend {};

template <typename... Args>
struct init {};

// TODO: support more customized tags
// struct kw_only {};
//
// struct pos_only {};
//
// struct default_arg {};
//
// struct arg {
//    const char* name;
//    const char* description;
// };
//
// struct default_arg {
//    const char* name;
//    const char* description;
//    const char* value;
// };

template <typename Fn,
          typename Extra,
          typename Args = callable_args_t<std::decay_t<Fn>>,
          typename IndexSequence = std::make_index_sequence<std::tuple_size_v<Args>>>
struct generator;

class function_record {
    union {
        void* data;
        char buffer[16];
    };

    // TODO: optimize the function_record size to reduce memory usage
    const char* name;
    function_record* next;
    void (*destructor)(function_record*);
    return_value_policy policy = return_value_policy::automatic;
    handle (*wrapper)(function_record&, pkpy::ArgsView, bool convert, handle parent);

    template <typename Fn, typename Extra, typename Args, typename IndexSequence>
    friend struct generator;

public:
    template <typename Fn, typename... Extras>
    function_record(Fn&& f, const char* name, const Extras&... extras) : name(name), next(nullptr) {

        if constexpr(sizeof(f) <= sizeof(buffer)) {
            new (buffer) auto(std::forward<Fn>(f));
            destructor = [](function_record* self) {
                reinterpret_cast<Fn*>(self->buffer)->~Fn();
            };
        } else {
            data = new auto(std::forward<Fn>(f));
            destructor = [](function_record* self) {
                delete static_cast<Fn*>(self->data);
            };
        }

        using Generator = generator<std::decay_t<Fn>, std::tuple<Extras...>>;
        Generator::initialize(*this, extras...);
        wrapper = Generator::generate();
    }

    ~function_record() { destructor(this); }

    template <typename Fn>
    auto& cast() {
        if constexpr(sizeof(Fn) <= sizeof(buffer)) {
            return *reinterpret_cast<Fn*>(buffer);
        } else {
            return *static_cast<Fn*>(data);
        }
    }

    void append(function_record* record) {
        function_record* p = this;
        while(p->next != nullptr) {
            p = p->next;
        }
        p->next = record;
    }

    handle operator() (pkpy::ArgsView view) {
        function_record* p = this;
        // foreach function record and call the function with not convert
        while(p != nullptr) {
            handle result = p->wrapper(*this, view, false, {});
            if(result) { return result; }
            p = p->next;
        }

        p = this;
        // foreach function record and call the function with convert
        while(p != nullptr) {
            handle result = p->wrapper(*this, view, true, {});
            if(result) { return result; }
            p = p->next;
        }

        vm->TypeError("no matching function found");
    }
};

template <typename Fn, std::size_t... Is, typename... Args>
handle invoke(Fn&& fn,
              std::index_sequence<Is...>,
              std::tuple<type_caster<Args>...>& casters,
              return_value_policy policy,
              handle parent) {
    using underlying_type = std::decay_t<Fn>;
    using ret = callable_return_t<underlying_type>;

    // if the return type is void, return None
    if constexpr(std::is_void_v<ret>) {
        // resolve the member function pointer
        if constexpr(std::is_member_function_pointer_v<underlying_type>) {
            [&](class_type_t<underlying_type>& self, auto&... args) {
                (self.*fn)(args...);
            }(std::get<Is>(casters).value...);
        } else {
            fn(std::get<Is>(casters).value...);
        }
        return vm->None;
    } else {
        // resolve the member function pointer
        if constexpr(std::is_member_function_pointer_v<remove_cvref_t<Fn>>) {
            return type_caster<ret>::cast(
                [&](class_type_t<underlying_type>& self, auto&... args) {
                    return (self.*fn)(args...);
                }(std::get<Is>(casters).value...),
                policy,
                parent);
        } else {
            return type_caster<ret>::cast(fn(std::get<Is>(casters).value...), policy, parent);
        }
    }
}

template <typename Fn, typename... Args, std::size_t... Is, typename... Extras>
struct generator<Fn, std::tuple<Extras...>, std::tuple<Args...>, std::index_sequence<Is...>> {
    static void initialize(function_record& record, const Extras&... extras) {}

    static auto generate() {
        return +[](function_record& self, pkpy::ArgsView view, bool convert, handle parent) {
            // FIXME:
            // Temporarily, args and kwargs must be at the end of the arguments list
            // Named arguments are not supported yet
            constexpr bool has_args = types_count_v<args, remove_cvref_t<Args>...> != 0;
            constexpr bool has_kwargs = types_count_v<kwargs, remove_cvref_t<Args>...> != 0;
            constexpr std::size_t count = sizeof...(Args) - has_args - has_kwargs;

            handle stack[sizeof...(Args)] = {};

            // initialize the stack

            if(!has_args && (view.size() != count)) { return handle(); }

            if(has_args && (view.size() < count)) { return handle(); }

            for(std::size_t i = 0; i < count; ++i) {
                stack[i] = view[i];
            }

            // pack the args and kwargs
            if constexpr(has_args) {
                const auto n = view.size() - count;
                pkpy::PyVar var = vm->new_object<pkpy::Tuple>(vm->tp_tuple, n);
                auto& tuple = var.obj_get<pkpy::Tuple>();
                for(std::size_t i = 0; i < n; ++i) {
                    tuple[i] = view[count + i];
                }
                stack[count] = var;
            }

            if constexpr(has_kwargs) {
                const auto n = vm->s_data._sp - view.end();
                pkpy::PyVar var = vm->new_object<pkpy::Dict>(vm->tp_dict);
                auto& dict = var.obj_get<pkpy::Dict>();

                for(std::size_t i = 0; i < n; i += 2) {
                    pkpy::i64 index = pkpy::_py_cast<pkpy::i64>(vm, view[count + i]);
                    pkpy::PyVar str = vm->new_object<pkpy::Str>(vm->tp_str, pkpy::StrName(index).sv());
                    dict.set(vm, str, view[count + i + 1]);
                }

                stack[count + 1] = var;
            }

            // check if all the arguments are not valid
            for(std::size_t i = 0; i < sizeof...(Args); ++i) {
                if(!stack[i]) { return handle(); }
            }

            // ok, all the arguments are valid, call the function
            std::tuple<type_caster<Args>...> casters;

            // check type compatibility
            if(((std::get<Is>(casters).load(stack[Is], convert)) && ...)) {
                return invoke(self.cast<Fn>(), std::index_sequence<Is...>{}, casters, self.policy, parent);
            }

            return handle();
        };
    }
};

constexpr inline static auto _wrapper = +[](pkpy::VM*, pkpy::ArgsView view) {
    auto& record = pkpy::lambda_get_userdata<function_record>(view.begin());
    return record(view).ptr();
};

class cpp_function : public function {
public:
    template <typename Fn, typename... Extras>
    cpp_function(Fn&& f, const Extras&... extras) {
        pkpy::any userdata = function_record(std::forward<Fn>(f), "anonymous", extras...);
        m_ptr = vm->bind_func(nullptr, "", -1, _wrapper, std::move(userdata));
        inc_ref();
    }
};

template <typename Fn, typename... Extras>
handle bind_function(const handle& obj, const char* name, Fn&& fn, pkpy::BindType type, const Extras&... extras) {
    // do not use cpp_function directly to avoid unnecessary reference count change
    pkpy::PyVar var = obj.ptr();
    pkpy::PyVar callable = var->attr().try_get(name);

    // if the function is not bound yet, bind it
    if(!callable) {
        pkpy::any userdata = function_record(std::forward<Fn>(fn), name, extras...);
        callable = vm->bind_func(var, name, -1, _wrapper, std::move(userdata));
    } else {
        auto& userdata = callable.obj_get<pkpy::NativeFunc>()._userdata;
        function_record* record = new function_record(std::forward<Fn>(fn), name, extras...);

        constexpr bool is_prepend = (types_count_v<prepend, Extras...> != 0);
        if constexpr(is_prepend) {
            // if prepend is specified, append the new record to the beginning of the list
            function_record* last = (function_record*)userdata.data;
            userdata.data = record;
            record->append(last);
        } else {
            // otherwise, append the new record to the end of the list
            function_record* last = (function_record*)userdata.data;
            last->append(record);
        }
    }
    return callable;
}

template <typename Getter_, typename Setter_, typename... Extras>
handle
    bind_property(const handle& obj, const char* name, Getter_&& getter_, Setter_&& setter_, const Extras&... extras) {
    pkpy::PyVar var = obj.ptr();
    pkpy::PyVar getter = vm->None;
    pkpy::PyVar setter = vm->None;
    using Getter = std::decay_t<Getter_>;
    using Setter = std::decay_t<Setter_>;

    getter = vm->new_object<pkpy::NativeFunc>(
        vm->tp_native_func,
        [](pkpy::VM* vm, pkpy::ArgsView view) -> pkpy::PyVar {
            auto& getter = pkpy::lambda_get_userdata<Getter>(view.begin());

            if constexpr(std::is_member_pointer_v<Getter>) {
                using Self = class_type_t<Getter>;
                auto& self = _builtin_cast<instance>(view[0]).cast<Self>();

                if constexpr(std::is_member_object_pointer_v<Getter>) {
                    return type_caster<member_type_t<Getter>>::cast(self.*getter,
                                                                    return_value_policy::reference_internal,
                                                                    view[0])
                        .ptr();
                } else {
                    return type_caster<callable_return_t<Getter>>::cast((self.*getter)(),
                                                                        return_value_policy::reference_internal,
                                                                        view[0])
                        .ptr();
                }

            } else {
                using Self = std::tuple_element_t<0, callable_args_t<Getter>>;
                auto& self = _builtin_cast<instance>(view[0]).cast<Self>();

                return type_caster<callable_return_t<Getter>>::cast(getter(self),
                                                                    return_value_policy::reference_internal,
                                                                    view[0])
                    .ptr();
            }
        },
        1,
        std::forward<Getter_>(getter_));

    if constexpr(!std::is_same_v<Setter, std::nullptr_t>) {
        setter = vm->new_object<pkpy::NativeFunc>(
            vm->tp_native_func,
            [](pkpy::VM* vm, pkpy::ArgsView view) -> pkpy::PyVar {
                auto& setter = pkpy::lambda_get_userdata<Setter>(view.begin());

                if constexpr(std::is_member_pointer_v<Setter>) {
                    using Self = class_type_t<Setter>;
                    auto& self = _builtin_cast<instance>(view[0]).cast<Self>();

                    if constexpr(std::is_member_object_pointer_v<Setter>) {
                        type_caster<member_type_t<Setter>> caster;
                        if(caster.load(view[1], true)) {
                            self.*setter = caster.value;
                            return vm->None;
                        }
                    } else {
                        type_caster<std::tuple_element_t<1, callable_args_t<Setter>>> caster;
                        if(caster.load(view[1], true)) {
                            (self.*setter)(caster.value);
                            return vm->None;
                        }
                    }
                } else {
                    using Self = std::tuple_element_t<0, callable_args_t<Setter>>;
                    auto& self = _builtin_cast<instance>(view[0]).cast<Self>();

                    type_caster<std::tuple_element_t<1, callable_args_t<Setter>>> caster;
                    if(caster.load(view[1], true)) {
                        setter(self, caster.value);
                        return vm->None;
                    }
                }

                vm->TypeError("invalid argument");
            },
            2,
            std::forward<Setter_>(setter_));
    }

    pkpy::PyVar property = vm->new_object<pkpy::Property>(vm->tp_property, getter, setter);
    var->attr().set(name, property);
    return property;
}

}  // namespace pybind11