Implemented rom preprocessing system and comptool support

5 files changed, 147 insertions, 5 deletions

diff --git a/src/Companion.cpp b/src/Companion.cpp
index 3fca4e3..199d649 100644
--- a/src/Companion.cpp
+++ b/src/Companion.cpp
@@ -79,6 +79,8 @@
 #include "factories/naudio/v1/BookFactory.h"
 #include "factories/naudio/v1/SequenceFactory.h"
 
+#include "preprocess/CompTool.h"
+
 using namespace std::chrono;
 namespace fs = std::filesystem;
 
@@ -1013,8 +1015,44 @@ void Companion::Process() {
     }
 
     auto rom = !isDirectoryMode ? config[this->gCartridge->GetHash()] : config;
-    auto cfg = rom["config"];
 
+    if(rom["preprocess"]) {
+        auto preprocess = rom["preprocess"];
+        for(auto job = preprocess.begin(); job != preprocess.end(); job++) {
+            auto name = job->first.as<std::string>();
+            auto item = job->second;
+            auto method = GetSafeNode<std::string>(item, "method");
+            if (method == "mio0-comptool") {
+                auto type = GetSafeNode<std::string>(item, "type");
+                auto target = GetSafeNode<std::string>(item, "target");
+                auto restart = GetSafeNode<bool>(item, "restart");
+
+                if (type == "decompress") {
+                    this->gRomData = CompTool::Decompress(this->gRomData);
+                    this->gCartridge = std::make_shared<N64::Cartridge>(this->gRomData);
+                    this->gCartridge->Initialize();
+
+                    auto hash = this->gCartridge->GetHash();
+
+                    SPDLOG_INFO("ROM decompressed to {}", hash);
+
+                    if (hash != target) {
+                        throw std::runtime_error("Hash mismatch");
+                    }
+
+                    if(restart){
+                        rom = config[this->gCartridge->GetHash()];
+                    }
+                } else {
+                    throw std::runtime_error("Only decompression is supported");
+                }
+            } else {
+                throw std::runtime_error("Invalid preprocess method");
+            }
+        }
+    }
+
+    auto cfg = rom["config"];
     if(!cfg) {
         SPDLOG_ERROR("No config found for {}", !isDirectoryMode ? this->gCartridge->GetHash() : GetSafeNode<std::string>(config, "folder"));
         return;
diff --git a/src/preprocess/CompTool.cpp b/src/preprocess/CompTool.cpp
new file mode 100644
index 0000000..9f585ef
--- /dev/null
+++ b/src/preprocess/CompTool.cpp
@@ -0,0 +1,86 @@
+#include "CompTool.h"
+#include "utils/Decompressor.h"
+#include "lib/binarytools/BinaryWriter.h"
+#include "lib/binarytools/BinaryReader.h"
+#include <fstream>
+
+uint32_t CompTool::FindFileTable(std::vector<uint8_t>& rom) {
+    uint8_t query_one[] = { 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x10, 0x50, 0x00, 0x00, 0x00, 0x00 };
+    uint8_t query_two[] = { 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x10, 0x60, 0x00, 0x00, 0x00, 0x00 };
+
+    for(size_t i = 0; i < rom.size() - sizeof(query_one); i++) {
+        if(memcmp(rom.data() + i, query_one, sizeof(query_one)) == 0){
+            return i;
+        }
+
+        if(memcmp(rom.data() + i, query_two, sizeof(query_two)) == 0){
+            return i;
+        }
+    }
+
+    throw std::runtime_error("Failed to find file table");
+}
+
+std::vector<uint8_t> CompTool::Decompress(std::vector<uint8_t> rom){
+    LUS::BinaryReader basefile((char*) rom.data(), rom.size());
+    basefile.SetEndianness(Torch::Endianness::Big);
+
+    LUS::BinaryWriter decompfile;
+    decompfile.SetEndianness(Torch::Endianness::Big);
+    decompfile.Write((uint8_t)0x80);
+
+    uint32_t table = CompTool::FindFileTable(rom);
+    uint32_t count = 0;
+    while (true){
+        auto entry = table + 0x10 * count;
+        basefile.Seek(entry, LUS::SeekOffsetType::Start);
+
+        auto v_begin = basefile.ReadInt32();
+        auto p_begin = basefile.ReadInt32();
+        auto p_end = basefile.ReadInt32();
+        auto comp_flag = basefile.ReadInt32();
+
+        auto p_size = p_end - p_begin;
+        auto v_size = (int32_t) 0;
+        DataChunk* decoded = nullptr;
+
+        if(v_begin == 0 && p_end == 0){
+            break;
+        }
+
+        basefile.Seek(p_begin, LUS::SeekOffsetType::Start);
+
+        auto bytes = new uint8_t[p_size];
+        basefile.Read((char*) bytes, p_size);
+
+        switch ((CompType) comp_flag) {
+            case CompType::UNCOMPRESSED:
+                v_size = p_size;
+                break;
+            case CompType::COMPRESSED:
+                decoded = Decompressor::Decode(std::vector(bytes, bytes + p_size), 0, CompressionType::MIO0, true);
+                bytes = decoded->data;
+                v_size = decoded->size;
+                break;
+            default:
+                throw std::runtime_error("Invalid compression flag. There may be a problem with your ROM.");
+        }
+
+        decompfile.Seek(v_begin, LUS::SeekOffsetType::Start);
+        decompfile.Write((char*) bytes, v_size);
+        auto v_end = v_begin + v_size;
+
+        decompfile.Seek(entry + 4, LUS::SeekOffsetType::Start);
+        decompfile.Write(v_begin);
+        decompfile.Write(v_end);
+        decompfile.Write((uint32_t) CompType::UNCOMPRESSED);
+        count++;
+    }
+
+    decompfile.Seek(0x10, LUS::SeekOffsetType::Start);
+    decompfile.Write(0xA7D5F194); // CRC1
+    decompfile.Write(0xFE3DF761); // CRC2
+
+    auto result = decompfile.ToVector();
+    return { (uint8_t*) result.data(), (uint8_t*) result.data() + result.size() };
+}
\ No newline at end of file
diff --git a/src/preprocess/CompTool.h b/src/preprocess/CompTool.h
new file mode 100644
index 0000000..15b7db5
--- /dev/null
+++ b/src/preprocess/CompTool.h
@@ -0,0 +1,18 @@
+#pragma once
+
+#include <cstdint>
+#include <string>
+#include <vector>
+
+enum class CompType {
+    UNCOMPRESSED,
+    COMPRESSED,
+    UNKNOWN
+};
+
+class CompTool {
+public:
+    static std::vector<uint8_t> Decompress(std::vector<uint8_t> rom);
+private:
+    static uint32_t FindFileTable(std::vector<uint8_t>& rom);
+};
\ No newline at end of file
diff --git a/src/utils/Decompressor.cpp b/src/utils/Decompressor.cpp
index 95af7b1..1fee382 100644
--- a/src/utils/Decompressor.cpp
+++ b/src/utils/Decompressor.cpp
@@ -12,9 +12,9 @@ extern "C" {
 
 std::unordered_map<uint32_t, DataChunk*> gCachedChunks;
 
-DataChunk* Decompressor::Decode(const std::vector<uint8_t>& buffer, const uint32_t offset, const CompressionType type) {
+DataChunk* Decompressor::Decode(const std::vector<uint8_t>& buffer, const uint32_t offset, const CompressionType type, bool ignoreCache) {
 
-    if(gCachedChunks.contains(offset)){
+    if(!ignoreCache && gCachedChunks.contains(offset)){
         return gCachedChunks[offset];
     }
 
@@ -208,4 +208,4 @@ void Decompressor::ClearCache() {
         delete value->data;
     }
     gCachedChunks.clear();
-}
+}
\ No newline at end of file
diff --git a/src/utils/Decompressor.h b/src/utils/Decompressor.h
index 4d2056a..4aff942 100644
--- a/src/utils/Decompressor.h
+++ b/src/utils/Decompressor.h
@@ -32,7 +32,7 @@ struct DecompressedData {
 
 class Decompressor {
 public:
-    static DataChunk* Decode(const std::vector<uint8_t>& buffer, uint32_t offset, CompressionType type);
+    static DataChunk* Decode(const std::vector<uint8_t>& buffer, uint32_t offset, CompressionType type, 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);