Asset Pre-counting (#212)

* Add atomic_ref versions of functions to allow asset count passback.
Add process variable to gate full processing (for pre-counting assets).

* Bump C++ standard back to 20 for atomic_ref.

2 files changed, 35 insertions, 12 deletions

diff --git a/src/Companion.cpp b/src/Companion.cpp
index 2ccbdd9..911e18e 100644
--- a/src/Companion.cpp
+++ b/src/Companion.cpp
@@ -130,6 +130,11 @@ static std::string GetTypeNode(YAML::Node& node) {
 Companion* Companion::Instance;
 
 void Companion::Init(const ExportType type) {
+    size_t assetCount = 0;
+    Init(type, std::atomic_ref<size_t>(assetCount));
+}
+
+void Companion::Init(const ExportType type, std::atomic_ref<size_t> assetCount) {
 
     spdlog::set_level(spdlog::level::debug);
     spdlog::set_pattern("[%Y-%m-%d %H:%M:%S.%e] [%l] %v");
@@ -224,7 +229,7 @@ void Companion::Init(const ExportType type) {
     this->RegisterFactory("NAUDIO:V1:SEQUENCE", std::make_shared<NSequenceFactory>());
 #endif
 #ifndef __EMSCRIPTEN__ // We call this manually
-    this->Process();
+    this->Process(assetCount);
 #endif
 }
 
@@ -356,7 +361,7 @@ void Companion::ParseModdingConfig() {
 }
 
 
-void Companion::ParseCurrentFileConfig(YAML::Node node) {
+void Companion::ParseCurrentFileConfig(YAML::Node node, std::atomic_ref<size_t> assetCount) {
     if (node["external_files"]) {
         auto externalFiles = node["external_files"];
         if (externalFiles.IsSequence() && externalFiles.size()) {
@@ -388,8 +393,10 @@ void Companion::ParseCurrentFileConfig(YAML::Node node) {
                     YAML::Node root = YAML::LoadFile(externalFileName);
 
                     if (!Torch::contains(this->gProcessedFiles, this->gCurrentFile)) {
-                        ProcessFile(root);
-                        this->gProcessedFiles.insert(this->gCurrentFile);
+                        ProcessFile(root, assetCount);
+                        if (process) {
+                            this->gProcessedFiles.insert(this->gCurrentFile);
+                        }
                     }
 
                     SPDLOG_INFO("Finishing processing of file: {}", currentFile);
@@ -610,6 +617,12 @@ void Companion::ProcessTables(YAML::Node& rom) {
 }
 
 void Companion::ProcessFile(YAML::Node root) {
+    size_t assetCount = 0;
+    ProcessFile(root, std::atomic_ref<size_t>(assetCount));
+}
+
+void Companion::ProcessFile(YAML::Node root, std::atomic_ref<size_t> assetCount) {
+    assetCount++;
     // Set compressed file offsets and compression type
     if (auto segments = root[":config"]["segments"]) {
         if (segments.IsSequence() && segments.size() > 0) {
@@ -672,10 +685,10 @@ void Companion::ProcessFile(YAML::Node root) {
     GFXDOverride::ClearVtx();
 
     if(root[":config"]) {
-        this->ParseCurrentFileConfig(root[":config"]);
+        this->ParseCurrentFileConfig(root[":config"], assetCount);
     }
 
-    if(!this->NodeHasChanges(this->gCurrentFile) && !this->gNodeForceProcessing) {
+    if(!process || (!this->NodeHasChanges(this->gCurrentFile) && !this->gNodeForceProcessing)) {
         return;
     }
 
@@ -1018,8 +1031,7 @@ void Companion::ProcessFile(YAML::Node root) {
     }
 }
 
-void Companion::Process() {
-
+void Companion::Process(std::atomic_ref<size_t> assetCount) {
     auto configPath = this->gSourceDirectory / "config.yml";
 
     if(!fs::exists(configPath)) {
@@ -1315,8 +1327,10 @@ void Companion::Process() {
         this->gCurrentFile = yamlPath;
 
         if (!Torch::contains(this->gProcessedFiles, this->gCurrentFile)) {
-            ProcessFile(root);
-            this->gProcessedFiles.insert(this->gCurrentFile);
+            ProcessFile(root, assetCount);
+            if (process) {
+                this->gProcessedFiles.insert(this->gCurrentFile);
+            }
         }
     }
 
@@ -1672,6 +1686,10 @@ std::optional<std::vector<std::tuple<std::string, YAML::Node>>> Companion::GetNo
 
 }
 
+void Companion::SetProcess(bool shouldProcess) {
+    this->process = shouldProcess;
+}
+
 void Companion::RegisterCompanionFile(const std::string path, std::vector<char> data) {
     this->gCompanionFiles[path] = data;
     SPDLOG_TRACE("Registered companion file {}", path);
diff --git a/src/Companion.h b/src/Companion.h
index 37b1508..8161ddb 100644
--- a/src/Companion.h
+++ b/src/Companion.h
@@ -4,6 +4,7 @@
 #include <optional>
 #include <filesystem>
 #include <vector>
+#include <atomic>
 #include <fstream>
 #include <unordered_map>
 #include <unordered_set>
@@ -139,10 +140,11 @@ public:
                        Companion(rom, otr, debug, false, srcDir, destPath) {}
 
     void Init(ExportType type);
+    void Init(ExportType type, std::atomic_ref<size_t> assetCount);
 
     bool NodeHasChanges(const std::string& string);
 
-    void Process();
+    void Process(std::atomic_ref<size_t> assetCount);
 
     bool IsOTRMode() const { return (this->gConfig.otrMode != ArchiveType::None); }
     bool IsDebug() const { return this->gConfig.debug; }
@@ -189,6 +191,7 @@ public:
     void SetAdditionalFiles(const std::vector<std::string>& files) { this->gAdditionalFiles = files; }
     void SetVersion(const std::string& version) { this->gVersion = version; }
 
+    void SetProcess(bool shouldProcess);
     TorchConfig& GetConfig() { return this->gConfig; }
     BinaryWrapper* GetCurrentWrapper() { return this->gCurrentWrapper; }
 
@@ -218,6 +221,7 @@ private:
     std::string gCurrentVirtualPath;
     std::string gFileHeader;
     bool gEnablePadGen = false;
+    bool process = true;
     uint32_t gCurrentPad = 0;
     uint32_t gCurrentFileOffset;
     uint32_t gCurrentSegmentNumber;
@@ -240,10 +244,11 @@ private:
     std::unordered_map<std::string, std::unordered_map<uint32_t, std::tuple<std::string, YAML::Node>>> gAddrMap;
 
     void ProcessFile(YAML::Node root);
+    void ProcessFile(YAML::Node root, std::atomic_ref<size_t> assetCount);
     void ParseEnums(std::string& file);
     void ParseHash();
     void ParseModdingConfig();
-    void ParseCurrentFileConfig(YAML::Node node);
+    void ParseCurrentFileConfig(YAML::Node node, std::atomic_ref<size_t> assetCount);
     void RegisterFactory(const std::string& type, const std::shared_ptr<BaseFactory>& factory);
     void ExtractNode(YAML::Node& node, std::string& name, BinaryWrapper* binary);
     void ProcessTables(YAML::Node& rom);