ams.mitm: greatly reduce memory requirements to build romfs

stratosphere/ams_mitm/source/fs_mitm/fsmitm_romfs.cpp

@@ -64,22 +64,168 @@ namespace ams::mitm::fs {
             };
             static_assert(std::is_pod<FileEntry>::value && sizeof(FileEntry) == 0x20);
 
-            constexpr inline DirectoryEntry *GetDirectoryEntry(void *dir_table, u32 offset) {
+            template<typename Entry>
-                return reinterpret_cast<DirectoryEntry *>(reinterpret_cast<uintptr_t>(dir_table) + offset);
+            class TableReader {
+                NON_COPYABLE(TableReader);
+                NON_MOVEABLE(TableReader);
+                private:
+                    static constexpr size_t MaxCachedSize = (1_MB / 4);
+                    static constexpr size_t FallbackCacheSize = 1_KB;
+                private:
+                    ams::fs::IStorage *storage;
+                    size_t offset;
+                    size_t size;
+                    size_t cache_idx;
+                    void *cache;
+                    u8 fallback_cache[FallbackCacheSize];
+                private:
+                    ALWAYS_INLINE void Read(size_t ofs, void *dst, size_t size) {
+                        R_ABORT_UNLESS(this->storage->Read(this->offset + ofs, dst, size));
+                    }
+                    ALWAYS_INLINE void ReloadCacheImpl(size_t idx) {
+                        const size_t rel_ofs = idx * MaxCachedSize;
+                        AMS_ABORT_UNLESS(rel_ofs < this->size);
+                        const size_t new_cache_size = std::min(this->size - rel_ofs, MaxCachedSize);
+                        this->Read(rel_ofs, this->cache, new_cache_size);
+                        this->cache_idx = idx;
                     }
 
-            constexpr inline FileEntry *GetFileEntry(void *file_table, u32 offset) {
+                    ALWAYS_INLINE void ReloadCache(size_t idx) {
-                return reinterpret_cast<FileEntry *>(reinterpret_cast<uintptr_t>(file_table) + offset);
+                        if (this->cache_idx != idx) {
+                            this->ReloadCacheImpl(idx);
+                        }
                     }
 
-            constexpr inline const DirectoryEntry *GetDirectoryEntry(const void *dir_table, u32 offset) {
+                    ALWAYS_INLINE size_t GetCacheIndex(u32 ofs) {
-                return reinterpret_cast<const DirectoryEntry *>(reinterpret_cast<uintptr_t>(dir_table) + offset);
+                        return ofs / MaxCachedSize;
+                    }
+                public:
+                    TableReader(ams::fs::IStorage *s, size_t ofs, size_t sz) : storage(s), offset(ofs), size(sz), cache_idx(0) {
+                        this->cache = std::malloc(std::min(sz, MaxCachedSize));
+                        AMS_ABORT_UNLESS(this->cache != nullptr);
+                        this->ReloadCacheImpl(0);
                     }
 
-            constexpr inline const FileEntry *GetFileEntry(const void *file_table, u32 offset) {
+                    ~TableReader() {
-                return reinterpret_cast<const FileEntry *>(reinterpret_cast<uintptr_t>(file_table) + offset);
+                        std::free(this->cache);
                     }
 
+                    const Entry *GetEntry(u32 entry_offset) {
+                        this->ReloadCache(this->GetCacheIndex(entry_offset));
+
+                        const size_t ofs = entry_offset % MaxCachedSize;
+
+                        const Entry *entry = reinterpret_cast<const Entry *>(reinterpret_cast<uintptr_t>(this->cache) + ofs);
+                        if (AMS_UNLIKELY(this->GetCacheIndex(entry_offset) != this->GetCacheIndex(entry_offset + sizeof(Entry) + entry->name_size + sizeof(u32)))) {
+                            this->Read(entry_offset, this->fallback_cache, std::min(this->size - entry_offset, FallbackCacheSize));
+                            entry = reinterpret_cast<const Entry *>(this->fallback_cache);
+                        }
+                        return entry;
+                    }
+            };
+
+            template<typename Entry>
+            class TableWriter {
+                NON_COPYABLE(TableWriter);
+                NON_MOVEABLE(TableWriter);
+                private:
+                    static constexpr size_t MaxCachedSize = (1_MB / 4);
+                    static constexpr size_t FallbackCacheSize = 1_KB;
+                private:
+                    ::FsFile *file;
+                    size_t offset;
+                    size_t size;
+                    size_t cache_idx;
+                    void *cache;
+                    u8 fallback_cache[FallbackCacheSize];
+                    size_t fallback_cache_entry_offset;
+                    size_t fallback_cache_entry_size;
+                    bool cache_dirty;
+                    bool fallback_cache_dirty;
+                private:
+                    ALWAYS_INLINE void Read(size_t ofs, void *dst, size_t sz) {
+                        u64 read_size;
+                        R_ABORT_UNLESS(fsFileRead(this->file, this->offset + ofs, dst, sz, 0, &read_size));
+                        AMS_ABORT_UNLESS(read_size == sz);
+                    }
+
+                    ALWAYS_INLINE void Write(size_t ofs, const void *src, size_t sz) {
+                        R_ABORT_UNLESS(fsFileWrite(this->file, this->offset + ofs, src, sz, FsWriteOption_None));
+                    }
+
+                    ALWAYS_INLINE void Flush() {
+                        AMS_ABORT_UNLESS(!(this->cache_dirty && this->fallback_cache_dirty));
+
+                        if (this->cache_dirty) {
+                            const size_t ofs = this->cache_idx * MaxCachedSize;
+                            this->Write(ofs, this->cache, std::min(this->size - ofs, MaxCachedSize));
+                            this->cache_dirty = false;
+                        }
+                        if (this->fallback_cache_dirty) {
+                            this->Write(this->fallback_cache_entry_offset, this->fallback_cache, this->fallback_cache_entry_size);
+                            this->fallback_cache_dirty = false;
+                        }
+                    }
+
+                    ALWAYS_INLINE size_t GetCacheIndex(u32 ofs) {
+                        return ofs / MaxCachedSize;
+                    }
+
+                    ALWAYS_INLINE void RefreshCacheImpl() {
+                        const size_t cur_cache = this->cache_idx * MaxCachedSize;
+                        this->Read(cur_cache, this->cache, std::min(this->size - cur_cache, MaxCachedSize));
+                    }
+
+                    ALWAYS_INLINE void RefreshCache(u32 entry_offset) {
+                        if (size_t idx = this->GetCacheIndex(entry_offset); idx != this->cache_idx || this->fallback_cache_dirty) {
+                            this->Flush();
+                            this->cache_idx = idx;
+                            this->RefreshCacheImpl();
+                        }
+                    }
+                public:
+                    TableWriter(::FsFile *f, size_t ofs, size_t sz) : file(f), offset(ofs), size(sz), cache_idx(0), fallback_cache_entry_offset(), fallback_cache_entry_size(), cache_dirty(), fallback_cache_dirty() {
+                        const size_t cache_size = std::min(sz, MaxCachedSize);
+                        this->cache = std::malloc(cache_size);
+                        AMS_ABORT_UNLESS(this->cache != nullptr);
+                        std::memset(this->cache, 0, cache_size);
+                        std::memset(this->fallback_cache, 0, sizeof(this->fallback_cache));
+                        for (size_t cur = 0; cur < this->size; cur += MaxCachedSize) {
+                            this->Write(cur, this->cache, std::min(this->size - cur, MaxCachedSize));
+                        }
+                    }
+
+                    ~TableWriter() {
+                        this->Flush();
+                    }
+
+                    Entry *GetEntry(u32 entry_offset, u32 name_len) {
+                        this->RefreshCache(entry_offset);
+
+                        const size_t ofs = entry_offset % MaxCachedSize;
+
+                        Entry *entry = reinterpret_cast<Entry *>(reinterpret_cast<uintptr_t>(this->cache) + ofs);
+                        if (ofs + sizeof(Entry) + util::AlignUp(name_len, sizeof(u32)) > MaxCachedSize) {
+                            this->Flush();
+
+                            this->fallback_cache_entry_offset = entry_offset;
+                            this->fallback_cache_entry_size   = sizeof(Entry) + util::AlignUp(name_len, sizeof(u32));
+                            this->Read(this->fallback_cache_entry_offset, this->fallback_cache, this->fallback_cache_entry_size);
+
+                            entry = reinterpret_cast<Entry *>(this->fallback_cache);
+                            this->fallback_cache_dirty = true;
+                        } else {
+                            this->cache_dirty = true;
+                        }
+
+                        return entry;
+                    }
+            };
+
+            using DirectoryTableWriter = TableWriter<DirectoryEntry>;
+            using FileTableWriter      = TableWriter<FileEntry>;
+
+
             constexpr inline u32 CalculatePathHash(u32 parent, const char *_path, u32 start, size_t path_len) {
                 const unsigned char *path = reinterpret_cast<const unsigned char *>(_path);
                 u32 hash = parent ^ 123456789;
@@ -214,34 +360,45 @@ namespace ams::mitm::fs {
 
         }
 
-        void Builder::VisitDirectory(BuildDirectoryContext *parent, u32 parent_offset, const void *dir_table, size_t dir_table_size, const void *file_table, size_t file_table_size) {
+        class DirectoryTableReader : public TableReader<DirectoryEntry> {
-            const DirectoryEntry *parent_entry = GetDirectoryEntry(dir_table, parent_offset);
+            public:
-            if (parent_entry->file != EmptyEntry) {
+                DirectoryTableReader(ams::fs::IStorage *s, size_t ofs, size_t sz) : TableReader(s, ofs, sz) { /* ... */ }
-                const FileEntry *cur_file = GetFileEntry(file_table, parent_entry->file);
+        };
-                while (true) {
+
+        class FileTableReader : public TableReader<FileEntry> {
+            public:
+                FileTableReader(ams::fs::IStorage *s, size_t ofs, size_t sz) : TableReader(s, ofs, sz) { /* ... */ }
+        };
+
+        void Builder::VisitDirectory(BuildDirectoryContext *parent, u32 parent_offset, DirectoryTableReader &dir_table, FileTableReader &file_table) {
+            const DirectoryEntry *parent_entry = dir_table.GetEntry(parent_offset);
+
+            u32 cur_file_offset = parent_entry->file;
+            while (cur_file_offset != EmptyEntry) {
+                const FileEntry *cur_file = file_table.GetEntry(cur_file_offset);
+
                 this->AddFile(parent, std::make_unique<BuildFileContext>(cur_file->name, cur_file->name_size, cur_file->size, cur_file->offset, this->cur_source_type));
-                    if (cur_file->sibling == EmptyEntry) {
+
-                        break;
+                cur_file_offset = cur_file->sibling;
             }
-                    cur_file = GetFileEntry(file_table, cur_file->sibling);
+
-                }
-            }
-            if (parent_entry->child != EmptyEntry) {
-                const DirectoryEntry *cur_child = GetDirectoryEntry(dir_table, parent_entry->child);
             u32 cur_child_offset = parent_entry->child;
-                while (true) {
+            while (cur_child_offset != EmptyEntry) {
                 BuildDirectoryContext *real_child = nullptr;
+                u32 next_child_offset = 0;
+                {
+                    const DirectoryEntry *cur_child = dir_table.GetEntry(cur_child_offset);
+
                     this->AddDirectory(&real_child, parent, std::make_unique<BuildDirectoryContext>(cur_child->name, cur_child->name_size));
                     AMS_ABORT_UNLESS(real_child != nullptr);
 
-                    this->VisitDirectory(real_child, cur_child_offset, dir_table, dir_table_size, file_table, file_table_size);
+                    next_child_offset = cur_child->sibling;
+                    __asm__ __volatile__("" ::: "memory");
+                }
 
-                    if (cur_child->sibling == EmptyEntry) {
+                this->VisitDirectory(real_child, cur_child_offset, dir_table, file_table);
-                        break;
+
-                    }
+                cur_child_offset = next_child_offset;
-                    cur_child_offset = cur_child->sibling;
-                    cur_child = GetDirectoryEntry(dir_table, cur_child_offset);
-                }
             }
         }
 
@@ -271,15 +428,11 @@ namespace ams::mitm::fs {
             AMS_ABORT_UNLESS(header.header_size == sizeof(Header));
 
             /* Read tables. */
-            void *tables = std::malloc(header.dir_table_size + header.file_table_size);
+            DirectoryTableReader dir_table(storage, header.dir_table_ofs, header.dir_table_size);
-            ON_SCOPE_EXIT { std::free(tables); };
+            FileTableReader file_table(storage, header.file_table_ofs, header.file_table_size);
-            void *dir_table  = tables;
-            void *file_table = reinterpret_cast<void *>(reinterpret_cast<uintptr_t>(tables) + header.dir_table_size);
-            R_ABORT_UNLESS(storage->Read(header.dir_table_ofs, dir_table, size_t(header.dir_table_size)));
-            R_ABORT_UNLESS(storage->Read(header.file_table_ofs, file_table, size_t(header.file_table_size)));
 
             this->cur_source_type = source_type;
-            this->VisitDirectory(this->root, 0x0, dir_table, size_t(header.dir_table_size), file_table, size_t(header.file_table_size));
+            this->VisitDirectory(this->root, 0x0, dir_table, file_table);
         }
 
         void Builder::Build(std::vector<SourceInfo> *out_infos) {
@@ -372,17 +525,23 @@ namespace ams::mitm::fs {
 
             /* Populate file tables. */
             {
-                void *ft_buf = std::malloc(this->file_table_size);
+                /* Allocate the hash table. */
                 void *fht_buf = std::malloc(this->file_hash_table_size);
-                ON_SCOPE_EXIT { std::free(fht_buf); std::free(ft_buf); };
+                AMS_ABORT_UNLESS(fht_buf != nullptr);
-
                 u32 *file_hash_table = reinterpret_cast<u32 *>(fht_buf);
-                FileEntry *file_table = reinterpret_cast<FileEntry *>(ft_buf);
                 std::memset(file_hash_table, 0xFF, this->file_hash_table_size);
+                ON_SCOPE_EXIT {
+                    R_ABORT_UNLESS(fsFileWrite(&metadata_file, this->dir_hash_table_size + this->dir_table_size, file_hash_table, this->file_hash_table_size, FsWriteOption_None));
+                    std::free(fht_buf);
+                };
+
+                /* Write the file table. */
+                {
+                    FileTableWriter file_table(&metadata_file, this->dir_hash_table_size + this->dir_table_size + this->file_hash_table_size, this->file_table_size);
 
                     for (const auto &it : this->files) {
                         BuildFileContext *cur_file = it.get();
-                    FileEntry *cur_entry = GetFileEntry(file_table, cur_file->entry_offset);
+                        FileEntry *cur_entry = file_table.GetEntry(cur_file->entry_offset, cur_file->path_len);
 
                         /* Set entry fields. */
                         cur_entry->parent = cur_file->parent->entry_offset;
@@ -429,25 +588,28 @@ namespace ams::mitm::fs {
                             AMS_UNREACHABLE_DEFAULT_CASE();
                         }
                     }
-
+                }
-                /* Write to file. */
-                R_ABORT_UNLESS(fsFileWrite(&metadata_file, this->dir_hash_table_size + this->dir_table_size, file_hash_table, this->file_hash_table_size, FsWriteOption_None));
-                R_ABORT_UNLESS(fsFileWrite(&metadata_file, this->dir_hash_table_size + this->dir_table_size + this->file_hash_table_size, file_table, this->file_table_size, FsWriteOption_None));
             }
 
             /* Populate directory tables. */
             {
-                void *dt_buf = std::malloc(this->dir_table_size);
+                /* Allocate the hash table. */
                 void *dht_buf = std::malloc(this->dir_hash_table_size);
-                ON_SCOPE_EXIT { std::free(dht_buf); std::free(dt_buf); };
+                AMS_ABORT_UNLESS(dht_buf != nullptr);
-
                 u32 *dir_hash_table = reinterpret_cast<u32 *>(dht_buf);
-                DirectoryEntry *dir_table = reinterpret_cast<DirectoryEntry *>(dt_buf);
                 std::memset(dir_hash_table, 0xFF, this->dir_hash_table_size);
+                ON_SCOPE_EXIT {
+                    R_ABORT_UNLESS(fsFileWrite(&metadata_file, 0, dir_hash_table, this->dir_hash_table_size, FsWriteOption_None));
+                    std::free(dht_buf);
+                };
+
+                /* Write the file table. */
+                {
+                    DirectoryTableWriter dir_table(&metadata_file, this->dir_hash_table_size, this->dir_table_size);
 
                     for (const auto &it : this->directories) {
                         BuildDirectoryContext *cur_dir = it.get();
-                    DirectoryEntry *cur_entry = GetDirectoryEntry(dir_table, cur_dir->entry_offset);
+                        DirectoryEntry *cur_entry = dir_table.GetEntry(cur_dir->entry_offset, cur_dir->path_len);
 
                         /* Set entry fields. */
                         cur_entry->parent = cur_dir == this->root ? 0 : cur_dir->parent->entry_offset;
@@ -470,10 +632,7 @@ namespace ams::mitm::fs {
                             }
                         }
                     }
-
+                }
-                /* Write to file. */
-                R_ABORT_UNLESS(fsFileWrite(&metadata_file, 0, dir_hash_table, this->dir_hash_table_size, FsWriteOption_None));
-                R_ABORT_UNLESS(fsFileWrite(&metadata_file, this->dir_hash_table_size, dir_table, this->dir_table_size, FsWriteOption_None));
             }
 
             /* Delete maps. */

stratosphere/ams_mitm/source/fs_mitm/fsmitm_romfs.hpp

@@ -199,6 +199,9 @@ namespace ams::mitm::fs::romfs {
         }
     };
 
+    class DirectoryTableReader;
+    class FileTableReader;
+
     struct Builder {
         NON_COPYABLE(Builder);
         NON_MOVEABLE(Builder);
@@ -243,7 +246,7 @@ namespace ams::mitm::fs::romfs {
             DataSourceType cur_source_type;
         private:
             void VisitDirectory(FsFileSystem *fs, BuildDirectoryContext *parent);
-            void VisitDirectory(BuildDirectoryContext *parent, u32 parent_offset, const void *dir_table, size_t dir_table_size, const void *file_table, size_t file_table_size);
+            void VisitDirectory(BuildDirectoryContext *parent, u32 parent_offset, DirectoryTableReader &dir_table, FileTableReader &file_table);
 
             void AddDirectory(BuildDirectoryContext **out, BuildDirectoryContext *parent_ctx, std::unique_ptr<BuildDirectoryContext> file_ctx);
             void AddFile(BuildDirectoryContext *parent_ctx, std::unique_ptr<BuildFileContext> file_ctx);