From 44aed24d3b53793dbd349c17cfb20e001b8b36af Mon Sep 17 00:00:00 2001
From: blueloveTH <blueloveth@foxmail.com>
Date: Thu, 9 May 2024 17:26:11 +0800
Subject: [PATCH] add `function`

---
 include/pocketpy/any.h | 28 +++++++++++++++++++++++++++-
 include/pocketpy/vm.h  |  8 +++++---
 src/pocketpy.cpp       |  2 +-
 src/vm.cpp             |  1 +
 4 files changed, 34 insertions(+), 5 deletions(-)

diff --git a/include/pocketpy/any.h b/include/pocketpy/any.h
index 40daaa9d..3b32a449 100644
--- a/include/pocketpy/any.h
+++ b/include/pocketpy/any.h
@@ -28,12 +28,13 @@ struct any{
 
     any() : data(nullptr), _vt(nullptr) {}
 
-    explicit operator bool() const { return _vt != nullptr; }
+    operator bool() const { return _vt != nullptr; }
 
     template<typename T>
     any(T&& value){
         using U = std::decay_t<T>;
         static_assert(!std::is_same_v<U, any>, "any(const any&) is deleted");
+        static_assert(sizeof(U) == sizeof(T));
         if constexpr (is_sso_v<U>){
             memcpy(&data, &value, sizeof(U));
         }else{
@@ -74,4 +75,29 @@ struct any{
     static void __bad_any_cast(const std::type_index expected, const std::type_index actual);
 };
 
+template<typename T>
+struct function;
+
+template<typename Ret, typename... Params>
+struct function<Ret(Params...)>{
+    any _impl;
+    Ret (*_wrapper)(const any&, Params...);
+
+    function(): _impl(), _wrapper(nullptr) {}
+
+    operator bool() const { return _wrapper != nullptr; }
+
+    template<typename F>
+    function(F&& f) : _impl(std::forward<F>(f)){
+        _wrapper = [](const any& impl, Params... params) -> Ret{
+            return impl.cast<std::decay_t<F>>()(std::forward<Params>(params)...);
+        };
+    }
+
+    Ret operator()(Params... params) const{
+        if(!_wrapper) throw std::runtime_error("empty function");
+        return _wrapper(_impl, std::forward<Params>(params)...);
+    }
+};
+
 } // namespace pkpy
\ No newline at end of file
diff --git a/include/pocketpy/vm.h b/include/pocketpy/vm.h
index 3f79a2ac..98ef8601 100644
--- a/include/pocketpy/vm.h
+++ b/include/pocketpy/vm.h
@@ -157,17 +157,19 @@ public:
     PyObject* __cached_object_new;
     std::map<std::string_view, CodeObject_> __cached_codes;
 
-    void (*_ceval_on_step)(VM*, Frame*, Bytecode bc) = nullptr;
-
 #if PK_ENABLE_PROFILER
     LineProfiler* _profiler = nullptr;
     NextBreakpoint _next_breakpoint;
 #endif
 
+    void (*_ceval_on_step)(VM*, Frame*, Bytecode bc);
     void(*_stdout)(const char*, int);
     void(*_stderr)(const char*, int);
     unsigned char* (*_import_handler)(const char*, int*);
-
+    // function<void(const char*, int)> _stdout;
+    // function<void(const char*, int)> _stderr;
+    // function<unsigned char*(const char*, int*)> _import_handler;
+    
     // for quick access
     static constexpr Type tp_object=Type(0), tp_type=Type(1);
     static constexpr Type tp_int=Type(kTpIntIndex), tp_float=Type(kTpFloatIndex), tp_bool=Type(4), tp_str=Type(5);
diff --git a/src/pocketpy.cpp b/src/pocketpy.cpp
index bb1da213..fa7b11ab 100644
--- a/src/pocketpy.cpp
+++ b/src/pocketpy.cpp
@@ -1638,7 +1638,7 @@ void VM::__post_init_builtin_types(){
     if(enable_os){
         add_module_io(this);
         add_module_os(this);
-        _import_handler = _default_import_handler;
+        _import_handler = &_default_import_handler;
     }
 
     add_module_csv(this);
diff --git a/src/vm.cpp b/src/vm.cpp
index 5c23521d..ff46163c 100644
--- a/src/vm.cpp
+++ b/src/vm.cpp
@@ -75,6 +75,7 @@ namespace pkpy{
     VM::VM(bool enable_os) : heap(this), enable_os(enable_os) {
         this->vm = this;
         this->__c.error = nullptr;
+        _ceval_on_step = nullptr;
         _stdout = [](const char* buf, int size) { std::cout.write(buf, size); };
         _stderr = [](const char* buf, int size) { std::cerr.write(buf, size); };
         _main = nullptr;