From b0c6bf497a1eabec14c116b710dcc757e77455bf Mon Sep 17 00:00:00 2001
From: Liam <byteslice@airmail.cc>
Date: Tue, 31 Oct 2023 20:11:14 -0400
Subject: romfs: fix extraction of single-directory root

---
 src/core/file_sys/romfs.cpp | 44 +++++++++++++++-----------------------------
 1 file changed, 15 insertions(+), 29 deletions(-)

(limited to 'src/core/file_sys/romfs.cpp')

diff --git a/src/core/file_sys/romfs.cpp b/src/core/file_sys/romfs.cpp
index 1c580de57..1eb1f439a 100644
--- a/src/core/file_sys/romfs.cpp
+++ b/src/core/file_sys/romfs.cpp
@@ -35,13 +35,14 @@ struct RomFSHeader {
 static_assert(sizeof(RomFSHeader) == 0x50, "RomFSHeader has incorrect size.");
 
 struct DirectoryEntry {
+    u32_le parent;
     u32_le sibling;
     u32_le child_dir;
     u32_le child_file;
     u32_le hash;
     u32_le name_length;
 };
-static_assert(sizeof(DirectoryEntry) == 0x14, "DirectoryEntry has incorrect size.");
+static_assert(sizeof(DirectoryEntry) == 0x18, "DirectoryEntry has incorrect size.");
 
 struct FileEntry {
     u32_le parent;
@@ -64,25 +65,22 @@ std::pair<Entry, std::string> GetEntry(const VirtualFile& file, std::size_t offs
     return {entry, string};
 }
 
-void ProcessFile(VirtualFile file, std::size_t file_offset, std::size_t data_offset,
-                 u32 this_file_offset, std::shared_ptr<VectorVfsDirectory> parent) {
-    while (true) {
+void ProcessFile(const VirtualFile& file, std::size_t file_offset, std::size_t data_offset,
+                 u32 this_file_offset, std::shared_ptr<VectorVfsDirectory>& parent) {
+    while (this_file_offset != ROMFS_ENTRY_EMPTY) {
         auto entry = GetEntry<FileEntry>(file, file_offset + this_file_offset);
 
         parent->AddFile(std::make_shared<OffsetVfsFile>(
             file, entry.first.size, entry.first.offset + data_offset, entry.second));
 
-        if (entry.first.sibling == ROMFS_ENTRY_EMPTY)
-            break;
-
         this_file_offset = entry.first.sibling;
     }
 }
 
-void ProcessDirectory(VirtualFile file, std::size_t dir_offset, std::size_t file_offset,
+void ProcessDirectory(const VirtualFile& file, std::size_t dir_offset, std::size_t file_offset,
                       std::size_t data_offset, u32 this_dir_offset,
-                      std::shared_ptr<VectorVfsDirectory> parent) {
-    while (true) {
+                      std::shared_ptr<VectorVfsDirectory>& parent) {
+    while (this_dir_offset != ROMFS_ENTRY_EMPTY) {
         auto entry = GetEntry<DirectoryEntry>(file, dir_offset + this_dir_offset);
         auto current = std::make_shared<VectorVfsDirectory>(
             std::vector<VirtualFile>{}, std::vector<VirtualDir>{}, entry.second);
@@ -97,14 +95,12 @@ void ProcessDirectory(VirtualFile file, std::size_t dir_offset, std::size_t file
         }
 
         parent->AddDirectory(current);
-        if (entry.first.sibling == ROMFS_ENTRY_EMPTY)
-            break;
         this_dir_offset = entry.first.sibling;
     }
 }
 } // Anonymous namespace
 
-VirtualDir ExtractRomFS(VirtualFile file, RomFSExtractionType type) {
+VirtualDir ExtractRomFS(VirtualFile file) {
     RomFSHeader header{};
     if (file->ReadObject(&header) != sizeof(RomFSHeader))
         return nullptr;
@@ -113,27 +109,17 @@ VirtualDir ExtractRomFS(VirtualFile file, RomFSExtractionType type) {
         return nullptr;
 
     const u64 file_offset = header.file_meta.offset;
-    const u64 dir_offset = header.directory_meta.offset + 4;
-
-    auto root =
-        std::make_shared<VectorVfsDirectory>(std::vector<VirtualFile>{}, std::vector<VirtualDir>{},
-                                             file->GetName(), file->GetContainingDirectory());
-
-    ProcessDirectory(file, dir_offset, file_offset, header.data_offset, 0, root);
+    const u64 dir_offset = header.directory_meta.offset;
 
-    VirtualDir out = std::move(root);
+    auto root_container = std::make_shared<VectorVfsDirectory>();
 
-    if (type == RomFSExtractionType::SingleDiscard)
-        return out->GetSubdirectories().front();
+    ProcessDirectory(file, dir_offset, file_offset, header.data_offset, 0, root_container);
 
-    while (out->GetSubdirectories().size() == 1 && out->GetFiles().empty()) {
-        if (Common::ToLower(out->GetSubdirectories().front()->GetName()) == "data" &&
-            type == RomFSExtractionType::Truncated)
-            break;
-        out = out->GetSubdirectories().front();
+    if (auto root = root_container->GetSubdirectory(""); root) {
+        return std::make_shared<CachedVfsDirectory>(std::move(root));
     }
 
-    return std::make_shared<CachedVfsDirectory>(std::move(out));
+    return nullptr;
 }
 
 VirtualFile CreateRomFS(VirtualDir dir, VirtualDir ext) {
-- 
cgit v1.2.3