From 7b83e15d8068c6fc18b52c90b0cbc425b5abd63f Mon Sep 17 00:00:00 2001
From: KiritoDv <kiritodev01@gmail.com>
Date: Wed, 29 Jan 2025 22:44:12 -0600
Subject: [PATCH] Implemented GetSafeNodeByAddr

---
 src/Companion.cpp                      | 18 ++++++++++++++++++
 src/Companion.h                        |  3 ++-
 src/factories/DisplayListOverrides.cpp | 18 +++++++++---------
 3 files changed, 29 insertions(+), 10 deletions(-)

diff --git a/src/Companion.cpp b/src/Companion.cpp
index 4a7a3fc..8b0def9 100644
--- a/src/Companion.cpp
+++ b/src/Companion.cpp
@@ -1398,6 +1398,24 @@ std::optional<std::tuple<std::string, YAML::Node>> Companion::GetNodeByAddr(cons
     return this->gAddrMap[this->gCurrentFile][addr];
 }
 
+std::optional<std::tuple<std::string, YAML::Node>> Companion::GetSafeNodeByAddr(const uint32_t addr, std::string type) {
+    auto node = this->GetNodeByAddr(addr);
+
+    if(!node.has_value()) {
+        return std::nullopt;
+    }
+
+    auto [name, n] = node.value();
+    auto n_type = GetSafeNode<std::string>(n, "type");
+
+    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));
+    }
+
+    return node;
+
+}
+
 std::optional<ParseResultData> Companion::GetParseDataByAddr(uint32_t addr) {
     if(!this->gParseResults.contains(this->gCurrentFile)){
         for (auto &file : this->gCurrentExternalFiles) {
diff --git a/src/Companion.h b/src/Companion.h
index 2cebae4..3d7f0f5 100644
--- a/src/Companion.h
+++ b/src/Companion.h
@@ -146,8 +146,9 @@ public:
     std::optional<ParseResultData> GetParseDataBySymbol(const std::string& symbol);
 
     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::tuple<std::string, YAML::Node>> GetNodeByAddr(uint32_t addr);
+    std::optional<std::tuple<std::string, YAML::Node>> GetSafeNodeByAddr(const uint32_t addr, std::string type);
     std::optional<std::vector<std::tuple<std::string, YAML::Node>>> GetNodesByType(const std::string& type);
 
     std::optional<std::uint32_t> GetFileOffset(void) const { return this->gCurrentFileOffset; };
diff --git a/src/factories/DisplayListOverrides.cpp b/src/factories/DisplayListOverrides.cpp
index 7284d29..1f7c07f 100644
--- a/src/factories/DisplayListOverrides.cpp
+++ b/src/factories/DisplayListOverrides.cpp
@@ -71,7 +71,7 @@ int Vtx(uint32_t ptr, int32_t num) {
         return 1;
     }
 
-    auto dec = Companion::Instance->GetNodeByAddr(ptr);
+    auto dec = Companion::Instance->GetSafeNodeByAddr(ptr, "VTX");
 
     if(dec.has_value()){
         auto node = std::get<1>(dec.value());
@@ -86,7 +86,7 @@ int Vtx(uint32_t ptr, int32_t num) {
 }
 
 int Texture(uint32_t ptr, int32_t fmt, int32_t siz, int32_t width, int32_t height, int32_t pal) {
-    auto dec = Companion::Instance->GetNodeByAddr(ptr);
+    auto dec = Companion::Instance->GetSafeNodeByAddr(ptr, "TEXTURE");
 
     if(dec.has_value()){
         auto node = std::get<1>(dec.value());
@@ -101,7 +101,7 @@ int Texture(uint32_t ptr, int32_t fmt, int32_t siz, int32_t width, int32_t heigh
 }
 
 int Palette(uint32_t ptr, int32_t idx, int32_t count) {
-    auto dec = Companion::Instance->GetNodeByAddr(ptr);
+    auto dec = Companion::Instance->GetSafeNodeByAddr(ptr, "TEXTURE");
 
     if(dec.has_value()){
         auto node = std::get<1>(dec.value());
@@ -116,7 +116,7 @@ int Palette(uint32_t ptr, int32_t idx, int32_t count) {
 }
 
 int Lights(uint32_t ptr, int32_t count) {
-    auto dec = Companion::Instance->GetNodeByAddr(ptr);
+    auto dec = Companion::Instance->GetSafeNodeByAddr(ptr, "LIGHTS");
 
     if(dec.has_value()){
         auto node = std::get<1>(dec.value());
@@ -131,7 +131,7 @@ int Lights(uint32_t ptr, int32_t count) {
 }
 
 int Light(uint32_t ptr) {
-    auto res = Companion::Instance->GetNodeByAddr(ptr);
+    auto res = Companion::Instance->GetSafeNodeByAddr(ptr, "LIGHTS");
 
     if(res.has_value()){
         auto node = std::get<1>(res.value());
@@ -141,7 +141,7 @@ int Light(uint32_t ptr) {
         return 1;
     }
 
-    res = Companion::Instance->GetNodeByAddr(ptr - 0x8);
+    res = Companion::Instance->GetSafeNodeByAddr(ptr - 0x8, "LIGHTS");
 
     if(res.has_value()){
         auto node = std::get<1>(res.value());
@@ -156,7 +156,7 @@ int Light(uint32_t ptr) {
 }
 
 int DisplayList(uint32_t ptr) {
-    auto dec = Companion::Instance->GetNodeByAddr(ptr);
+    auto dec = Companion::Instance->GetSafeNodeByAddr(ptr, "GFX");
 
     if(dec.has_value()){
         auto node = std::get<1>(dec.value());
@@ -171,7 +171,7 @@ int DisplayList(uint32_t ptr) {
 }
 
 int Viewport(uint32_t ptr) {
-    auto dec = Companion::Instance->GetNodeByAddr(ptr);
+    auto dec = Companion::Instance->GetSafeNodeByAddr(ptr, "VP");
 
     if(dec.has_value()){
         auto node = std::get<1>(dec.value());
@@ -186,7 +186,7 @@ int Viewport(uint32_t ptr) {
 }
 
 int Matrix(uint32_t ptr) {
-    auto dec = Companion::Instance->GetNodeByAddr(ptr);
+    auto dec = Companion::Instance->GetSafeNodeByAddr(ptr, "MTX");
 
     if(dec.has_value()){
         auto node = std::get<1>(dec.value());