Added better memory handling on dependenciesimpl/sm64-romhacks

1 files changed, 59 insertions, 8 deletions

diff --git a/src/Companion.cpp b/src/Companion.cpp
index 06ed9f0..180fad2 100644
--- a/src/Companion.cpp
+++ b/src/Companion.cpp
@@ -1752,15 +1752,17 @@ void Companion::Process(std::atomic<size_t>& assetCount) {
                path.find("config.yml") == std::string::npos && path.find("hashes.yaml") == std::string::npos;
     };
 
-    // A file another yaml lists under external_files is looked up again when that yaml is processed,
-    // and a processed file is never reparsed, so only files nothing depends on are freed after export.
-    std::unordered_set<std::string> externalDependencies;
+    // A file another yaml lists under external_files is looked up again when that yaml is processed, and a
+    // processed file is never reparsed, so it is only freed once every yaml that lists it has been processed.
+    std::unordered_map<std::string, std::vector<std::string>> externalFilesOf;
+    std::unordered_map<std::string, int> pendingDependents;
     if (this->gConfig.exporterType == ExportType::Binary) {
         for (const auto& entry : entries) {
             if (!isAssetYaml(entry)) {
                 continue;
             }
-            const YAML::Node root = YAML::LoadFile(entry.path().generic_string());
+            const auto yamlPath = entry.path().generic_string();
+            const YAML::Node root = YAML::LoadFile(yamlPath);
             const auto config = root[":config"];
             if (!config) {
                 continue;
@@ -1771,14 +1773,55 @@ void Companion::Process(std::atomic<size_t>& assetCount) {
             }
             for (const auto& externalFile : externalFiles) {
                 if (externalFile.IsScalar()) {
-                    externalDependencies.insert(
-                        (this->gSourceDirectory / externalFile.as<std::string>()).generic_string());
+                    const auto path = (this->gSourceDirectory / externalFile.as<std::string>()).generic_string();
+                    externalFilesOf[yamlPath].push_back(path);
+                    pendingDependents[path]++;
                 }
             }
         }
     }
+    std::unordered_set<std::string> finishedFiles;
 
-    for (const auto& entry : entries) {
+    // Process a yaml right after the files it lists under external_files, when it lists nothing else not yet
+    // reached, so those files can be freed instead of staying in memory until it comes up in directory order.
+    std::vector<fs::directory_entry> orderedEntries;
+    if (this->gConfig.exporterType == ExportType::Binary) {
+        std::unordered_map<std::string, std::vector<const fs::directory_entry*>> dependentsOf;
+        for (const auto& entry : entries) {
+            for (const auto& external : externalFilesOf[entry.path().generic_string()]) {
+                dependentsOf[external].push_back(&entry);
+            }
+        }
+        std::unordered_set<std::string> placed;
+        std::function<void(const fs::directory_entry&)> place = [&](const fs::directory_entry& entry) {
+            const auto path = entry.path().generic_string();
+            if (!placed.insert(path).second) {
+                return;
+            }
+            orderedEntries.push_back(entry);
+            const auto& ownExternals = externalFilesOf[path];
+            for (const auto* dependent : dependentsOf[path]) {
+                bool ready = true;
+                for (const auto& external : externalFilesOf[dependent->path().generic_string()]) {
+                    if (!placed.contains(external) &&
+                        std::find(ownExternals.begin(), ownExternals.end(), external) == ownExternals.end()) {
+                        ready = false;
+                        break;
+                    }
+                }
+                if (ready) {
+                    place(*dependent);
+                }
+            }
+        };
+        for (const auto& entry : entries) {
+            place(entry);
+        }
+    } else {
+        orderedEntries = entries;
+    }
+
+    for (const auto& entry : orderedEntries) {
         if (!isAssetYaml(entry)) {
             continue;
         }
@@ -1910,7 +1953,7 @@ void Companion::Process(std::atomic<size_t>& assetCount) {
             // Keep memory usage low during extraction
             if (this->gConfig.exporterType == ExportType::Binary) {
                 auto freeFile = [&](const std::string& file) {
-                    if (externalDependencies.contains(file)) {
+                    if (pendingDependents[file] > 0) {
                         return;
                     }
                     this->gParseResults.erase(file);
@@ -1918,10 +1961,18 @@ void Companion::Process(std::atomic<size_t>& assetCount) {
                     this->gNodesByTypeCache.erase(file);
                     this->gNodesByTypeCacheValid.erase(file);
                 };
+                finishedFiles.insert(yamlPath);
                 freeFile(yamlPath);
                 for (const auto& subFile : this->gSubFileList) {
                     freeFile(subFile);
                 }
+                for (const auto& external : externalFilesOf[yamlPath]) {
+                    if (--pendingDependents[external] == 0 &&
+                        (finishedFiles.contains(external) || this->gProcessedFiles.contains(external) ||
+                         this->gParseResults.contains(external))) {
+                        freeFile(external);
+                    }
+                }
             }
         }
     }