From 1e7fe289989314c876232782e5f9cc7f8037c8d8 Mon Sep 17 00:00:00 2001
From: BLUELOVETH <blueloveTH@foxmail.com>
Date: Thu, 3 Aug 2023 11:42:35 +0800
Subject: [PATCH] native modules are lazy by default

---
 docs/quick-start/modules.md |  8 ++++----
 include/pocketpy/vm.h       |  2 ++
 src/pocketpy.cpp            | 25 +++++++++++++------------
 src/vm.cpp                  |  9 +++++++++
 4 files changed, 28 insertions(+), 16 deletions(-)

diff --git a/docs/quick-start/modules.md b/docs/quick-start/modules.md
index 3a8c76ff..663be989 100644
--- a/docs/quick-start/modules.md
+++ b/docs/quick-start/modules.md
@@ -25,7 +25,7 @@ print(test.pi)  # 3.14
 ### Native modules
 
 A native module is a module written in c++ or mixed c++/python.
-Native modules are always compiled and executed when the VM is created.
+Native modules can be added into `_modules` or `_lazy_native_modules`.
 
 To creata a native module,
 use `vm->new_module(...)`.
@@ -52,9 +52,9 @@ print(test.add(1, 2))  # 3
 When you do `import` a module, the VM will try to find it in the following order:
 
 1. Search `vm->_modules`, if found, return it.
-2. Search `vm->_lazy_modules`, if found, compile and execute it, then return it.
-3. Try `vm->_import_handler`.
-
+2. Search `vm->_lazy_native_modules`, if found, compile and execute it, then return it.
+3. Search `vm->_lazy_modules`, if found, compile and execute it, then return it.
+4. Try `vm->_import_handler`.
 
 ### Customized import handler
 
diff --git a/include/pocketpy/vm.h b/include/pocketpy/vm.h
index adde17e3..142f579e 100644
--- a/include/pocketpy/vm.h
+++ b/include/pocketpy/vm.h
@@ -47,6 +47,7 @@ namespace pkpy{
 
 
 typedef PyObject* (*BinaryFuncC)(VM*, PyObject*, PyObject*);
+typedef void (*NativeModuleInitializer)(VM*);
 
 struct PyTypeInfo{
     PyObject* obj;
@@ -116,6 +117,7 @@ public:
     
     NameDict _modules;                                 // loaded modules
     std::map<StrName, Str> _lazy_modules;              // lazy loaded modules
+    std::map<StrName, NativeModuleInitializer> _lazy_native_modules;  // lazy loaded native modules
 
     struct{
         PyObject* error;
diff --git a/src/pocketpy.cpp b/src/pocketpy.cpp
index c8a96b34..d5e32330 100644
--- a/src/pocketpy.cpp
+++ b/src/pocketpy.cpp
@@ -1554,17 +1554,18 @@ void VM::post_init(){
 
 #if !PK_DEBUG_NO_BUILTINS
     add_module_sys(this);
-    add_module_traceback(this);
-    add_module_time(this);
-    add_module_json(this);
-    add_module_math(this);
-    add_module_re(this);
-    add_module_dis(this);
     add_module_c(this);
-    add_module_gc(this);
-    add_module_random(this);
-    add_module_base64(this);
-    add_module_timeit(this);
+    // lazy modules
+    vm->_lazy_native_modules["traceback"] = &add_module_traceback;
+    vm->_lazy_native_modules["time"] = &add_module_time;
+    vm->_lazy_native_modules["json"] = &add_module_json;
+    vm->_lazy_native_modules["math"] = &add_module_math;
+    vm->_lazy_native_modules["re"] = &add_module_re;
+    vm->_lazy_native_modules["dis"] = &add_module_dis;
+    vm->_lazy_native_modules["gc"] = &add_module_gc;
+    vm->_lazy_native_modules["random"] = &add_module_random;
+    vm->_lazy_native_modules["base64"] = &add_module_base64;
+    vm->_lazy_native_modules["timeit"] = &add_module_timeit;
 
     for(const char* name: {"this", "functools", "collections", "heapq", "bisect", "pickle", "_long", "colorsys"}){
         _lazy_modules[name] = kPythonLibs[name];
@@ -1587,8 +1588,8 @@ void VM::post_init(){
         _import_handler = _default_import_handler;
     }
 
-    add_module_linalg(this);
-    add_module_easing(this);
+    vm->_lazy_native_modules["linalg"] = &add_module_linalg;
+    vm->_lazy_native_modules["easing"] = &add_module_easing;
 #endif
 }
 
diff --git a/src/vm.cpp b/src/vm.cpp
index 8381dc1a..113e7b8e 100644
--- a/src/vm.cpp
+++ b/src/vm.cpp
@@ -239,6 +239,14 @@ namespace pkpy{
         }
         PyObject* ext_mod = _modules.try_get(name);
         if(ext_mod == nullptr){
+            auto it2 = _lazy_native_modules.find(name);
+            if(it2 != _lazy_native_modules.end()){
+                it2->second(this);  // init
+                ext_mod = _modules.try_get(name);
+                PK_ASSERT(ext_mod != nullptr);
+                _lazy_native_modules.erase(it2);
+                return ext_mod;
+            }
             Str source;
             auto it = _lazy_modules.find(name);
             if(it == _lazy_modules.end()){
@@ -271,6 +279,7 @@ namespace pkpy{
         _all_types.clear();
         _modules.clear();
         _lazy_modules.clear();
+        _lazy_native_modules.clear();
     }
 
 PyObject* VM::py_negate(PyObject* obj){