Add bk_zip library and BK64 decompression foundation

3 files changed, 111 insertions, 33 deletions

diff --git a/src/utils/Decompressor.cpp b/src/utils/Decompressor.cpp
index 2cf9412..f9c800c 100644
--- a/src/utils/Decompressor.cpp
+++ b/src/utils/Decompressor.cpp
@@ -2,6 +2,7 @@
 #include "TorchUtils.h"
 
 #include <stdexcept>
+#include <mutex>
 #include "spdlog/spdlog.h"
 #include <Companion.h>
 
@@ -12,13 +13,19 @@ extern "C" {
 #include <libmio0/tkmk00.h>
 }
 
+#include <bk_zip/bk_unzip.h>
+
 std::unordered_map<uint32_t, DataChunk*> gCachedChunks;
+std::mutex gDecompCacheMutex;
 
 DataChunk* Decompressor::Decode(const std::vector<uint8_t>& buffer, const uint32_t offset, const CompressionType type,
-                                bool ignoreCache) {
+                                const uint32_t in_size, bool ignoreCache) {
 
-    if (!ignoreCache && Torch::contains(gCachedChunks, offset)) {
-        return gCachedChunks[offset];
+    {
+        std::lock_guard<std::mutex> lock(gDecompCacheMutex);
+        if (!ignoreCache && Torch::contains(gCachedChunks, offset)) {
+            return gCachedChunks[offset];
+        }
     }
 
     const unsigned char* in_buf = buffer.data() + offset;
@@ -32,8 +39,11 @@ DataChunk* Decompressor::Decode(const std::vector<uint8_t>& buffer, const uint32
 
             const auto decompressed = new uint8_t[head.dest_size];
             mio0_decode(in_buf, decompressed, nullptr);
-            gCachedChunks[offset] = new DataChunk{ decompressed, head.dest_size };
-            return gCachedChunks[offset];
+            {
+                std::lock_guard<std::mutex> lock(gDecompCacheMutex);
+                gCachedChunks[offset] = new DataChunk{ decompressed, head.dest_size };
+                return gCachedChunks[offset];
+            }
         }
         case CompressionType::YAY0: {
             uint32_t size = 0;
@@ -43,8 +53,11 @@ DataChunk* Decompressor::Decode(const std::vector<uint8_t>& buffer, const uint32
                 throw std::runtime_error("Failed to decode YAY0");
             }
 
-            gCachedChunks[offset] = new DataChunk{ decompressed, size };
-            return gCachedChunks[offset];
+            {
+                std::lock_guard<std::mutex> lock(gDecompCacheMutex);
+                gCachedChunks[offset] = new DataChunk{ decompressed, size };
+                return gCachedChunks[offset];
+            }
         }
         case CompressionType::YAY1: {
             uint32_t size = 0;
@@ -54,8 +67,30 @@ DataChunk* Decompressor::Decode(const std::vector<uint8_t>& buffer, const uint32
                 throw std::runtime_error("Failed to decode YAY1");
             }
 
-            gCachedChunks[offset] = new DataChunk{ decompressed, size };
-            return gCachedChunks[offset];
+            {
+                std::lock_guard<std::mutex> lock(gDecompCacheMutex);
+                gCachedChunks[offset] = new DataChunk{ decompressed, size };
+                return gCachedChunks[offset];
+            }
+        }
+        case CompressionType::BKZIP: {
+            uint32_t size = in_size;
+            try {
+                uint8_t* decompressed = BK64::bk_unzip(in_buf, &size);
+
+                if (!decompressed) {
+                    throw std::runtime_error("bk_unzip returned null");
+                }
+
+                {
+                    std::lock_guard<std::mutex> lock(gDecompCacheMutex);
+                    gCachedChunks[offset] = new DataChunk{ decompressed, size };
+                    return gCachedChunks[offset];
+                }
+            } catch (const std::exception& e) {
+                throw std::runtime_error(std::string(e.what()) + " (ROM offset 0x" + Torch::to_hex(offset, false) +
+                                         " compressed size 0x" + Torch::to_hex(in_size, false) + ")");
+            }
         }
         default:
             throw std::runtime_error("Unknown compression type");
@@ -64,8 +99,11 @@ DataChunk* Decompressor::Decode(const std::vector<uint8_t>& buffer, const uint32
 
 DataChunk* Decompressor::DecodeTKMK00(const std::vector<uint8_t>& buffer, const uint32_t offset, const uint32_t size,
                                       const uint32_t alpha) {
-    if (Torch::contains(gCachedChunks, offset)) {
-        return gCachedChunks[offset];
+    {
+        std::lock_guard<std::mutex> lock(gDecompCacheMutex);
+        if (Torch::contains(gCachedChunks, offset)) {
+            return gCachedChunks[offset];
+        }
     }
 
     const uint8_t* in_buf = buffer.data() + offset;
@@ -73,8 +111,11 @@ DataChunk* Decompressor::DecodeTKMK00(const std::vector<uint8_t>& buffer, const
     const auto decompressed = new uint8_t[size];
     const auto rgba = new uint8_t[size];
     tkmk00_decode(in_buf, decompressed, rgba, alpha);
-    gCachedChunks[offset] = new DataChunk{ rgba, size };
-    return gCachedChunks[offset];
+    {
+        std::lock_guard<std::mutex> lock(gDecompCacheMutex);
+        gCachedChunks[offset] = new DataChunk{ rgba, size };
+        return gCachedChunks[offset];
+    }
 }
 
 DecompressedData Decompressor::AutoDecode(YAML::Node& node, std::vector<uint8_t>& buffer,
@@ -84,16 +125,14 @@ DecompressedData Decompressor::AutoDecode(YAML::Node& node, std::vector<uint8_t>
     CompressionType type = Companion::Instance->GetCurrCompressionType();
 
     auto fileOffset = TranslateAddr(offset, true);
+    auto offsetFromFile = TranslateAddr(offset, false) - fileOffset;
 
     // Check if an asset in a yaml file is mio0 compressed and extract.
     if (node["mio0"]) {
         auto assetPtr = ASSET_PTR(offset);
         auto gameSize = Companion::Instance->GetRomData().size();
 
-        auto fileOffset = TranslateAddr(offset, true);
-        offset = ASSET_PTR(offset);
-
-        auto decoded = Decode(buffer, fileOffset + offset, CompressionType::MIO0);
+        auto decoded = Decode(buffer, fileOffset + offsetFromFile, CompressionType::MIO0);
         size_t decodedSize = decoded->size - offset;
         size_t size;
 
@@ -144,7 +183,12 @@ DecompressedData Decompressor::AutoDecode(YAML::Node& node, std::vector<uint8_t>
                         size, decodedSize, assetPtr);
             size = decodedSize;
         }
-
+        return { .root = decoded, .segment = { decoded->data, size } };
+    }
+    if (node["bkzip"]) {
+        const auto compressedSize = GetSafeNode<uint32_t>(node, "compressed_size");
+        auto decoded = Decode(buffer, fileOffset + offsetFromFile, CompressionType::BKZIP, compressedSize);
+        auto size = node["size"] ? node["size"].as<size_t>() : manualSize.value_or(decoded->size);
         return { .root = decoded, .segment = { decoded->data, size } };
     }
 
@@ -176,6 +220,20 @@ DecompressedData Decompressor::AutoDecode(YAML::Node& node, std::vector<uint8_t>
 
             return { .root = decoded, .segment = { decoded->data + offset, size } };
         }
+        case CompressionType::BKZIP: {
+            const auto sizeEntry = Companion::Instance->GetCurrentCompressedSize();
+
+            if (!sizeEntry.has_value()) {
+                throw std::runtime_error("Auto decode missing file compressed size.");
+            }
+
+            auto decoded = Decode(buffer, fileOffset, type, sizeEntry.value());
+            auto size = node["size"] ? node["size"].as<size_t>() : manualSize.value_or(decoded->size - offsetFromFile);
+            SPDLOG_INFO("AutoDecode BKZIP: offset=0x{:X} fileOffset=0x{:X} offsetFromFile=0x{:X} compSize=0x{:X} "
+                        "decodedSize=0x{:X} segSize=0x{:X}",
+                        offset, fileOffset, offsetFromFile, sizeEntry.value(), decoded->size, size);
+            return { .root = decoded, .segment = { decoded->data + offsetFromFile, size } };
+        }
         case CompressionType::YAZ0:
             throw std::runtime_error(
                 "Found compressed yaz0 segment.\nDecompression of yaz0 has not been implemented yet.");
@@ -203,10 +261,9 @@ DecompressedData Decompressor::AutoDecode(YAML::Node& node, std::vector<uint8_t>
 
             return { .root = nullptr, .segment = { buffer.data() + fileOffset, size } };
         }
+        default:
+            throw std::runtime_error("Auto decode could not find a supported compression type.");
     }
-
-    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.");
 }
 
 DecompressedData Decompressor::AutoDecode(uint32_t offset, std::optional<size_t> size, std::vector<uint8_t>& buffer) {
@@ -220,9 +277,15 @@ uint32_t Decompressor::TranslateAddr(uint32_t addr, bool baseAddress) {
     if (IS_SEGMENTED(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;
+            const auto compressedSegmentPair =
+                Companion::Instance->GetFileOffsetFromCompressedSegmentedAddr(SEGMENT_NUMBER(addr));
+            if (!compressedSegmentPair.has_value()) {
+                SPDLOG_ERROR("Segment data missing from game config\nPlease add an entry for segment {}",
+                             SEGMENT_NUMBER(addr));
+                return 0;
+            }
+            return compressedSegmentPair.value().first +
+                   (!baseAddress ? (compressedSegmentPair.value().second + SEGMENT_OFFSET(addr)) : 0);
         }
 
         return segment.value() + (!baseAddress ? SEGMENT_OFFSET(addr) : 0);
@@ -234,7 +297,7 @@ uint32_t Decompressor::TranslateAddr(uint32_t addr, bool baseAddress) {
         const auto vram = vramEntry.value();
 
         if (addr >= vram.addr) {
-            return vram.offset + (addr - vram.addr);
+            return vram.offset + (!baseAddress ? (addr - vram.addr) : 0);
         }
     }
 
@@ -273,11 +336,14 @@ bool Decompressor::IsSegmented(uint32_t 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 false;
+            const auto compressedSegmentPair =
+                Companion::Instance->GetFileOffsetFromCompressedSegmentedAddr(SEGMENT_NUMBER(addr));
+            if (!compressedSegmentPair.has_value()) {
+                SPDLOG_ERROR("Segment data missing from game config\nPlease add an entry for segment {}",
+                             SEGMENT_NUMBER(addr));
+                return false;
+            }
         }
-
         return true;
     }
 
@@ -285,6 +351,7 @@ bool Decompressor::IsSegmented(uint32_t addr) {
 }
 
 void Decompressor::ClearCache() {
+    std::lock_guard<std::mutex> lock(gDecompCacheMutex);
     for (auto& [key, value] : gCachedChunks) {
         delete[] value->data;
     }
diff --git a/src/utils/Decompressor.h b/src/utils/Decompressor.h
index 22f4b24..cc07193 100644
--- a/src/utils/Decompressor.h
+++ b/src/utils/Decompressor.h
@@ -14,6 +14,7 @@ enum class CompressionType {
     YAY0,
     YAY1,
     YAZ0,
+    BKZIP,
 };
 
 
@@ -33,7 +34,7 @@ struct DecompressedData {
 
 class Decompressor {
 public:
-    static DataChunk* Decode(const std::vector<uint8_t>& buffer, uint32_t offset, CompressionType type, bool ignoreCache = false);
+    static DataChunk* Decode(const std::vector<uint8_t>& buffer, uint32_t offset, CompressionType type, const uint32_t in_size = 0, bool ignoreCache = false);
     static DataChunk* DecodeTKMK00(const std::vector<uint8_t>& buffer, const uint32_t offset, const uint32_t size, const uint32_t alpha);
     static DecompressedData AutoDecode(YAML::Node& node, std::vector<uint8_t>& buffer, std::optional<size_t> size = std::nullopt);
     static DecompressedData AutoDecode(uint32_t offset, std::optional<size_t> size, std::vector<uint8_t>& buffer);
diff --git a/src/utils/TorchUtils.cpp b/src/utils/TorchUtils.cpp
index 86c836f..ec95c4a 100644
--- a/src/utils/TorchUtils.cpp
+++ b/src/utils/TorchUtils.cpp
@@ -8,12 +8,22 @@
 namespace fs = std::filesystem;
 
 uint32_t Torch::translate(const uint32_t offset) {
-    if (SEGMENT_NUMBER(offset) > 0x01) {
+    // Segment 0 is a raw offset. BK64 also translates segment 1 (it stores assets
+    // in compressed segments); other games leave segment 1 raw.
+    const uint32_t firstSegment =
+        Companion::Instance->GetGBIMinorVersion() == GBIMinorVersion::BK64 ? 0x00 : 0x01;
+    if (SEGMENT_NUMBER(offset) > firstSegment) {
         auto segment = SEGMENT_NUMBER(offset);
         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");
+            // Fall back to the compressed-segment map
+            const auto compressedSegmentPair = Companion::Instance->GetFileOffsetFromCompressedSegmentedAddr(segment);
+
+            if (!compressedSegmentPair.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");
+            }
+            return compressedSegmentPair.value().first + compressedSegmentPair.value().second + SEGMENT_OFFSET(offset);
         }
 
         return addr.value() + SEGMENT_OFFSET(offset);