From 5a2d389c22f6ecdbe819bf77400ebdc12c80a33e Mon Sep 17 00:00:00 2001
From: blueloveTH <blueloveth@hclcat.games>
Date: Wed, 30 Nov 2022 21:25:08 +0800
Subject: [PATCH] change src module to lazy

---
 src/obj.h       |  2 +-
 src/pocketpy.h  |  5 +++--
 src/vm.h        | 29 +++++++++++++++++++++++------
 tests/random.py |  6 ++++++
 4 files changed, 33 insertions(+), 9 deletions(-)
 create mode 100644 tests/random.py

diff --git a/src/obj.h b/src/obj.h
index d38c5b81..ab3b82cf 100644
--- a/src/obj.h
+++ b/src/obj.h
@@ -10,7 +10,7 @@ const _Int _Int_MAX_NEG = -9223372036854775807LL;
 const _Float _FLOAT_INF_POS = INFINITY;
 const _Float _FLOAT_INF_NEG = -INFINITY;
 
-#define PK_VERSION "0.3.4"
+#define PK_VERSION "0.3.5"
 
 class CodeObject;
 class BasePointer;
diff --git a/src/pocketpy.h b/src/pocketpy.h
index e9c1913e..c90ac45d 100644
--- a/src/pocketpy.h
+++ b/src/pocketpy.h
@@ -699,10 +699,11 @@ extern "C" {
 
     __EXPORT
     bool pkpy_add_module(VM* vm, const char* name, const char* source){
+        // compile the module but don't execute it
         _Code code = compile(vm, source, name + _Str(".py"));
         if(code == nullptr) return false;
-        PyVar _m = vm->newModule(name);
-        return vm->exec(code, _m) != nullptr;
+        vm->addLazyModule(name, code);
+        return true;
     }
 
     void __vm_init(VM* vm){
diff --git a/src/vm.h b/src/vm.h
index 090d7352..e7a94107 100644
--- a/src/vm.h
+++ b/src/vm.h
@@ -46,7 +46,8 @@ typedef void(*PrintFn)(const VM*, const char*);
 class VM: public PkExportedResource{
 protected:
     std::deque< std::unique_ptr<Frame> > callstack;
-    PyVarDict _modules;       // 3rd modules
+    PyVarDict _modules;                     // loaded modules
+    std::map<_Str, _Code> _lazyModules;     // lazy loaded modules
     PyVar __py2py_call_signal;
 
     PyVar runFrame(Frame* frame){
@@ -329,8 +330,20 @@ protected:
                 {
                     const _Str& name = frame->code->co_names[byte.arg]->name;
                     auto it = _modules.find(name);
-                    if(it == _modules.end()) _error("ImportError", "module '" + name + "' not found");
-                    else frame->push(it->second); 
+                    if(it == _modules.end()){
+                        auto it2 = _lazyModules.find(name);
+                        if(it2 == _lazyModules.end()){
+                            _error("ImportError", "module '" + name + "' not found");
+                        }else{
+                            _Code code = it2->second;
+                            PyVar _m = newModule(name);
+                            _exec(code, _m, {});
+                            frame->push(_m);
+                            _lazyModules.erase(it2);
+                        }
+                    }else{
+                        frame->push(it->second);
+                    }
                 } break;
             default:
                 systemError(_Str("opcode ") + OP_NAMES[byte.op] + " is not implemented");
@@ -585,13 +598,17 @@ public:
         return obj;
     }
 
-    PyVar newModule(_Str name, bool saveToPath=true) {
+    PyVar newModule(_Str name) {
         PyVar obj = newObject(_tp_module, (_Int)-2);
         setAttr(obj, __name__, PyStr(name));
-        if(saveToPath) _modules[name] = obj;
+        _modules[name] = obj;
         return obj;
     }
 
+    void addLazyModule(_Str name, _Code code){
+        _lazyModules[name] = code;
+    }
+
     PyVarOrNull getAttr(const PyVar& obj, const _Str& name, bool throw_err=true) {
         PyVarDict::iterator it;
         PyObject* cls;
@@ -792,7 +809,7 @@ public:
         this->True = newObject(_tp_bool, true);
         this->False = newObject(_tp_bool, false);
         this->builtins = newModule("builtins");
-        this->_main = newModule("__main__"_c, false);
+        this->_main = newModule("__main__"_c);
 
         setAttr(_tp_type, __base__, _tp_object);
         _tp_type->setType(_tp_type);
diff --git a/tests/random.py b/tests/random.py
new file mode 100644
index 00000000..64abcef9
--- /dev/null
+++ b/tests/random.py
@@ -0,0 +1,6 @@
+import random
+
+for _ in range(100):
+    i = random.randint(1, 10)
+    assert i <= 10
+    assert i >= 1
\ No newline at end of file