Detect Compression Type and Improve Yaml Structure (#21)

* Auto detect mio0

* Altered decompressor

* Fix bug

* Bug fix

* fix

* debugging

* fix debug

* Remove logging

* Fix spacing

* node change

* Cleanup

* Track compression offset using global var

* Implementation

* Renames

* Simplification

* Fixes

* Fix debugging output

* yaml config segments uses array style

* Fix to_hex zero fill

* Bug fixes

---------

Co-authored-by: = <=>

10 files changed, 138 insertions, 72 deletions

diff --git a/src/Companion.cpp b/src/Companion.cpp
index 4df7fef..66dfbda 100644
--- a/src/Companion.cpp
+++ b/src/Companion.cpp
@@ -204,11 +204,30 @@ void Companion::ParseModdingConfig() {
 
 void Companion::ParseCurrentFileConfig(YAML::Node node) {
     if(node["segments"]) {
-        for(auto segment = node["segments"].begin(); segment != node["segments"].end(); ++segment) {
-            const auto id = std::stoi(segment->first.as<std::string>().substr(3));
-            const auto replacement = segment->second.as<uint32_t>();
-            this->gConfig.segment.local[id] = replacement;
-            SPDLOG_DEBUG("Segment {} replaced with 0x{:X}", id, replacement);
+        auto segments = node["segments"];
+
+        // Set global variables for segmented data
+        if (segments.IsSequence() && segments.size()) {
+            if (segments[0].IsSequence() && segments[0].size() == 2) {
+                gCurrentSegmentNumber = segments[0][0].as<uint32_t>();
+                gCurrentFileOffset = segments[0][1].as<uint32_t>();
+                gCurrentCompressionType = GetCompressionType(this->gRomData, gCurrentFileOffset);
+            } else {
+                throw std::runtime_error("Incorrect yaml syntax for segments.\n\nThe yaml expects:\n:config:\n  segments:\n  - [<segment>, <offset_of_compressed_file>]\n\nLike so:\nsegments:\n  - [0x06, 0x821D10]");
+            }
+        }
+
+        // Set file offset for later use.
+        for(size_t i = 0; i < segments.size(); i++) {
+            auto segment = segments[i];
+            if (segment.IsSequence() && segment.size() == 2) {
+                const auto id = segment[0].as<uint32_t>();
+                const auto replacement = segment[1].as<uint32_t>();
+                this->gConfig.segment.local[id] = replacement;
+                SPDLOG_DEBUG("Segment {} replaced with 0x{:X}", id, replacement);
+            } else {
+                throw std::runtime_error("Incorrect yaml syntax for segments.\n\nThe yaml expects:\n:config:\n  segments:\n  - [<segment>, <offset_of_compressed_file>]\n\nLike so:\nsegments:\n  - [0x06, 0x821D10]");
+            }
         }
     }
     if(node["header"]) {
@@ -381,6 +400,19 @@ void Companion::Process() {
         this->gCurrentDirectory = relative(entry.path(), path).replace_extension("");
         this->gCurrentFile = yamlPath;
 
+        // Set compressed file offsets and compression type
+        if (auto segments = root[":config"]["segments"]) {
+            if (segments.IsSequence() && segments.size() > 0) {
+                if (segments[0].IsSequence() && segments[0].size() == 2) {
+                    gCurrentSegmentNumber = segments[0][0].as<uint32_t>();
+                    gCurrentFileOffset = segments[0][1].as<uint32_t>();
+                    gCurrentCompressionType = GetCompressionType(this->gRomData, gCurrentFileOffset);
+                } else {
+                    throw std::runtime_error("Incorrect yaml syntax for segments.\n\nThe yaml expects:\n:config:\n  segments:\n  - [<segment>, <offset_of_compressed_file>]\n\nLike so:\nsegments:\n  - [0x06, 0x821D10]");
+                }
+            }
+        }
+
         for(auto asset = root.begin(); asset != root.end(); ++asset){
             auto node = asset->second;
             auto entryName = asset->first.as<std::string>();
@@ -406,16 +438,18 @@ void Companion::Process() {
                         continue;
                     }
 
-                    if(segment != -1) {
+                    if(segment != -1 || gCurrentSegmentNumber) {
                         assetNode["offset"] = (segment << 24) | assetNode["offset"].as<uint32_t>();
                     }
 
                     this->gAddrMap[this->gCurrentFile][assetNode["offset"].as<uint32_t>()] = std::make_tuple(output, assetNode);
                 }
             } else {
+
                 auto output = (this->gCurrentDirectory / entryName).string();
                 std::replace(output.begin(), output.end(), '\\', '/');
 
+
                 if(node["type"]){
                     const auto type = GetSafeNode<std::string>(node, "type");
                     if(type == "SAMPLE"){
@@ -427,6 +461,12 @@ void Companion::Process() {
                     continue;
                 }
 
+                if(gCurrentSegmentNumber) {
+                    if (IS_SEGMENTED(node["offset"].as<uint32_t>()) == false) {
+                        node["offset"] = (gCurrentSegmentNumber << 24) | node["offset"].as<uint32_t>();
+                    }
+                }
+
                 this->gAddrMap[this->gCurrentFile][node["offset"].as<uint32_t>()] = std::make_tuple(output, node);
             }
         }
@@ -454,6 +494,7 @@ void Companion::Process() {
                 continue;
             }
 
+
             // Parse horizontal assets
             if(assetNode["files"]){
                 auto segment = assetNode["segment"] ? assetNode["segment"].as<uint8_t>() : -1;
@@ -466,7 +507,7 @@ void Companion::Process() {
                         continue;
                     }
 
-                    if(segment != -1) {
+                    if(segment != -1 || gCurrentFileOffset) {
                         node["offset"] = (segment << 24) | node["offset"].as<uint32_t>();
                     }
 
@@ -476,6 +517,11 @@ void Companion::Process() {
                     this->ExtractNode(node, output, wrapper);
                 }
             } else {
+                if(gCurrentFileOffset) {
+                    if (IS_SEGMENTED(assetNode["offset"].as<uint32_t>()) == false) {
+                        assetNode["offset"] = (gCurrentSegmentNumber << 24) | assetNode["offset"].as<uint32_t>();
+                    }
+                }
                 std::string output = (this->gCurrentDirectory / entryName).string();
                 std::replace(output.begin(), output.end(), '\\', '/');
                 this->gConfig.segment.temporal.clear();
@@ -677,7 +723,30 @@ std::optional<std::shared_ptr<BaseFactory>> Companion::GetFactory(const std::str
     return this->gFactories[type];
 }
 
-std::optional<std::uint32_t> Companion::GetSegmentedAddr(const uint8_t segment) const {
+/**
+ * @param offset Rom offset of compressed mio0 file.
+ * @returns CompressionType
+ */
+CompressionType Companion::GetCompressionType(std::vector<uint8_t>& buffer, const uint32_t offset) {
+    if (offset) {
+        LUS::BinaryReader reader((char*) buffer.data() + offset, sizeof(uint32_t));
+        reader.SetEndianness(LUS::Endianness::Big);
+
+        const std::string header = reader.ReadCString();
+
+        // Check if a compressed header exists
+        if (header == "MIO0") {
+            return CompressionType::MIO0;
+        } else if (header == "YAY0") {
+            return CompressionType::YAY0;
+        } else if (header == "YAZ0") {
+            return CompressionType::YAZ0;
+        }
+    }
+    return CompressionType::None;
+}
+
+std::optional<std::uint32_t> Companion::GetFileOffsetFromSegmentedAddr(const uint8_t segment) const {
 
     auto segments = this->gConfig.segment;
 
diff --git a/src/Companion.h b/src/Companion.h
index b82b83d..fbde351 100644
--- a/src/Companion.h
+++ b/src/Companion.h
@@ -9,6 +9,7 @@
 #include <variant>
 #include "factories/BaseFactory.h"
 #include "n64/Cartridge.h"
+#include "utils/Decompressor.h"
 
 class SWrapper;
 namespace fs = std::filesystem;
@@ -72,10 +73,15 @@ public:
     GBIVersion GetGBIVersion() const { return this->gConfig.gbi.version; }
     GBIMinorVersion GetGBIMinorVersion() const { return  this->gConfig.gbi.subversion; }
 
-    std::optional<std::uint32_t> GetSegmentedAddr(uint8_t segment) const;
+    std::optional<std::uint32_t> GetFileOffsetFromSegmentedAddr(uint8_t segment) const;
     std::optional<std::tuple<std::string, YAML::Node>> GetNodeByAddr(uint32_t addr);
     std::optional<std::shared_ptr<BaseFactory>> GetFactory(const std::string& type);
 
+    std::optional<std::uint32_t> GetFileOffset(void) const { return this->gCurrentFileOffset; };
+    std::optional<std::uint32_t> GetCurrSegmentNumber(void) const { return this->gCurrentSegmentNumber; };
+    CompressionType GetCurrCompressionType(void) const { return this->gCurrentCompressionType; };
+    CompressionType GetCompressionType(std::vector<uint8_t>& buffer, const uint32_t offset);
+
     static std::string CalculateHash(const std::vector<uint8_t>& data);
     static void Pack(const std::string& folder, const std::string& output);
     std::string NormalizeAsset(const std::string& name) const;
@@ -101,6 +107,9 @@ private:
     std::unordered_map<std::string, std::map<std::string, std::pair<YAML::Node, bool>>> gAssetDependencies;
     std::unordered_map<std::string, std::map<std::string, std::vector<std::pair<uint32_t, std::string>>>> gWriteMap;
     std::unordered_map<std::string, std::unordered_map<uint32_t, std::tuple<std::string, YAML::Node>>> gAddrMap;
+    uint32_t gCurrentFileOffset;
+    uint32_t gCurrentSegmentNumber;
+    CompressionType gCurrentCompressionType;
 
     void ParseCurrentFileConfig(YAML::Node node);
     void ParseModdingConfig();
diff --git a/src/factories/DisplayListFactory.cpp b/src/factories/DisplayListFactory.cpp
index bf13826..d5197cc 100644
--- a/src/factories/DisplayListFactory.cpp
+++ b/src/factories/DisplayListFactory.cpp
@@ -77,7 +77,7 @@ void DListHeaderExporter::Export(std::ostream &write, std::shared_ptr<IParsedDat
 void DListCodeExporter::Export(std::ostream &write, std::shared_ptr<IParsedData> raw, std::string& entryName, YAML::Node &node, std::string* replacement ) {
     const auto cmds = std::static_pointer_cast<DListData>(raw)->mGfxs;
     const auto symbol = GetSafeNode(node, "symbol", entryName);
-    const auto offset = GetSafeNode<uint32_t>(node, "offset");
+    auto offset = GetSafeNode<uint32_t>(node, "offset");
 
     char out[0xFFFF] = {0};
 
@@ -113,6 +113,9 @@ void DListCodeExporter::Export(std::ostream &write, std::shared_ptr<IParsedData>
     GFXDSetGBIVersion();
 
     if (Companion::Instance->IsDebug()) {
+        if (IS_SEGMENTED(offset)) {
+            offset = SEGMENT_OFFSET(offset);
+        }
         write << "// 0x" << std::hex << std::uppercase << offset << "\n";
     }
 
@@ -267,8 +270,6 @@ std::optional<std::shared_ptr<IParsedData>> DListFactory::parse(std::vector<uint
     auto [_, segment] = Decompressor::AutoDecode(node, raw_buffer);
     LUS::BinaryReader reader(segment.data, segment.size);
     reader.SetEndianness(LUS::Endianness::Big);
-    // Validate this
-    // reader.Seek(offset, LUS::SeekOffsetType::Start);
 
     std::vector<uint32_t> gfxs;
     auto processing = true;
@@ -305,7 +306,6 @@ std::optional<std::shared_ptr<IParsedData>> DListFactory::parse(std::vector<uint
                     output = Companion::Instance->NormalizeAsset("dl_" + Torch::to_hex(ptr, false));
                 }
 
-                Decompressor::CopyCompression(node, dl);
                 dl["type"] = "GFX";
                 dl["offset"] = ptr;
                 dl["symbol"] = output;
@@ -363,7 +363,6 @@ std::optional<std::shared_ptr<IParsedData>> DListFactory::parse(std::vector<uint
                     output = Companion::Instance->NormalizeAsset("lights1_" + Torch::to_hex(w1, false));
                 }
 
-                Decompressor::CopyCompression(node, light);
                 light["type"] = "lights";
                 light["offset"] = ptr;
                 light["symbol"] = output;
@@ -411,7 +410,6 @@ std::optional<std::shared_ptr<IParsedData>> DListFactory::parse(std::vector<uint
                     output = Companion::Instance->NormalizeAsset("vtx_" + Torch::to_hex(w1, false));
                 }
 
-                Decompressor::CopyCompression(node, vtx);
                 vtx["type"] = "VTX";
                 vtx["offset"] = ptr;
                 vtx["count"] = nvtx;
diff --git a/src/factories/TextureFactory.cpp b/src/factories/TextureFactory.cpp
index 6aa4f82..a9d0c44 100644
--- a/src/factories/TextureFactory.cpp
+++ b/src/factories/TextureFactory.cpp
@@ -261,7 +261,6 @@ std::optional<std::shared_ptr<IParsedData>> TextureFactory::parse(std::vector<ui
     }
     SPDLOG_INFO("Size: {}", size);
     SPDLOG_INFO("Offset: 0x{:X}", offset);
-    SPDLOG_INFO("Is Compressed: {}", Decompressor::IsCompressed(node) ? "true" : "false");
 
     if(result.size() == 0){
         return std::nullopt;
@@ -360,7 +359,6 @@ std::optional<std::shared_ptr<IParsedData>> TextureFactory::parse_modding(std::v
     }
     SPDLOG_INFO("Size: {}", size);
     SPDLOG_INFO("Offset: 0x{:X}", offset);
-    SPDLOG_INFO("Is Compressed: {}", Decompressor::IsCompressed(node) ? "true" : "false");
 
     if(result.size() == 0){
         return std::nullopt;
diff --git a/src/factories/VtxFactory.cpp b/src/factories/VtxFactory.cpp
index 7b74681..d506a39 100644
--- a/src/factories/VtxFactory.cpp
+++ b/src/factories/VtxFactory.cpp
@@ -1,4 +1,5 @@
 #include "VtxFactory.h"
+#include "spdlog/spdlog.h"
 
 #include "Companion.h"
 #include "utils/Decompressor.h"
@@ -20,7 +21,11 @@ void VtxHeaderExporter::Export(std::ostream &write, std::shared_ptr<IParsedData>
 void VtxCodeExporter::Export(std::ostream &write, std::shared_ptr<IParsedData> raw, std::string& entryName, YAML::Node &node, std::string* replacement ) {
     auto vtx = std::static_pointer_cast<VtxData>(raw)->mVtxs;
     const auto symbol = GetSafeNode(node, "symbol", entryName);
-    const auto offset = GetSafeNode<uint32_t>(node, "offset");
+    auto offset = GetSafeNode<uint32_t>(node, "offset");
+
+    if (IS_SEGMENTED(offset)) {
+        offset = SEGMENT_OFFSET(offset);
+    }
 
     if (Companion::Instance->IsDebug()) {
         write << "// 0x" << std::hex << std::uppercase << offset << "\n";
diff --git a/src/factories/sm64/GeoLayoutFactory.cpp b/src/factories/sm64/GeoLayoutFactory.cpp
index 844b8fb..f0d4788 100644
--- a/src/factories/sm64/GeoLayoutFactory.cpp
+++ b/src/factories/sm64/GeoLayoutFactory.cpp
@@ -16,7 +16,7 @@ uint64_t RegisterOrLoadGeoLayout(uint32_t ptr) {
     }
 
     SPDLOG_INFO("Addr to Geo Layout at 0x{:X} not in yaml, autogenerating it", ptr);
-    auto addr = Companion::Instance->GetSegmentedAddr(SEGMENT_NUMBER(ptr));
+    auto addr = Companion::Instance->GetFileOffsetFromSegmentedAddr(SEGMENT_NUMBER(ptr));
     if(!addr.has_value()) {
         SPDLOG_ERROR("Segment data missing from game config\nPlease add an entry for segment {}", SEGMENT_NUMBER(ptr));
         return 0;
@@ -56,7 +56,7 @@ uint64_t RegisterOrLoadDisplayList(uint32_t ptr) {
     }
 
     SPDLOG_INFO("Addr to Display List at 0x{:X} not in yaml, autogenerating it", ptr);
-    auto addr = Companion::Instance->GetSegmentedAddr(SEGMENT_NUMBER(ptr));
+    auto addr = Companion::Instance->GetFileOffsetFromSegmentedAddr(SEGMENT_NUMBER(ptr));
     if(!addr.has_value()) {
         SPDLOG_ERROR("Segment data missing from game config\nPlease add an entry for segment {}", SEGMENT_NUMBER(ptr));
         return 0;
diff --git a/src/utils/Decompressor.cpp b/src/utils/Decompressor.cpp
index b4ca7a4..6d41910 100644
--- a/src/utils/Decompressor.cpp
+++ b/src/utils/Decompressor.cpp
@@ -36,46 +36,53 @@ DataChunk* Decompressor::Decode(const std::vector<uint8_t>& buffer, const uint32
 }
 
 DecompressedData Decompressor::AutoDecode(YAML::Node& node, std::vector<uint8_t>& buffer, std::optional<size_t> manualSize) {
+
     if(!node["offset"]){
         throw std::runtime_error("Failed to find offset");
     }
 
-    auto offset = node["offset"].as<uint32_t>();    
-
-    if(node["mio0"]){
-        const auto mio0 = node["mio0"].as<uint32_t>();
-
-        offset = TranslateAddr(offset, IS_SEGMENTED(offset));        
+    auto offset = node["offset"].as<uint32_t>();
 
-        auto decoded = Decode(buffer, offset, CompressionType::MIO0);
-        auto size = node["size"] ? node["size"].as<size_t>() : manualSize.value_or(decoded->size - mio0);
-
-        return {
-            .root = decoded,
-            .segment = { decoded->data + mio0, size }
-        };
-    }
+    CompressionType type = Companion::Instance->GetCurrCompressionType();
 
-    if(node["yay0"]){
-        throw std::runtime_error("Yay0 not implemented");
-    }
+    switch(type) {
+        case CompressionType::MIO0:
+        {
+            auto fileOffset = TranslateAddr(offset, true);
+            offset = IS_SEGMENTED(offset) ? SEGMENT_OFFSET(offset) : offset;
 
-    if(node["yaz0"]){
-        throw std::runtime_error("Yaz0 not implemented");
+            auto decoded = Decode(buffer, fileOffset, CompressionType::MIO0);
+            auto size = node["size"] ? node["size"].as<size_t>() : manualSize.value_or(decoded->size - offset);
+            return {
+                .root = decoded,
+                .segment = { decoded->data + offset, size }
+            };
+        }
+        case CompressionType::YAY0:
+            throw std::runtime_error("Found compressed yay0 segment.\nDecompression of yay0 has not been implemented yet.");
+            break;
+        case CompressionType::YAZ0:
+            throw std::runtime_error("Found compressed yaz0 segment.\nDecompression of yaz0 has not been implemented yet.");
+            break;
+        case CompressionType::None:
+        {
+            auto fileOffset = TranslateAddr(offset);
+
+            auto size = node["size"] ? node["size"].as<size_t>() : manualSize.value_or(buffer.size() - fileOffset);
+
+            return {
+                .root = nullptr,
+                .segment = { buffer.data() + fileOffset, size }
+            };
+        }
     }
 
-    offset = TranslateAddr(offset);
-
-    return {
-        .root = nullptr,
-        .segment = { buffer.data() + offset, node["size"] ? node["size"].as<size_t>() : manualSize.value_or(buffer.size() - offset) }
-    };
+    throw std::runtime_error("Auto decode could not find a compression type nor uncompressed segment.\nThis is one of those issues that should never really happen.");
 }
 
 uint32_t Decompressor::TranslateAddr(uint32_t addr, bool baseAddress){
     if(IS_SEGMENTED(addr)){
-        const auto segment = Companion::Instance->GetSegmentedAddr(SEGMENT_NUMBER(addr));
-
+        const auto segment = Companion::Instance->GetFileOffsetFromSegmentedAddr(SEGMENT_NUMBER(addr));
         if(!segment.has_value()) {
             SPDLOG_ERROR("Segment data missing from game config\nPlease add an entry for segment {}", SEGMENT_NUMBER(addr));
             return 0;
@@ -89,7 +96,7 @@ uint32_t Decompressor::TranslateAddr(uint32_t addr, bool baseAddress){
 
 bool Decompressor::IsSegmented(uint32_t addr) {
     if(IS_SEGMENTED(addr)){
-        const auto segment = Companion::Instance->GetSegmentedAddr(SEGMENT_NUMBER(addr));
+        const auto segment = Companion::Instance->GetFileOffsetFromSegmentedAddr(SEGMENT_NUMBER(addr));
 
         if(!segment.has_value()) {
             SPDLOG_ERROR("Segment data missing from game config\nPlease add an entry for segment {}", SEGMENT_NUMBER(addr));
@@ -102,24 +109,6 @@ bool Decompressor::IsSegmented(uint32_t addr) {
     return false;
 }
 
-bool Decompressor::IsCompressed(YAML::Node& node) {
-    return node["mio0"] || node["yay0"] || node["yaz0"];
-}
-
-void Decompressor::CopyCompression(YAML::Node& from, YAML::Node& to){
-    if(from["mio0"]){
-        to["mio0"] = from["mio0"];
-    }
-
-    if(from["yay0"]){
-        to["yay0"] = from["yay0"];
-    }
-
-    if(from["yaz0"]){
-        to["yaz0"] = from["yaz0"];
-    }
-}
-
 void Decompressor::ClearCache() {
     for(auto& [key, value] : gCachedChunks){
         delete value->data;
diff --git a/src/utils/Decompressor.h b/src/utils/Decompressor.h
index 718b726..4e84cde 100644
--- a/src/utils/Decompressor.h
+++ b/src/utils/Decompressor.h
@@ -9,8 +9,8 @@
 
 enum class CompressionType {
     MIO0,
-    Yay0,
-    Yaz0,
+    YAY0,
+    YAZ0,
     None,
 };
 
@@ -31,8 +31,6 @@ public:
     static DecompressedData AutoDecode(YAML::Node& node, std::vector<uint8_t>& buffer, std::optional<size_t> size = std::nullopt);
     static uint32_t TranslateAddr(uint32_t addr, bool baseAddress = false);
     static bool IsSegmented(uint32_t addr);
-    static bool IsCompressed(YAML::Node& node);
-    static void CopyCompression(YAML::Node& from, YAML::Node& to);
 
     static void ClearCache();
 };
\ No newline at end of file
diff --git a/src/utils/TorchUtils.cpp b/src/utils/TorchUtils.cpp
index d63fa0b..5dc4621 100644
--- a/src/utils/TorchUtils.cpp
+++ b/src/utils/TorchUtils.cpp
@@ -7,7 +7,7 @@
 uint32_t Torch::translate(const uint32_t offset) {
     if(SEGMENT_NUMBER(offset) > 0x01) {
         auto segment = SEGMENT_NUMBER(offset);
-        const auto addr = Companion::Instance->GetSegmentedAddr(segment);
+        const auto addr = Companion::Instance->GetFileOffsetFromSegmentedAddr(segment);
         if(!addr.has_value()) {
             SPDLOG_ERROR("Segment data missing from game config\nPlease add an entry for segment {}", segment);
             throw std::runtime_error("Failed to find offset");
diff --git a/src/utils/TorchUtils.h b/src/utils/TorchUtils.h
index 36ce46d..271ab98 100644
--- a/src/utils/TorchUtils.h
+++ b/src/utils/TorchUtils.h
@@ -14,7 +14,7 @@ std::string to_hex(T number, const bool append0x = true) {
     if(append0x) {
         stream << "0x";
     }
-    stream << std::setfill ('0') << std::setw(sizeof(T)*2)
+    stream << std::setfill ('0')
            << std::hex << number;
 
     auto format = stream.str();