diff --git a/stratosphere/fs_mitm/source/fsmitm_romfsbuild.cpp b/stratosphere/fs_mitm/source/fsmitm_romfsbuild.cpp index 321453aa5..c5d1d1f90 100644 --- a/stratosphere/fs_mitm/source/fsmitm_romfsbuild.cpp +++ b/stratosphere/fs_mitm/source/fsmitm_romfsbuild.cpp @@ -10,6 +10,8 @@ void RomFSBuildContext::VisitDirectory(FsFileSystem *filesys, RomFSBuildDirector FsDir dir; Result rc; + std::vector child_dirs; + /* Open the current parent directory. */ if (R_FAILED((rc = Utils::OpenRomFSDir(filesys, this->title_id, parent->path, &dir)))) { fatalSimple(rc); @@ -33,6 +35,8 @@ void RomFSBuildContext::VisitDirectory(FsFileSystem *filesys, RomFSBuildDirector if (!this->AddDirectory(parent, child, NULL)) { delete child->path; delete child; + } else { + child_dirs.push_back(child); } } else if (this->dir_entry.type == ENTRYTYPE_FILE) { RomFSBuildFileContext *child = new RomFSBuildFileContext({0}); @@ -61,10 +65,8 @@ void RomFSBuildContext::VisitDirectory(FsFileSystem *filesys, RomFSBuildDirector } fsDirClose(&dir); - RomFSBuildDirectoryContext *cur_child = parent->child; - while (cur_child != NULL) { - this->VisitDirectory(filesys, cur_child); - cur_child = cur_child->sibling; + for (auto &child : child_dirs) { + this->VisitDirectory(filesys, child); } } @@ -179,57 +181,22 @@ void RomFSBuildContext::MergeRomStorage(IROStorage *storage, RomFSDataSource sou this->VisitDirectory(this->root, 0x0, dir_table.get(), (size_t)header.dir_table_size, file_table.get(), (size_t)header.file_table_size); } -bool RomFSBuildContext::AddDirectory(RomFSBuildDirectoryContext *parent_dir_ctx, RomFSBuildDirectoryContext *dir_ctx, RomFSBuildDirectoryContext **out_dir_ctx) { - /* Ordered insertion on child. */ - int cmp_val; - if (parent_dir_ctx->child == NULL || (cmp_val = strcmp(dir_ctx->path, parent_dir_ctx->child->path)) < 0) { - dir_ctx->sibling = parent_dir_ctx->child; - parent_dir_ctx->child = dir_ctx; - } else if (cmp_val == 0) { - /* This directory already exists! */ +bool RomFSBuildContext::AddDirectory(RomFSBuildDirectoryContext *parent_dir_ctx, RomFSBuildDirectoryContext *dir_ctx, RomFSBuildDirectoryContext **out_dir_ctx) { + /* Check whether it's already in the known directories. */ + auto existing = this->directories.find(dir_ctx->path); + if (existing != this->directories.end()) { if (out_dir_ctx) { - *out_dir_ctx = parent_dir_ctx->child; + *out_dir_ctx = existing->second; } return false; - } else { - RomFSBuildDirectoryContext *child, *prev; - prev = parent_dir_ctx->child; - child = prev->sibling; - while (child != NULL) { - cmp_val = strcmp(dir_ctx->path, child->path); - if (cmp_val < 0) { - break; - } else if (cmp_val == 0) { - /* This directory already exists! */ - if (out_dir_ctx) { - *out_dir_ctx = child; - } - return false; - } - prev = child; - child = child->sibling; - } - - prev->sibling = dir_ctx; - dir_ctx->sibling = child; } - /* If we got this far, we're definitely adding a new directory. */ + /* Add a new directory. */ this->num_dirs++; this->dir_table_size += sizeof(RomFSDirectoryEntry) + ((dir_ctx->path_len - dir_ctx->cur_path_ofs + 3) & ~3); dir_ctx->parent = parent_dir_ctx; + this->directories.insert({dir_ctx->path, dir_ctx}); - /* Ordered insertion on next */ - RomFSBuildDirectoryContext *cur = parent_dir_ctx->next, *prev = parent_dir_ctx; - while (cur != NULL) { - if (strcmp(dir_ctx->path, cur->path) < 0) { - break; - } - prev = cur; - cur = cur->next; - } - prev->next = dir_ctx; - dir_ctx->next = cur; if (out_dir_ctx) { *out_dir_ctx = dir_ctx; } @@ -237,55 +204,18 @@ bool RomFSBuildContext::AddDirectory(RomFSBuildDirectoryContext *parent_dir_ctx, } bool RomFSBuildContext::AddFile(RomFSBuildDirectoryContext *parent_dir_ctx, RomFSBuildFileContext *file_ctx) { - /* Ordered insertion on sibling */ - int cmp_val; - if (parent_dir_ctx->file == NULL || (cmp_val = strcmp(file_ctx->path, parent_dir_ctx->file->path)) < 0) { - file_ctx->sibling = parent_dir_ctx->file; - parent_dir_ctx->file = file_ctx; - } else if (cmp_val == 0) { - /* This file already exists! */ + /* Check whether it's already in the known files. */ + auto existing = this->files.find(file_ctx->path); + if (existing != this->files.end()) { return false; - } else { - RomFSBuildFileContext *child, *prev; - prev = parent_dir_ctx->file; - child = prev->sibling; - while (child != NULL) { - cmp_val = strcmp(file_ctx->path, child->path); - if (cmp_val < 0) { - break; - } else if (cmp_val == 0) { - /* This directory already exists! */ - return false; - } - prev = child; - child = child->sibling; - } - - prev->sibling = file_ctx; - file_ctx->sibling = child; } - /* If we got this far, we're definitely adding a new file. */ + /* Add a new file. */ this->num_files++; this->file_table_size += sizeof(RomFSFileEntry) + ((file_ctx->path_len - file_ctx->cur_path_ofs + 3) & ~3); file_ctx->parent = parent_dir_ctx; + this->files.insert({file_ctx->path, file_ctx}); - /* Ordered insertion on next */ - if (this->files == NULL || strcmp(file_ctx->path, this->files->path) < 0) { - file_ctx->next = this->files; - this->files = file_ctx; - } else { - RomFSBuildFileContext *cur = this->files->next, *prev = this->files; - while (cur != NULL) { - if (strcmp(file_ctx->path, cur->path) < 0) { - break; - } - prev = cur; - cur = cur->next; - } - prev->next = file_ctx; - file_ctx->next = cur; - } return true; } @@ -319,10 +249,10 @@ void RomFSBuildContext::Build(std::vector *out_infos) { out_infos->emplace_back(0, sizeof(*header), header, RomFSDataSource::Memory); /* Determine file offsets. */ - cur_file = this->files; entry_offset = 0; RomFSBuildFileContext *prev_file = NULL; - while (cur_file != NULL) { + for (const auto &it : this->files) { + cur_file = it.second; this->file_partition_size = (this->file_partition_size + 0xFULL) & ~0xFULL; /* Check for extra padding in the original romfs source and preserve it, to help ourselves later. */ if (prev_file && prev_file->source == cur_file->source && (prev_file->source == RomFSDataSource::BaseRomFS || prev_file->source == RomFSDataSource::FileRomFS)) { @@ -339,23 +269,45 @@ void RomFSBuildContext::Build(std::vector *out_infos) { this->file_partition_size += cur_file->size; cur_file->entry_offset = entry_offset; entry_offset += sizeof(RomFSFileEntry) + ((cur_file->path_len - cur_file->cur_path_ofs + 3) & ~3); + + /* Assign deferred parent/sibling ownership. */ + if (cur_file->parent->file == NULL) { + cur_file->parent->file = cur_file; + } else { + RomFSBuildFileContext *temp = cur_file->parent->file; + while (temp->sibling != NULL) { + temp = temp->sibling; + } + temp->sibling = cur_file; + } + prev_file = cur_file; - cur_file = cur_file->next; } /* Determine directory offsets. */ - cur_dir = this->root; entry_offset = 0; - while (cur_dir != NULL) { - cur_dir->entry_offset = entry_offset; + for (const auto &it : this->directories) { + cur_dir = it.second; entry_offset += sizeof(RomFSDirectoryEntry) + ((cur_dir->path_len - cur_dir->cur_path_ofs + 3) & ~3); - cur_dir = cur_dir->next; + + /* Assign deferred parent/sibling ownership. */ + if (cur_dir != this->root) { + if (cur_dir->parent->child == NULL) { + cur_dir->parent->child = cur_dir; + } else { + RomFSBuildDirectoryContext *temp = cur_dir->parent->child; + while (temp->sibling != NULL) { + temp = temp->sibling; + } + temp->sibling = cur_dir; + } + } } /* Populate file tables. */ - cur_file = this->files; - while (cur_file != NULL) { + for (const auto &it : this->files) { + cur_file = it.second; RomFSFileEntry *cur_entry = romfs_get_fentry(file_table, cur_file->entry_offset); cur_entry->parent = cur_file->parent->entry_offset; @@ -367,13 +319,11 @@ void RomFSBuildContext::Build(std::vector *out_infos) { u32 hash = romfs_calc_path_hash(cur_file->parent->entry_offset, (unsigned char *)cur_file->path + cur_file->cur_path_ofs, 0, name_size); cur_entry->hash = file_hash_table[hash % file_hash_table_entry_count]; file_hash_table[hash % file_hash_table_entry_count] = cur_file->entry_offset; - - + cur_entry->name_size = name_size; memset(cur_entry->name, 0, (cur_entry->name_size + 3) & ~3); memcpy(cur_entry->name, cur_file->path + cur_file->cur_path_ofs, name_size); - switch (cur_file->source) { case RomFSDataSource::BaseRomFS: case RomFSDataSource::FileRomFS: @@ -395,15 +345,14 @@ void RomFSBuildContext::Build(std::vector *out_infos) { fatalSimple(0xF601); } - RomFSBuildFileContext *temp = cur_file; - cur_file = cur_file->next; - delete temp->path; - delete temp; + delete cur_file->path; + delete cur_file; } + this->files.clear(); /* Populate dir tables. */ - cur_dir = this->root; - while (cur_dir != NULL) { + for (const auto &it : this->directories) { + cur_dir = it.second; RomFSDirectoryEntry *cur_entry = romfs_get_direntry(dir_table, cur_dir->entry_offset); cur_entry->parent = cur_dir == this->root ? 0 : cur_dir->parent->entry_offset; cur_entry->sibling = (cur_dir->sibling == NULL) ? ROMFS_ENTRY_EMPTY : cur_dir->sibling->entry_offset; @@ -419,11 +368,11 @@ void RomFSBuildContext::Build(std::vector *out_infos) { memset(cur_entry->name, 0, (cur_entry->name_size + 3) & ~3); memcpy(cur_entry->name, cur_dir->path + cur_dir->cur_path_ofs, name_size); - RomFSBuildDirectoryContext *temp = cur_dir; - cur_dir = cur_dir->next; - delete temp->path; - delete temp; + delete cur_dir->path; + delete cur_dir; } + this->root = NULL; + this->directories.clear(); /* Set header fields. */ header->header_size = sizeof(*header); diff --git a/stratosphere/fs_mitm/source/fsmitm_romfsbuild.hpp b/stratosphere/fs_mitm/source/fsmitm_romfsbuild.hpp index aafd2685a..004bf5426 100644 --- a/stratosphere/fs_mitm/source/fsmitm_romfsbuild.hpp +++ b/stratosphere/fs_mitm/source/fsmitm_romfsbuild.hpp @@ -1,6 +1,6 @@ #pragma once #include -#include +#include #include "fsmitm_romstorage.hpp" @@ -137,6 +137,13 @@ static_assert(sizeof(RomFSFileEntry) == 0x20, "Incorrect RomFSFileEntry definiti struct RomFSBuildFileContext; +/* Used as comparator for std::map */ +struct build_ctx_cmp { + bool operator()(const char *a, const char *b) const { + return strcmp(a, b) < 0; + } +}; + struct RomFSBuildDirectoryContext { char *path; u32 cur_path_ofs; @@ -146,7 +153,6 @@ struct RomFSBuildDirectoryContext { RomFSBuildDirectoryContext *child = NULL; RomFSBuildDirectoryContext *sibling = NULL; RomFSBuildFileContext *file = NULL; - RomFSBuildDirectoryContext *next = NULL; }; struct RomFSBuildFileContext { @@ -158,7 +164,6 @@ struct RomFSBuildFileContext { u64 size = 0; RomFSBuildDirectoryContext *parent = NULL; RomFSBuildFileContext *sibling = NULL; - RomFSBuildFileContext *next = NULL; RomFSDataSource source{0}; u64 orig_offset = 0; }; @@ -167,7 +172,8 @@ class RomFSBuildContext { private: u64 title_id; RomFSBuildDirectoryContext *root; - RomFSBuildFileContext *files; + std::map directories; + std::map files; u64 num_dirs; u64 num_files; u64 dir_table_size; @@ -185,10 +191,11 @@ class RomFSBuildContext { bool AddDirectory(RomFSBuildDirectoryContext *parent_dir_ctx, RomFSBuildDirectoryContext *dir_ctx, RomFSBuildDirectoryContext **out_dir_ctx); bool AddFile(RomFSBuildDirectoryContext *parent_dir_ctx, RomFSBuildFileContext *file_ctx); public: - RomFSBuildContext(u64 tid) : title_id(tid), root(NULL), files(NULL), num_dirs(0), num_files(0), dir_table_size(0), file_table_size(0), dir_hash_table_size(0), file_hash_table_size(0), file_partition_size(0) { + RomFSBuildContext(u64 tid) : title_id(tid), num_dirs(0), num_files(0), dir_table_size(0), file_table_size(0), dir_hash_table_size(0), file_hash_table_size(0), file_partition_size(0) { this->root = new RomFSBuildDirectoryContext({0}); this->root->path = new char[1]; this->root->path[0] = '\x00'; + this->directories.insert({this->root->path, this->root}); this->num_dirs = 1; this->dir_table_size = 0x18; }