Add option to extract 1 yml (#237)

3 files changed, 40 insertions, 10 deletions

diff --git a/src/Companion.cpp b/src/Companion.cpp
index bb58dcf..ee244a0 100644
--- a/src/Companion.cpp
+++ b/src/Companion.cpp
@@ -903,7 +903,6 @@ void Companion::ProcessExportFile() {
                             AliasManager::Instance->WriteAliases(result.name, this->gCurrentWrapper, dataVec);
                         }
 
-
                         for (auto& entry : this->gCompanionFiles) {
                             auto output = (this->gCurrentDirectory / entry.first).string();
                             std::replace(output.begin(), output.end(), '\\', '/');
@@ -1297,6 +1296,17 @@ void Companion::ProcessFile(YAML::Node root, std::atomic<size_t>& assetCount) {
     ProcessExportFile();
 }
 
+std::vector<fs::directory_entry> Companion::GetAssetYMLs(YAML::Node& rom) const {
+    if (!this->gSingleYMLPath.empty()) {
+        std::vector<fs::directory_entry> single;
+        fs::path assetPath = this->gAssetPath;
+        fs::directory_entry a(assetPath / this->gSingleYMLPath);
+        single.emplace_back(a);
+        return single;
+    }
+    return Torch::getRecursiveEntries(this->gAssetPath);
+}
+
 void Companion::Process(std::atomic<size_t>& assetCount) {
     this->gAssetCounter = &assetCount;
     // gAssetTotal is set by the caller (port-side) if progress reporting is needed
@@ -1623,7 +1633,9 @@ void Companion::Process(std::atomic<size_t>& assetCount) {
         vWriter.Write((uint32_t)0);
     }
 
-    for (const auto& entry : Torch::getRecursiveEntries(this->gAssetPath)) {
+    std::vector<fs::directory_entry> entries = GetAssetYMLs(rom);
+
+    for (const auto& entry : entries) {
         if (entry.is_directory()) {
             continue;
         }
@@ -2070,8 +2082,8 @@ uint32_t Companion::PatchVirtualAddr(uint32_t addr) {
 // then scanning all entries in the file for one that resolves to the same address.
 // Used when multiple segments map to the same ROM data, so the same asset may be
 // registered under a different segment number than the one being looked up.
-std::optional<std::tuple<std::string, YAML::Node>> Companion::FindNodeInOverlaySegments(
-        uint32_t addr, const std::string& file) {
+std::optional<std::tuple<std::string, YAML::Node>> Companion::FindNodeInOverlaySegments(uint32_t addr,
+                                                                                        const std::string& file) {
     // Only applies to files with virtual address mappings (overlays).
     if (!Torch::contains(gVirtualAddrMap, file)) {
         return std::nullopt;
@@ -2100,8 +2112,8 @@ std::optional<std::tuple<std::string, YAML::Node>> Companion::FindNodeInOverlayS
 // Check if this address is a VRAM pointer belonging to the given external file.
 // If it falls within the file's VRAM range, convert to a virtual segment address
 // and look it up in the file's gAddrMap.
-std::optional<std::tuple<std::string, YAML::Node>> Companion::FindInExternalByVRAM(
-        uint32_t addr, const std::string& file) {
+std::optional<std::tuple<std::string, YAML::Node>> Companion::FindInExternalByVRAM(uint32_t addr,
+                                                                                   const std::string& file) {
     // Can't resolve if the external file doesn't have a virtual address mapping.
     if (!Torch::contains(gVirtualAddrMap, file)) {
         return std::nullopt;
@@ -2506,10 +2518,8 @@ std::optional<YAML::Node> Companion::AddAsset(YAML::Node asset) {
     // With strict_declarations, every asset must be pre-declared in the YAML.
     // Throw if an undeclared asset is encountered to catch declaration gaps.
     if (!decl.has_value() && this->gConfig.strictDeclarations) {
-        throw std::runtime_error(
-            "AddAsset: undeclared " + type + " at " + Torch::to_hex(offset, false) +
-            " (symbol: " + symbol + ") in " + this->gCurrentFile +
-            " — YAML declarations incomplete");
+        throw std::runtime_error("AddAsset: undeclared " + type + " at " + Torch::to_hex(offset, false) + " (symbol: " +
+                                 symbol + ") in " + this->gCurrentFile + " — YAML declarations incomplete");
     }
 
     if (decl.has_value()) {
diff --git a/src/Companion.h b/src/Companion.h
index 5bdfb49..b9466ec 100644
--- a/src/Companion.h
+++ b/src/Companion.h
@@ -252,6 +252,8 @@ public:
     void SetCompressedSegment(uint32_t segmentId, uint32_t compressedFileOffset, uint32_t offset);
     bool GetCompressedSegmentOffset(uint32_t* addr);
 
+    void SetSingleYMLPath(const std::string& path) { this->gSingleYMLPath = path; }
+
 #ifdef BUILD_UI
     void RegisterUIFactory(const std::string& type, const std::shared_ptr<BaseFactoryUI>& factory);
     std::optional<std::shared_ptr<BaseFactoryUI>> GetUIFactory(const std::string& type);
@@ -266,6 +268,7 @@ private:
     std::string gCurrentHash;
     std::string gAssetPath;
     std::string gVersion;
+    std::string gSingleYMLPath;
     std::vector<uint8_t> gRomData;
     std::optional<std::filesystem::path> gRomPath;
     bool gNodeForceProcessing = false;
@@ -336,5 +339,6 @@ private:
     void ExtractNode(YAML::Node& node, std::string& name, BinaryWrapper* binary);
     void ProcessTables(YAML::Node& rom);
     void LoadYAMLRecursively(const std::string &dirPath, std::vector<YAML::Node> &result, bool skipRoot);
+    std::vector<fs::directory_entry> GetAssetYMLs(YAML::Node& rom) const;
     std::optional<ParseResultData> ParseNode(YAML::Node& node, std::string& name);
 };
diff --git a/src/main.cpp b/src/main.cpp
index 4145e93..561acca 100644
--- a/src/main.cpp
+++ b/src/main.cpp
@@ -105,6 +105,7 @@ int main(int argc, char* argv[]) {
     std::string folder;
     std::string archive;
     std::string version;
+    std::string singlePath;
     ArchiveType otrMode = ArchiveType::None;
     bool otrModeSelected = false;
     bool xmlMode = false;
@@ -124,9 +125,14 @@ int main(int argc, char* argv[]) {
                     "Set source directory to locate config.yml and asset metadata for processing")
         ->check(CLI::ExistingDirectory);
     otr->add_option("-d,--destdir", destdir, "Set destination directory for export");
+    otr->add_option("-S,--single", singlePath,
+                    "Extracts 1 yml rather than everything relative to the source directory");
 
     otr->parse_complete_callback([&] {
         const auto instance = Companion::Instance = new Companion(filename, ArchiveType::OTR, debug, srcdir, destdir);
+        if (!singlePath.empty()) {
+            instance->SetSingleYMLPath(singlePath);
+        }
         instance->Init(ExportType::Binary);
     });
 
@@ -143,11 +149,16 @@ int main(int argc, char* argv[]) {
                     "Additional files to include in the o2r archive (e.g., mods.toml)")
         ->check(CLI::ExistingFile);
     o2r->add_option("-u,--version", version, "Version to set in the o2r archive");
+    o2r->add_option("-S,--single", singlePath,
+                    "Extracts 1 yml rather than everything relative to the source directory");
 
     o2r->parse_complete_callback([&] {
         const auto instance = Companion::Instance = new Companion(filename, ArchiveType::O2R, debug, srcdir, destdir);
         instance->SetAdditionalFiles(additionalFiles);
         instance->SetVersion(version);
+        if (!singlePath.empty()) {
+            instance->SetSingleYMLPath(singlePath);
+        }
         instance->Init(ExportType::Binary);
     });
 
@@ -160,9 +171,14 @@ int main(int argc, char* argv[]) {
                      "Set source directory to locate config.yml and asset metadata for processing")
         ->check(CLI::ExistingDirectory);
     code->add_option("-d,--destdir", destdir, "Set destination directory to place C code to");
+    code->add_option("-S,--single", singlePath,
+                     "Extracts 1 yml rather than everything relative to the source directory");
 
     code->parse_complete_callback([&]() {
         const auto instance = Companion::Instance = new Companion(filename, ArchiveType::None, debug, srcdir, destdir);
+        if (!singlePath.empty()) {
+            instance->SetSingleYMLPath(singlePath);
+        }
         instance->Init(ExportType::Code);
     });