From 4b11dc1fe70053335bf530fad98ebf94a669bdd0 Mon Sep 17 00:00:00 2001
From: Serg Creevanose <sergeikrivonos@gmail.com>
Date: Wed, 20 Mar 2024 07:48:51 -0700
Subject: [PATCH] Skrypt: WaitForModule if IsModuleLoading to get
 MappedModuleVariable

---
 libskrypt/skrypt.cpp | 49 ++++++++++++++++++++++----------------------
 libskrypt/skrypt.h   |  7 +++++++
 2 files changed, 31 insertions(+), 25 deletions(-)

diff --git a/libskrypt/skrypt.cpp b/libskrypt/skrypt.cpp
index 668242c..feebda4 100644
--- a/libskrypt/skrypt.cpp
+++ b/libskrypt/skrypt.cpp
@@ -55,6 +55,8 @@ const ::omnn::math::Variable& Skrypt::MappedModuleVariable(const ::omnn::math::V
         );
         if (!module) {
             module = Module(moduleName);
+        } else if (IsModuleLoading(moduleName)) {
+            WaitForModule(moduleName);
         }
         auto& moduleVarNames = module->GetVarNames();
         auto moduleVarIt = moduleVarNames.find(name);
@@ -411,6 +413,17 @@ bool Skrypt::IsModuleLoading(std::string_view name) const {
     return modulesLoading.contains(name);
 }
 
+void Skrypt::WaitForModule(std::string_view name) {
+    for (;;)
+    {
+        auto loaded = modulesLoadingQueue.PeekNextResult();
+        auto it = loaded.find(name);
+        if (it != loaded.end()) {
+            break;
+        }
+    }
+}
+
 Skrypt::module_t Skrypt::Module(std::string_view name) {
     auto module = GetLoadedModule(name);
     if (!module) {
@@ -449,31 +462,21 @@ Skrypt::module_t Skrypt::Module(std::string_view name) {
             module->Load(FindModulePath(name));
 
 			std::unique_lock lock(modulesLoadingMutex);
-            auto loading = modulesLoading.find(std::string(name));
+            auto loading = modulesLoading.find(name);
             if (loading != modulesLoading.end()) {
 				modulesLoading.erase(loading);
             } else {
                 std::cerr << "Module " << name << " was not found in the loading map" << std::endl; // race condition
             }
         } else {
-            // wait for the module to be loaded
-            bool waiting = true;
-            do {
-                auto loaded = modulesLoadingQueue.PeekNextResult();
-                auto it = loaded.find(name);
-                if (it != loaded.end()) {
-                    if (module != it->second.get()) {
-						IMPLEMENT
-					}
-					waiting = false;
-				}
-            } while (waiting);
+            WaitForModule(name);
         }
     }
 	return module;
 }
 
 Skrypt::loading_module_t Skrypt::StartLoadingModule(std::string_view name) {
+    std::cout << "Module " << name << " loading started" << std::endl;
     return std::async(
         std::launch::async, [this, name]() {
             auto module = Module(name);
@@ -489,20 +492,16 @@ Skrypt::loading_modules_t Skrypt::LoadModules(const ::omnn::math::Valuable& v) {
         if (dot != std::string::npos) {
             auto moduleName = name.substr(0, dot);
             auto loading = loadingModules.find(std::string(moduleName));
-            if (loading == loadingModules.end()) {
-                // not loading yet
+            if (loading == loadingModules.end()) { // not loading here yet
+                auto loaded = GetLoadedModule(moduleName);
+                if (loaded)
                 {
-                    std::shared_lock lock(modulesMapMutex);
-                    auto loaded = modules.find(moduleName);
-                    if (loaded != modules.end()) {
-                        // already loaded
-                        std::promise<module_t> promise;
-                        promise.set_value(loaded->second);
-                        loadingModules.emplace(moduleName, promise.get_future());
-                        continue;
-                    }
+                    std::promise<module_t> promise;
+                    promise.set_value(loaded);
+                    loadingModules.emplace(moduleName, promise.get_future());
+                } else {
+                    loadingModules.emplace(moduleName, StartLoadingModule(moduleName));
                 }
-                loadingModules.emplace(moduleName, StartLoadingModule(moduleName));
             }
         }
     }
diff --git a/libskrypt/skrypt.h b/libskrypt/skrypt.h
index c9f961b..baee50f 100644
--- a/libskrypt/skrypt.h
+++ b/libskrypt/skrypt.h
@@ -91,6 +91,13 @@ class Skrypt
 		moduleFileSearchAdditionalPaths.emplace_back(std::forward<T>(p));
     }
     bool IsModuleLoading(std::string_view name) const;
+
+    /// <summary>
+    /// Wait for the module to be loaded
+    /// </summary>
+    /// <param name="name"></param>
+    /// <returns></returns>
+    void WaitForModule(std::string_view name);
     loading_module_t StartLoadingModule(std::string_view name);
     loading_modules_t LoadModules(const ::omnn::math::Valuable& v);
     loading_modules_future_t StartLoadingModules(const ::omnn::math::Valuable& v);