add manual_segments

3 files changed, 81 insertions, 24 deletions

diff --git a/src/Companion.cpp b/src/Companion.cpp
index 231b24c..d5cb096 100644
--- a/src/Companion.cpp
+++ b/src/Companion.cpp
@@ -400,6 +400,23 @@ void Companion::ParseCurrentFileConfig(YAML::Node node) {
         }
     }
 
+    if(node["manual_segments"]) {
+        auto manualSegments = node["manual_segments"];
+        if (manualSegments.IsSequence() && manualSegments.size()) {
+            for(size_t i = 0; i < manualSegments.size(); i++) {
+                auto segment = manualSegments[i];
+                if (segment.IsSequence() && segment.size() == 2) {
+                    const auto id = segment[0].as<uint32_t>();
+                    const auto replacement = segment[1].as<std::string>();
+                    this->gManualSegments[id] = replacement;
+                    SPDLOG_DEBUG("Manual Segment {} replaced with {}", id, replacement);
+                } else {
+                    throw std::runtime_error("Incorrect yaml syntax for manual segments.\n\nThe yaml expects:\n:config:\n  manual_segments:\n  - [<addr>, <replacement>]\n\nLike so:\nmanual_segments:\n  - [0x05000000, \"textures/other_textures/texture_6447C4\"]");
+                }
+            }
+        }
+    }
+
     if(node["segments"]) {
         auto segments = node["segments"];
 
@@ -646,6 +663,7 @@ void Companion::ProcessFile(YAML::Node root) {
     this->gCurrentFileOffset = 0;
     this->gTables.clear();
     this->gCurrentExternalFiles.clear();
+    this->gManualSegments.clear();
     GFXDOverride::ClearVtx();
 
     if(root[":config"]) {
@@ -1482,6 +1500,20 @@ std::optional<std::tuple<std::string, YAML::Node>> Companion::GetNodeByAddr(uint
     return this->gAddrMap[this->gCurrentFile][addr];
 }
 
+std::optional<std::string> Companion::GetStringByAddr(const uint32_t addr) {
+    if(this->gManualSegments.contains(addr)) {
+        return this->gManualSegments[addr];
+    }
+
+    auto node = this->GetNodeByAddr(addr);
+
+    if(!node.has_value()) {
+        return std::nullopt;
+    }
+
+    return std::get<0>(node.value());
+}
+
 std::optional<std::tuple<std::string, YAML::Node>> Companion::GetSafeNodeByAddr(const uint32_t addr, std::string type) {
     auto node = this->GetNodeByAddr(addr);
 
@@ -1500,6 +1532,27 @@ std::optional<std::tuple<std::string, YAML::Node>> Companion::GetSafeNodeByAddr(
 
 }
 
+std::optional<std::string> Companion::GetSafeStringByAddr(const uint32_t addr, std::string type) {
+    if(this->gManualSegments.contains(addr)) {
+        return this->gManualSegments[addr];
+    }
+
+    auto node = this->GetNodeByAddr(addr);
+
+    if(!node.has_value()) {
+        return std::nullopt;
+    }
+
+    auto [name, n] = node.value();
+    auto n_type = GetTypeNode(n);
+
+    if(n_type != type) {
+        throw std::runtime_error("Requested node type does not match with the target node type at " + Torch::to_hex(addr, false) + " Found: " + n_type + " Expected: " + type);
+    }
+
+    return std::get<0>(node.value());
+}
+
 std::string Companion::GetSymbolFromAddr(uint32_t address, bool validZero) {
     auto dec = Companion::Instance->GetNodeByAddr(address);
     std::ostringstream outSymbol;
diff --git a/src/Companion.h b/src/Companion.h
index 95da2a4..2b6fc09 100644
--- a/src/Companion.h
+++ b/src/Companion.h
@@ -166,7 +166,9 @@ public:
     std::optional<std::shared_ptr<BaseFactory>> GetFactory(const std::string& type);
     uint32_t PatchVirtualAddr(uint32_t addr);
     std::optional<std::tuple<std::string, YAML::Node>> GetNodeByAddr(uint32_t addr);
+    std::optional<std::string> GetStringByAddr(uint32_t addr);
     std::optional<std::tuple<std::string, YAML::Node>> GetSafeNodeByAddr(const uint32_t addr, std::string type);
+    std::optional<std::string> GetSafeStringByAddr(const uint32_t addr, std::string type);
     std::optional<std::vector<std::tuple<std::string, YAML::Node>>> GetNodesByType(const std::string& type);
     std::string GetSymbolFromAddr(uint32_t addr, bool validZero = false);
 
@@ -219,6 +221,7 @@ private:
     CompressionType gCurrentCompressionType = CompressionType::None;
     std::vector<Table> gTables;
     std::vector<std::string> gCurrentExternalFiles;
+    std::unordered_map<int, std::string> gManualSegments;
     std::unordered_set<std::string> gProcessedFiles;
 
     std::unordered_map<std::string, std::vector<char>> gCompanionFiles;
diff --git a/src/factories/DisplayListFactory.cpp b/src/factories/DisplayListFactory.cpp
index 11bc7f8..faf6a67 100644
--- a/src/factories/DisplayListFactory.cpp
+++ b/src/factories/DisplayListFactory.cpp
@@ -315,15 +315,15 @@ ExportResult DListBinaryExporter::Export(std::ostream &write, std::shared_ptr<IP
                 SPDLOG_WARN("Could not find vtx at 0x{:X}", ptr);
             }
 
-            auto dec = Companion::Instance->GetNodeByAddr(ptr);
+            auto dec = Companion::Instance->GetSafeStringByAddr(ptr, "VTX");
 
             if(dec.has_value()){
-                uint64_t hash = CRC64(std::get<0>(dec.value()).c_str());
+                uint64_t hash = CRC64(dec.value().c_str());
                 if(hash == 0) {
-                    throw std::runtime_error("Vtx hash is 0 for " + std::get<0>(dec.value()));
+                    throw std::runtime_error("Vtx hash is 0 for " + dec.value());
                 }
 
-                SPDLOG_INFO("Found vtx: 0x{:X} Hash: 0x{:X} Path: {}", ptr, hash, std::get<0>(dec.value()));
+                SPDLOG_INFO("Found vtx: 0x{:X} Hash: 0x{:X} Path: {}", ptr, hash, dec.value());
 
                 N64Gfx value = gsSPVertexOTR(0, nvtx, didx);
 
@@ -345,7 +345,7 @@ ExportResult DListBinaryExporter::Export(std::ostream &write, std::shared_ptr<IP
         if(opcode == GBI(G_DL)) {
             N64Gfx value;
             auto ptr = w1;
-            auto dec = Companion::Instance->GetNodeByAddr(ptr);
+            auto dec = Companion::Instance->GetSafeStringByAddr(ptr, "GFX");
             auto branch = (w0 >> 16) & G_DL_NO_PUSH;
 
             // Export displaylist segment addresses as an index into a buffer of gfx
@@ -363,8 +363,8 @@ ExportResult DListBinaryExporter::Export(std::ostream &write, std::shared_ptr<IP
             writer.Write(w1);
 
             if(dec.has_value()){
-                uint64_t hash = CRC64(std::get<0>(dec.value()).c_str());
-                SPDLOG_INFO("Found display list: 0x{:X} Hash: 0x{:X} Path: {}", ptr, hash, std::get<0>(dec.value()));
+                uint64_t hash = CRC64(dec.value().c_str());
+                SPDLOG_INFO("Found display list: 0x{:X} Hash: 0x{:X} Path: {}", ptr, hash, dec.value());
                 w0 = hash >> 32;
                 w1 = hash & 0xFFFFFFFF;
             } else {
@@ -404,10 +404,10 @@ ExportResult DListBinaryExporter::Export(std::ostream &write, std::shared_ptr<IP
                     break;
             }
                         
-            auto res = Companion::Instance->GetNodeByAddr(ptr);
+            auto res = Companion::Instance->GetStringByAddr(ptr);
 
             if(!res.has_value()){
-                res = Companion::Instance->GetNodeByAddr(ptr - 0x8);
+                res = Companion::Instance->GetStringByAddr(ptr - 0x8);
                 hasOffset = res.has_value();
 
                 if(!hasOffset){
@@ -424,8 +424,8 @@ ExportResult DListBinaryExporter::Export(std::ostream &write, std::shared_ptr<IP
             writer.Write(w1);
 
             if(res.has_value()){
-                uint64_t hash = CRC64(std::get<0>(res.value()).c_str());
-                SPDLOG_INFO("Found movemem: 0x{:X} Hash: 0x{:X} Path: {}", ptr, hash, std::get<0>(res.value()));
+                uint64_t hash = CRC64(res.value().c_str());
+                SPDLOG_INFO("Found movemem: 0x{:X} Hash: 0x{:X} Path: {}", ptr, hash, res.value());
                 w0 = hash >> 32;
                 w1 = hash & 0xFFFFFFFF;
             } else {
@@ -435,14 +435,15 @@ ExportResult DListBinaryExporter::Export(std::ostream &write, std::shared_ptr<IP
 
         if(opcode == GBI(G_SETTIMG)) {
             auto ptr = w1;
-            auto dec = Companion::Instance->GetNodeByAddr(ptr);
+            auto dec = Companion::Instance->GetSafeStringByAddr(ptr, "TEXTURE");
 
             // Export texture segment addresses as segmented addresses
-            if ((Companion::Instance->GetGBIMinorVersion() == GBIMinorVersion::Mk64) && ((SEGMENT_NUMBER(w1) == 0x03) || (SEGMENT_NUMBER(w1) == 0x05))) {
-                w1 |= 1;
-                writer.Write(w0);
-                writer.Write(w1);
-            } else {
+            // if ((Companion::Instance->GetGBIMinorVersion() == GBIMinorVersion::Mk64) && ((SEGMENT_NUMBER(w1) == 0x03) || (SEGMENT_NUMBER(w1) == 0x05))) {
+            //     w1 |= 1;
+            //     writer.Write(w0);
+            //     writer.Write(w1);
+            // } else
+            {
                 N64Gfx value = gsDPSetTextureOTRImage(C0(21, 3), C0(19, 2), C0(0, 10), ptr);
                 w0 = value.words.w0;
                 w1 = value.words.w1;
@@ -451,13 +452,13 @@ ExportResult DListBinaryExporter::Export(std::ostream &write, std::shared_ptr<IP
                 writer.Write(w1);
 
                 if(dec.has_value()){
-                    uint64_t hash = CRC64(std::get<0>(dec.value()).c_str());
+                    uint64_t hash = CRC64(dec.value().c_str());
 
                     if(hash == 0){
-                        throw std::runtime_error("Texture hash is 0 for " + std::get<0>(dec.value()));
+                        throw std::runtime_error("Texture hash is 0 for " + dec.value());
                     }
 
-                    SPDLOG_INFO("Found texture: 0x{:X} Hash: 0x{:X} Path: {}", ptr, hash, std::get<0>(dec.value()));
+                    SPDLOG_INFO("Found texture: 0x{:X} Hash: 0x{:X} Path: {}", ptr, hash, dec.value());
                     w0 = hash >> 32;
                     w1 = hash & 0xFFFFFFFF;
                 } else {
@@ -468,7 +469,7 @@ ExportResult DListBinaryExporter::Export(std::ostream &write, std::shared_ptr<IP
 
         if(opcode == GBI(G_MTX)) {
             auto ptr = w1;
-            auto dec = Companion::Instance->GetNodeByAddr(ptr);
+            auto dec = Companion::Instance->GetSafeStringByAddr(ptr, "MTX");
 
             w0 &= 0x00FFFFFF;
             w0 += G_MTX_OTR << 24;
@@ -478,13 +479,13 @@ ExportResult DListBinaryExporter::Export(std::ostream &write, std::shared_ptr<IP
             writer.Write(w1);
 
             if(dec.has_value()){
-                uint64_t hash = CRC64(std::get<0>(dec.value()).c_str());
+                uint64_t hash = CRC64(dec.value().c_str());
 
                 if(hash == 0){
-                    throw std::runtime_error("Matrix hash is 0 for " + std::get<0>(dec.value()));
+                    throw std::runtime_error("Matrix hash is 0 for " + dec.value());
                 }
 
-                SPDLOG_INFO("Found matrix: 0x{:X} Hash: 0x{:X} Path: {}", ptr, hash, std::get<0>(dec.value()));
+                SPDLOG_INFO("Found matrix: 0x{:X} Hash: 0x{:X} Path: {}", ptr, hash, dec.value());
                 w0 = hash >> 32;
                 w1 = hash & 0xFFFFFFFF;
             } else {