fs.mitm: Use std::map for RomFS building

stratosphere/fs_mitm/source/fsmitm_romfsbuild.cpp

@@ -10,6 +10,8 @@ void RomFSBuildContext::VisitDirectory(FsFileSystem *filesys, RomFSBuildDirector
     FsDir dir;
     Result rc;
     
+    std::vector<RomFSBuildDirectoryContext *> 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;
+    for (auto &child : child_dirs) {
-    while (cur_child != NULL) {
+        this->VisitDirectory(filesys, child);
-        this->VisitDirectory(filesys, cur_child);
-        cur_child = cur_child->sibling;
     }
 }
 
@@ -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) {    
+bool RomFSBuildContext::AddDirectory(RomFSBuildDirectoryContext *parent_dir_ctx, RomFSBuildDirectoryContext *dir_ctx, RomFSBuildDirectoryContext **out_dir_ctx) {
-    /* Ordered insertion on child. */
+    /* Check whether it's already in the known directories. */
-    int cmp_val;
+    auto existing = this->directories.find(dir_ctx->path);
-    if (parent_dir_ctx->child == NULL || (cmp_val = strcmp(dir_ctx->path, parent_dir_ctx->child->path)) < 0) {
+    if (existing != this->directories.end()) {
-        dir_ctx->sibling = parent_dir_ctx->child;
-        parent_dir_ctx->child = dir_ctx;
-    } else if (cmp_val == 0) {
-        /* This directory already exists! */
         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 */
+    /* Check whether it's already in the known files. */
-    int cmp_val;
+    auto existing = this->files.find(file_ctx->path);
-    if (parent_dir_ctx->file == NULL || (cmp_val = strcmp(file_ctx->path, parent_dir_ctx->file->path)) < 0) {
+    if (existing != this->files.end()) {
-        file_ctx->sibling = parent_dir_ctx->file;
-        parent_dir_ctx->file = file_ctx;
-    } else if (cmp_val == 0) {
-        /* This file already exists! */
         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<RomFSSourceInfo> *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<RomFSSourceInfo> *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) {
+    for (const auto &it : this->directories) {
-        cur_dir->entry_offset = entry_offset;
+        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;
+    for (const auto &it : this->files) {
-    while (cur_file != NULL) {
+        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<RomFSSourceInfo> *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<RomFSSourceInfo> *out_infos) {
                 fatalSimple(0xF601);
         }
         
-        RomFSBuildFileContext *temp = cur_file;
+        delete cur_file->path;
-        cur_file = cur_file->next;
+        delete cur_file;
-        delete temp->path;
-        delete temp;
     }
+    this->files.clear();
         
     /* Populate dir tables. */
-    cur_dir = this->root;
+    for (const auto &it : this->directories) {
-    while (cur_dir != NULL) {
+        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<RomFSSourceInfo> *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;
+        delete cur_dir->path;
-        cur_dir = cur_dir->next;
+        delete cur_dir;
-        delete temp->path;
-        delete temp;
     }
+    this->root = NULL;
+    this->directories.clear();
     
     /* Set header fields. */
     header->header_size = sizeof(*header);

stratosphere/fs_mitm/source/fsmitm_romfsbuild.hpp

@@ -1,6 +1,6 @@
 #pragma once
 #include <switch.h>
-#include <variant>
+#include <map>
 
 #include "fsmitm_romstorage.hpp"
 
@@ -137,6 +137,13 @@ static_assert(sizeof(RomFSFileEntry) == 0x20, "Incorrect RomFSFileEntry definiti
 
 struct RomFSBuildFileContext;
 
+/* Used as comparator for std::map<char *, RomFSBuild*Context> */
+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<char *, RomFSBuildDirectoryContext *, build_ctx_cmp> directories;
+        std::map<char *, RomFSBuildFileContext *, build_ctx_cmp> 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;
         }