From ded9fbf6769b07f99e6a47592aa90532b92d9b5f Mon Sep 17 00:00:00 2001
From: itsmattkc <34096995+itsmattkc@users.noreply.github.com>
Date: Sun, 3 Jul 2022 22:45:05 -0500
Subject: [PATCH] lib: cleanup and add find helper functions

---
 lib/chunk.cpp | 86 +++++++++++++++++++++++++++++++++++----------------
 lib/chunk.h   |  6 ++++
 2 files changed, 65 insertions(+), 27 deletions(-)

diff --git a/lib/chunk.cpp b/lib/chunk.cpp
index 92a02ed..76fb03f 100644
--- a/lib/chunk.cpp
+++ b/lib/chunk.cpp
@@ -54,33 +54,15 @@ bool Chunk::Read(std::ifstream &f, u32 &version, u32 &alignment)
   // Read custom data from this chunk
   Clear();
 
-  switch (type()) {
-  case TYPE_RIFF:
-    RIFF().Read(f, data_, version, size);
-    break;
-  case TYPE_MxHd:
-    MxHd().Read(f, data_, version, size);
-    version = data_["Version"];
-    alignment = data_["BufferSize"];
-    break;
-  case TYPE_LIST:
-    LIST().Read(f, data_, version, size);
-    break;
-  case TYPE_MxSt:
-    MxSt().Read(f, data_, version, size);
-    break;
-  case TYPE_MxCh:
-    MxCh().Read(f, data_, version, size);
-    break;
-  case TYPE_MxOf:
-    MxOf().Read(f, data_, version, size);
-    break;
-  case TYPE_pad_:
-    pad_().Read(f, data_, version, size);
-    break;
-  case TYPE_MxOb:
-    MxOb().Read(f, data_, version, size);
-    break;
+  if (RIFF *reader = GetReaderFromType(type())) {
+    reader->Read(f, data_, version, size);
+
+    if (type() == TYPE_MxHd) {
+      version = data_["Version"];
+      alignment = data_["BufferSize"];
+    }
+
+    delete reader;
   }
 
   // Assume any remaining data is this chunk's children
@@ -109,6 +91,30 @@ bool Chunk::Read(std::ifstream &f, u32 &version, u32 &alignment)
   return true;
 }
 
+RIFF *Chunk::GetReaderFromType(Type type)
+{
+  switch (type) {
+  case TYPE_RIFF:
+    return new RIFF();
+  case TYPE_MxHd:
+    return new MxHd();
+  case TYPE_LIST:
+    return new LIST();
+  case TYPE_MxSt:
+    return new MxSt();
+  case TYPE_MxCh:
+    return new MxCh();
+  case TYPE_MxOf:
+    return new MxOf();
+  case TYPE_pad_:
+    return new pad_();
+  case TYPE_MxOb:
+    return new MxOb();
+  }
+
+  return NULL;
+}
+
 void Chunk::Clear()
 {
   // Delete data
@@ -192,4 +198,30 @@ const char *Chunk::GetTypeDescription(Type type)
   return "Unknown";
 }
 
+Chunk *Chunk::FindChildWithType(Type type) const
+{
+  for (Chunk *child : children_) {
+    if (child->type() == type) {
+      return child;
+    } else if (Chunk *grandchild = child->FindChildWithType(type)) {
+      return grandchild;
+    }
+  }
+
+  return NULL;
+}
+
+Chunk *Chunk::FindChildWithOffset(u32 offset) const
+{
+  for (Chunk *child : children_) {
+    if (child->offset() == offset) {
+      return child;
+    } else if (Chunk *grandchild = child->FindChildWithOffset(offset)) {
+      return grandchild;
+    }
+  }
+
+  return NULL;
+}
+
 }
diff --git a/lib/chunk.h b/lib/chunk.h
index 90ac3fa..8a809ff 100644
--- a/lib/chunk.h
+++ b/lib/chunk.h
@@ -6,6 +6,7 @@
 #include <memory>
 
 #include "common.h"
+#include "sitypes.h"
 #include "types.h"
 
 namespace si {
@@ -57,6 +58,9 @@ public:
     return GetTypeDescription(type());
   }
 
+  LIBWEAVER_EXPORT Chunk *FindChildWithType(Type type) const;
+  LIBWEAVER_EXPORT Chunk *FindChildWithOffset(u32 offset) const;
+
 private:
   // Disable copy
   Chunk(const Chunk& other);
@@ -64,6 +68,8 @@ private:
 
   bool Read(std::ifstream &f, u32 &version, u32 &alignment);
 
+  static RIFF *GetReaderFromType(Type type);
+
   u32 id_;
   u32 offset_;
   std::map<std::string, Data> data_;