Merge pull request #12208 from liamwhite/romfs

romfs: optimize parsing and building

6 files changed, 137 insertions, 113 deletions

diff --git a/src/core/file_sys/fsmitm_romfsbuild.cpp b/src/core/file_sys/fsmitm_romfsbuild.cpp
index f1d3e4129..dd9cca103 100644
--- a/src/core/file_sys/fsmitm_romfsbuild.cpp
+++ b/src/core/file_sys/fsmitm_romfsbuild.cpp
@@ -2,6 +2,7 @@
 // SPDX-License-Identifier: GPL-2.0-or-later
 
 #include <cstring>
+#include <span>
 #include <string_view>
 #include "common/alignment.h"
 #include "common/assert.h"
@@ -134,7 +135,7 @@ void RomFSBuildContext::VisitDirectory(VirtualDir romfs_dir, VirtualDir ext_dir,
 
         child->size = child->source->GetSize();
 
-        AddFile(parent, child);
+        AddFile(parent, std::move(child));
     }
 
     for (auto& child_romfs_dir : romfs_dir->GetSubdirectories()) {
@@ -163,36 +164,24 @@ void RomFSBuildContext::VisitDirectory(VirtualDir romfs_dir, VirtualDir ext_dir,
 
 bool RomFSBuildContext::AddDirectory(std::shared_ptr<RomFSBuildDirectoryContext> parent_dir_ctx,
                                      std::shared_ptr<RomFSBuildDirectoryContext> dir_ctx) {
-    // Check whether it's already in the known directories.
-    const auto [it, is_new] = directories.emplace(dir_ctx->path, nullptr);
-    if (!is_new) {
-        return false;
-    }
-
     // Add a new directory.
     num_dirs++;
     dir_table_size +=
         sizeof(RomFSDirectoryEntry) + Common::AlignUp(dir_ctx->path_len - dir_ctx->cur_path_ofs, 4);
-    dir_ctx->parent = parent_dir_ctx;
-    it->second = dir_ctx;
+    dir_ctx->parent = std::move(parent_dir_ctx);
+    directories.emplace_back(std::move(dir_ctx));
 
     return true;
 }
 
 bool RomFSBuildContext::AddFile(std::shared_ptr<RomFSBuildDirectoryContext> parent_dir_ctx,
                                 std::shared_ptr<RomFSBuildFileContext> file_ctx) {
-    // Check whether it's already in the known files.
-    const auto [it, is_new] = files.emplace(file_ctx->path, nullptr);
-    if (!is_new) {
-        return false;
-    }
-
     // Add a new file.
     num_files++;
     file_table_size +=
         sizeof(RomFSFileEntry) + Common::AlignUp(file_ctx->path_len - file_ctx->cur_path_ofs, 4);
-    file_ctx->parent = parent_dir_ctx;
-    it->second = file_ctx;
+    file_ctx->parent = std::move(parent_dir_ctx);
+    files.emplace_back(std::move(file_ctx));
 
     return true;
 }
@@ -201,7 +190,7 @@ RomFSBuildContext::RomFSBuildContext(VirtualDir base_, VirtualDir ext_)
     : base(std::move(base_)), ext(std::move(ext_)) {
     root = std::make_shared<RomFSBuildDirectoryContext>();
     root->path = "\0";
-    directories.emplace(root->path, root);
+    directories.emplace_back(root);
     num_dirs = 1;
     dir_table_size = 0x18;
 
@@ -210,28 +199,43 @@ RomFSBuildContext::RomFSBuildContext(VirtualDir base_, VirtualDir ext_)
 
 RomFSBuildContext::~RomFSBuildContext() = default;
 
-std::multimap<u64, VirtualFile> RomFSBuildContext::Build() {
+std::vector<std::pair<u64, VirtualFile>> RomFSBuildContext::Build() {
     const u64 dir_hash_table_entry_count = romfs_get_hash_table_count(num_dirs);
     const u64 file_hash_table_entry_count = romfs_get_hash_table_count(num_files);
     dir_hash_table_size = 4 * dir_hash_table_entry_count;
     file_hash_table_size = 4 * file_hash_table_entry_count;
 
-    // Assign metadata pointers
+    // Assign metadata pointers.
     RomFSHeader header{};
 
-    std::vector<u32> dir_hash_table(dir_hash_table_entry_count, ROMFS_ENTRY_EMPTY);
-    std::vector<u32> file_hash_table(file_hash_table_entry_count, ROMFS_ENTRY_EMPTY);
-
-    std::vector<u8> dir_table(dir_table_size);
-    std::vector<u8> file_table(file_table_size);
-
-    std::shared_ptr<RomFSBuildFileContext> cur_file;
+    std::vector<u8> metadata(file_hash_table_size + file_table_size + dir_hash_table_size +
+                             dir_table_size);
+    u32* const dir_hash_table_pointer = reinterpret_cast<u32*>(metadata.data());
+    u8* const dir_table_pointer = metadata.data() + dir_hash_table_size;
+    u32* const file_hash_table_pointer =
+        reinterpret_cast<u32*>(metadata.data() + dir_hash_table_size + dir_table_size);
+    u8* const file_table_pointer =
+        metadata.data() + dir_hash_table_size + dir_table_size + file_hash_table_size;
+
+    std::span<u32> dir_hash_table(dir_hash_table_pointer, dir_hash_table_entry_count);
+    std::span<u32> file_hash_table(file_hash_table_pointer, file_hash_table_entry_count);
+    std::span<u8> dir_table(dir_table_pointer, dir_table_size);
+    std::span<u8> file_table(file_table_pointer, file_table_size);
+
+    // Initialize hash tables.
+    std::memset(dir_hash_table.data(), 0xFF, dir_hash_table.size_bytes());
+    std::memset(file_hash_table.data(), 0xFF, file_hash_table.size_bytes());
+
+    // Sort tables by name.
+    std::sort(files.begin(), files.end(),
+              [](const auto& a, const auto& b) { return a->path < b->path; });
+    std::sort(directories.begin(), directories.end(),
+              [](const auto& a, const auto& b) { return a->path < b->path; });
 
     // Determine file offsets.
     u32 entry_offset = 0;
     std::shared_ptr<RomFSBuildFileContext> prev_file = nullptr;
-    for (const auto& it : files) {
-        cur_file = it.second;
+    for (const auto& cur_file : files) {
         file_partition_size = Common::AlignUp(file_partition_size, 16);
         cur_file->offset = file_partition_size;
         file_partition_size += cur_file->size;
@@ -243,34 +247,48 @@ std::multimap<u64, VirtualFile> RomFSBuildContext::Build() {
     }
     // Assign deferred parent/sibling ownership.
     for (auto it = files.rbegin(); it != files.rend(); ++it) {
-        cur_file = it->second;
+        auto& cur_file = *it;
         cur_file->sibling = cur_file->parent->file;
         cur_file->parent->file = cur_file;
     }
 
-    std::shared_ptr<RomFSBuildDirectoryContext> cur_dir;
-
     // Determine directory offsets.
     entry_offset = 0;
-    for (const auto& it : directories) {
-        cur_dir = it.second;
+    for (const auto& cur_dir : directories) {
         cur_dir->entry_offset = entry_offset;
         entry_offset +=
             static_cast<u32>(sizeof(RomFSDirectoryEntry) +
                              Common::AlignUp(cur_dir->path_len - cur_dir->cur_path_ofs, 4));
     }
     // Assign deferred parent/sibling ownership.
-    for (auto it = directories.rbegin(); it->second != root; ++it) {
-        cur_dir = it->second;
+    for (auto it = directories.rbegin(); (*it) != root; ++it) {
+        auto& cur_dir = *it;
         cur_dir->sibling = cur_dir->parent->child;
         cur_dir->parent->child = cur_dir;
     }
 
-    std::multimap<u64, VirtualFile> out;
+    // Create output map.
+    std::vector<std::pair<u64, VirtualFile>> out;
+    out.reserve(num_files + 2);
+
+    // Set header fields.
+    header.header_size = sizeof(RomFSHeader);
+    header.file_hash_table_size = file_hash_table_size;
+    header.file_table_size = file_table_size;
+    header.dir_hash_table_size = dir_hash_table_size;
+    header.dir_table_size = dir_table_size;
+    header.file_partition_ofs = ROMFS_FILEPARTITION_OFS;
+    header.dir_hash_table_ofs = Common::AlignUp(header.file_partition_ofs + file_partition_size, 4);
+    header.dir_table_ofs = header.dir_hash_table_ofs + header.dir_hash_table_size;
+    header.file_hash_table_ofs = header.dir_table_ofs + header.dir_table_size;
+    header.file_table_ofs = header.file_hash_table_ofs + header.file_hash_table_size;
+
+    std::vector<u8> header_data(sizeof(RomFSHeader));
+    std::memcpy(header_data.data(), &header, header_data.size());
+    out.emplace_back(0, std::make_shared<VectorVfsFile>(std::move(header_data)));
 
     // Populate file tables.
-    for (const auto& it : files) {
-        cur_file = it.second;
+    for (const auto& cur_file : files) {
         RomFSFileEntry cur_entry{};
 
         cur_entry.parent = cur_file->parent->entry_offset;
@@ -287,7 +305,7 @@ std::multimap<u64, VirtualFile> RomFSBuildContext::Build() {
 
         cur_entry.name_size = name_size;
 
-        out.emplace(cur_file->offset + ROMFS_FILEPARTITION_OFS, std::move(cur_file->source));
+        out.emplace_back(cur_file->offset + ROMFS_FILEPARTITION_OFS, std::move(cur_file->source));
         std::memcpy(file_table.data() + cur_file->entry_offset, &cur_entry, sizeof(RomFSFileEntry));
         std::memset(file_table.data() + cur_file->entry_offset + sizeof(RomFSFileEntry), 0,
                     Common::AlignUp(cur_entry.name_size, 4));
@@ -296,8 +314,7 @@ std::multimap<u64, VirtualFile> RomFSBuildContext::Build() {
     }
 
     // Populate dir tables.
-    for (const auto& it : directories) {
-        cur_dir = it.second;
+    for (const auto& cur_dir : directories) {
         RomFSDirectoryEntry cur_entry{};
 
         cur_entry.parent = cur_dir == root ? 0 : cur_dir->parent->entry_offset;
@@ -323,34 +340,13 @@ std::multimap<u64, VirtualFile> RomFSBuildContext::Build() {
                     cur_dir->path.data() + cur_dir->cur_path_ofs, name_size);
     }
 
-    // Set header fields.
-    header.header_size = sizeof(RomFSHeader);
-    header.file_hash_table_size = file_hash_table_size;
-    header.file_table_size = file_table_size;
-    header.dir_hash_table_size = dir_hash_table_size;
-    header.dir_table_size = dir_table_size;
-    header.file_partition_ofs = ROMFS_FILEPARTITION_OFS;
-    header.dir_hash_table_ofs = Common::AlignUp(header.file_partition_ofs + file_partition_size, 4);
-    header.dir_table_ofs = header.dir_hash_table_ofs + header.dir_hash_table_size;
-    header.file_hash_table_ofs = header.dir_table_ofs + header.dir_table_size;
-    header.file_table_ofs = header.file_hash_table_ofs + header.file_hash_table_size;
-
-    std::vector<u8> header_data(sizeof(RomFSHeader));
-    std::memcpy(header_data.data(), &header, header_data.size());
-    out.emplace(0, std::make_shared<VectorVfsFile>(std::move(header_data)));
+    // Write metadata.
+    out.emplace_back(header.dir_hash_table_ofs,
+                     std::make_shared<VectorVfsFile>(std::move(metadata)));
 
-    std::vector<u8> metadata(file_hash_table_size + file_table_size + dir_hash_table_size +
-                             dir_table_size);
-    std::size_t index = 0;
-    std::memcpy(metadata.data(), dir_hash_table.data(), dir_hash_table.size() * sizeof(u32));
-    index += dir_hash_table.size() * sizeof(u32);
-    std::memcpy(metadata.data() + index, dir_table.data(), dir_table.size());
-    index += dir_table.size();
-    std::memcpy(metadata.data() + index, file_hash_table.data(),
-                file_hash_table.size() * sizeof(u32));
-    index += file_hash_table.size() * sizeof(u32);
-    std::memcpy(metadata.data() + index, file_table.data(), file_table.size());
-    out.emplace(header.dir_hash_table_ofs, std::make_shared<VectorVfsFile>(std::move(metadata)));
+    // Sort the output.
+    std::sort(out.begin(), out.end(),
+              [](const auto& a, const auto& b) { return a.first < b.first; });
 
     return out;
 }
diff --git a/src/core/file_sys/fsmitm_romfsbuild.h b/src/core/file_sys/fsmitm_romfsbuild.h
index 06e5d5a47..f387c79f1 100644
--- a/src/core/file_sys/fsmitm_romfsbuild.h
+++ b/src/core/file_sys/fsmitm_romfsbuild.h
@@ -22,14 +22,14 @@ public:
     ~RomFSBuildContext();
 
     // This finalizes the context.
-    std::multimap<u64, VirtualFile> Build();
+    std::vector<std::pair<u64, VirtualFile>> Build();
 
 private:
     VirtualDir base;
     VirtualDir ext;
     std::shared_ptr<RomFSBuildDirectoryContext> root;
-    std::map<std::string, std::shared_ptr<RomFSBuildDirectoryContext>, std::less<>> directories;
-    std::map<std::string, std::shared_ptr<RomFSBuildFileContext>, std::less<>> files;
+    std::vector<std::shared_ptr<RomFSBuildDirectoryContext>> directories;
+    std::vector<std::shared_ptr<RomFSBuildFileContext>> files;
     u64 num_dirs = 0;
     u64 num_files = 0;
     u64 dir_table_size = 0;
diff --git a/src/core/file_sys/romfs.cpp b/src/core/file_sys/romfs.cpp
index 6de2103a0..6182598ae 100644
--- a/src/core/file_sys/romfs.cpp
+++ b/src/core/file_sys/romfs.cpp
@@ -55,44 +55,68 @@ struct FileEntry {
 };
 static_assert(sizeof(FileEntry) == 0x20, "FileEntry has incorrect size.");
 
-template <typename Entry>
-std::pair<Entry, std::string> GetEntry(const VirtualFile& file, std::size_t offset) {
-    Entry entry{};
-    if (file->ReadObject(&entry, offset) != sizeof(Entry))
-        return {};
-    std::string string(entry.name_length, '\0');
-    if (file->ReadArray(&string[0], string.size(), offset + sizeof(Entry)) != string.size())
+struct RomFSTraversalContext {
+    RomFSHeader header;
+    VirtualFile file;
+    std::vector<u8> directory_meta;
+    std::vector<u8> file_meta;
+};
+
+template <typename EntryType, auto Member>
+std::pair<EntryType, std::string> GetEntry(const RomFSTraversalContext& ctx, size_t offset) {
+    const size_t entry_end = offset + sizeof(EntryType);
+    const std::vector<u8>& vec = ctx.*Member;
+    const size_t size = vec.size();
+    const u8* data = vec.data();
+    EntryType entry{};
+
+    if (entry_end > size) {
         return {};
-    return {entry, string};
+    }
+    std::memcpy(&entry, data + offset, sizeof(EntryType));
+
+    const size_t name_length = std::min(entry_end + entry.name_length, size) - entry_end;
+    std::string name(reinterpret_cast<const char*>(data + entry_end), name_length);
+
+    return {entry, std::move(name)};
+}
+
+std::pair<DirectoryEntry, std::string> GetDirectoryEntry(const RomFSTraversalContext& ctx,
+                                                         size_t directory_offset) {
+    return GetEntry<DirectoryEntry, &RomFSTraversalContext::directory_meta>(ctx, directory_offset);
+}
+
+std::pair<FileEntry, std::string> GetFileEntry(const RomFSTraversalContext& ctx,
+                                               size_t file_offset) {
+    return GetEntry<FileEntry, &RomFSTraversalContext::file_meta>(ctx, file_offset);
 }
 
-void ProcessFile(const VirtualFile& file, std::size_t file_offset, std::size_t data_offset,
-                 u32 this_file_offset, std::shared_ptr<VectorVfsDirectory>& parent) {
+void ProcessFile(const RomFSTraversalContext& ctx, 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);
+        auto entry = GetFileEntry(ctx, this_file_offset);
 
-        parent->AddFile(std::make_shared<OffsetVfsFile>(
-            file, entry.first.size, entry.first.offset + data_offset, entry.second));
+        parent->AddFile(std::make_shared<OffsetVfsFile>(ctx.file, entry.first.size,
+                                                        entry.first.offset + ctx.header.data_offset,
+                                                        std::move(entry.second)));
 
         this_file_offset = entry.first.sibling;
     }
 }
 
-void ProcessDirectory(const VirtualFile& file, std::size_t dir_offset, std::size_t file_offset,
-                      std::size_t data_offset, u32 this_dir_offset,
+void ProcessDirectory(const RomFSTraversalContext& ctx, u32 this_dir_offset,
                       std::shared_ptr<VectorVfsDirectory>& parent) {
     while (this_dir_offset != ROMFS_ENTRY_EMPTY) {
-        auto entry = GetEntry<DirectoryEntry>(file, dir_offset + this_dir_offset);
+        auto entry = GetDirectoryEntry(ctx, this_dir_offset);
         auto current = std::make_shared<VectorVfsDirectory>(
             std::vector<VirtualFile>{}, std::vector<VirtualDir>{}, entry.second);
 
         if (entry.first.child_file != ROMFS_ENTRY_EMPTY) {
-            ProcessFile(file, file_offset, data_offset, entry.first.child_file, current);
+            ProcessFile(ctx, entry.first.child_file, current);
         }
 
         if (entry.first.child_dir != ROMFS_ENTRY_EMPTY) {
-            ProcessDirectory(file, dir_offset, file_offset, data_offset, entry.first.child_dir,
-                             current);
+            ProcessDirectory(ctx, entry.first.child_dir, current);
         }
 
         parent->AddDirectory(current);
@@ -107,22 +131,25 @@ VirtualDir ExtractRomFS(VirtualFile file) {
         return root_container;
     }
 
-    RomFSHeader header{};
-    if (file->ReadObject(&header) != sizeof(RomFSHeader)) {
-        return root_container;
+    RomFSTraversalContext ctx{};
+
+    if (file->ReadObject(&ctx.header) != sizeof(RomFSHeader)) {
+        return nullptr;
     }
 
-    if (header.header_size != sizeof(RomFSHeader)) {
-        return root_container;
+    if (ctx.header.header_size != sizeof(RomFSHeader)) {
+        return nullptr;
     }
 
-    const u64 file_offset = header.file_meta.offset;
-    const u64 dir_offset = header.directory_meta.offset;
+    ctx.file = file;
+    ctx.directory_meta =
+        file->ReadBytes(ctx.header.directory_meta.size, ctx.header.directory_meta.offset);
+    ctx.file_meta = file->ReadBytes(ctx.header.file_meta.size, ctx.header.file_meta.offset);
 
-    ProcessDirectory(file, dir_offset, file_offset, header.data_offset, 0, root_container);
+    ProcessDirectory(ctx, 0, root_container);
 
     if (auto root = root_container->GetSubdirectory(""); root) {
-        return std::make_shared<CachedVfsDirectory>(std::move(root));
+        return root;
     }
 
     ASSERT(false);
diff --git a/src/core/file_sys/vfs_concat.cpp b/src/core/file_sys/vfs_concat.cpp
index 168b9cbec..7c7298527 100644
--- a/src/core/file_sys/vfs_concat.cpp
+++ b/src/core/file_sys/vfs_concat.cpp
@@ -59,8 +59,8 @@ VirtualFile ConcatenatedVfsFile::MakeConcatenatedFile(std::string&& name,
     return VirtualFile(new ConcatenatedVfsFile(std::move(name), std::move(concatenation_map)));
 }
 
-VirtualFile ConcatenatedVfsFile::MakeConcatenatedFile(u8 filler_byte, std::string&& name,
-                                                      std::multimap<u64, VirtualFile>&& files) {
+VirtualFile ConcatenatedVfsFile::MakeConcatenatedFile(
+    u8 filler_byte, std::string&& name, std::vector<std::pair<u64, VirtualFile>>&& files) {
     // Fold trivial cases.
     if (files.empty()) {
         return nullptr;
diff --git a/src/core/file_sys/vfs_concat.h b/src/core/file_sys/vfs_concat.h
index cbddd12bd..b5f3d72e3 100644
--- a/src/core/file_sys/vfs_concat.h
+++ b/src/core/file_sys/vfs_concat.h
@@ -37,7 +37,7 @@ public:
     /// Convenience function that turns a map of offsets to files into a concatenated file, filling
     /// gaps with a given filler byte.
     static VirtualFile MakeConcatenatedFile(u8 filler_byte, std::string&& name,
-                                            std::multimap<u64, VirtualFile>&& files);
+                                            std::vector<std::pair<u64, VirtualFile>>&& files);
 
     std::string GetName() const override;
     std::size_t GetSize() const override;
diff --git a/src/core/file_sys/vfs_layered.cpp b/src/core/file_sys/vfs_layered.cpp
index 08daca397..5551743fb 100644
--- a/src/core/file_sys/vfs_layered.cpp
+++ b/src/core/file_sys/vfs_layered.cpp
@@ -3,6 +3,7 @@
 
 #include <algorithm>
 #include <set>
+#include <unordered_set>
 #include <utility>
 #include "core/file_sys/vfs_layered.h"
 
@@ -59,13 +60,12 @@ std::string LayeredVfsDirectory::GetFullPath() const {
 
 std::vector<VirtualFile> LayeredVfsDirectory::GetFiles() const {
     std::vector<VirtualFile> out;
-    std::set<std::string, std::less<>> out_names;
+    std::unordered_set<std::string> out_names;
 
     for (const auto& layer : dirs) {
         for (auto& file : layer->GetFiles()) {
-            auto file_name = file->GetName();
-            if (!out_names.contains(file_name)) {
-                out_names.emplace(std::move(file_name));
+            const auto [it, is_new] = out_names.emplace(file->GetName());
+            if (is_new) {
                 out.emplace_back(std::move(file));
             }
         }
@@ -75,18 +75,19 @@ std::vector<VirtualFile> LayeredVfsDirectory::GetFiles() const {
 }
 
 std::vector<VirtualDir> LayeredVfsDirectory::GetSubdirectories() const {
-    std::vector<std::string> names;
+    std::vector<VirtualDir> out;
+    std::unordered_set<std::string> out_names;
+
     for (const auto& layer : dirs) {
         for (const auto& sd : layer->GetSubdirectories()) {
-            if (std::find(names.begin(), names.end(), sd->GetName()) == names.end())
-                names.push_back(sd->GetName());
+            out_names.emplace(sd->GetName());
         }
     }
 
-    std::vector<VirtualDir> out;
-    out.reserve(names.size());
-    for (const auto& subdir : names)
+    out.reserve(out_names.size());
+    for (const auto& subdir : out_names) {
         out.emplace_back(GetSubdirectory(subdir));
+    }
 
     return out;
 }