diff --git a/stratosphere/ams_mitm/source/fs_mitm/fsmitm_layered_romfs_storage.hpp b/stratosphere/ams_mitm/source/fs_mitm/fsmitm_layered_romfs_storage.hpp index ad22120c0..2a4863340 100644 --- a/stratosphere/ams_mitm/source/fs_mitm/fsmitm_layered_romfs_storage.hpp +++ b/stratosphere/ams_mitm/source/fs_mitm/fsmitm_layered_romfs_storage.hpp @@ -22,7 +22,7 @@ namespace ams::mitm::fs { class LayeredRomfsStorageImpl { private: - std::vector m_source_infos; + romfs::Builder::SourceInfoVector m_source_infos; std::unique_ptr m_storage_romfs; std::unique_ptr m_file_romfs; os::Event m_initialize_event; diff --git a/stratosphere/ams_mitm/source/fs_mitm/fsmitm_romfs.cpp b/stratosphere/ams_mitm/source/fs_mitm/fsmitm_romfs.cpp index 542d5e202..87daef48b 100644 --- a/stratosphere/ams_mitm/source/fs_mitm/fsmitm_romfs.cpp +++ b/stratosphere/ams_mitm/source/fs_mitm/fsmitm_romfs.cpp @@ -23,6 +23,27 @@ namespace ams::mitm::fs { namespace romfs { + namespace { + + /* TODO: Fancy Dynamic allocation globals. */ + + } + + void *AllocateTracked(AllocationType type, size_t size) { + AMS_UNUSED(type); + + /* TODO: Fancy dynamic allocation with memory stealing from application pool. */ + return std::malloc(size); + } + + void FreeTracked(AllocationType type, void *p, size_t size) { + AMS_UNUSED(type); + AMS_UNUSED(size); + + /* TODO: Fancy dynamic allocation with memory stealing from application pool. */ + return std::free(p); + } + namespace { constexpr u32 EmptyEntry = 0xFFFFFFFF; @@ -71,22 +92,23 @@ namespace ams::mitm::fs { static constexpr size_t MaxCachedSize = (1_MB / 4); private: size_t m_cache_bitsize; + size_t m_cache_size; protected: void *m_cache; protected: DynamicTableCache(size_t sz) { - size_t cache_size = util::CeilingPowerOfTwo(std::min(sz, MaxCachedSize)); - m_cache = std::malloc(cache_size); + m_cache_size = util::CeilingPowerOfTwo(std::min(sz, MaxCachedSize)); + m_cache = AllocateTracked(AllocationType_TableCache, m_cache_size); while (m_cache == nullptr) { - cache_size >>= 1; - AMS_ABORT_UNLESS(cache_size >= 16_KB); - m_cache = std::malloc(cache_size); + m_cache_size >>= 1; + AMS_ABORT_UNLESS(m_cache_size >= 16_KB); + m_cache = AllocateTracked(AllocationType_TableCache, m_cache_size); } - m_cache_bitsize = util::CountTrailingZeros(cache_size); + m_cache_bitsize = util::CountTrailingZeros(m_cache_size); } ~DynamicTableCache() { - std::free(m_cache); + FreeTracked(AllocationType_TableCache, m_cache, m_cache_size); } ALWAYS_INLINE size_t GetCacheSize() const { return static_cast(1) << m_cache_bitsize; } @@ -293,7 +315,7 @@ namespace ams::mitm::fs { } Builder::Builder(ncm::ProgramId pr_id) : m_program_id(pr_id), m_num_dirs(0), m_num_files(0), m_dir_table_size(0), m_file_table_size(0), m_dir_hash_table_size(0), m_file_hash_table_size(0), m_file_partition_size(0) { - auto res = m_directories.emplace(std::make_unique(BuildDirectoryContext::RootTag{})); + auto res = m_directories.emplace(std::unique_ptr(AllocateTyped(AllocationType_BuildDirContext, BuildDirectoryContext::RootTag{}))); AMS_ABORT_UNLESS(res.second); m_root = res.first->get(); m_num_dirs = 1; @@ -347,9 +369,9 @@ namespace ams::mitm::fs { AMS_ABORT_UNLESS(num_child_dirs >= 0); { - BuildDirectoryContext **child_dirs = num_child_dirs != 0 ? reinterpret_cast(std::malloc(sizeof(BuildDirectoryContext *) * num_child_dirs)) : nullptr; + BuildDirectoryContext **child_dirs = num_child_dirs != 0 ? reinterpret_cast(AllocateTracked(AllocationType_DirPointerArray, sizeof(BuildDirectoryContext *) * num_child_dirs)) : nullptr; AMS_ABORT_UNLESS(num_child_dirs == 0 || child_dirs != nullptr); - ON_SCOPE_EXIT { std::free(child_dirs); }; + ON_SCOPE_EXIT { if (child_dirs != nullptr) { FreeTracked(AllocationType_DirPointerArray, child_dirs, sizeof(BuildDirectoryContext *) * num_child_dirs); } }; s64 cur_child_dir_ind = 0; { @@ -368,12 +390,12 @@ namespace ams::mitm::fs { AMS_ABORT_UNLESS(child_dirs != nullptr); BuildDirectoryContext *real_child = nullptr; - this->AddDirectory(std::addressof(real_child), parent, std::make_unique(m_dir_entry.name, strlen(m_dir_entry.name))); + this->AddDirectory(std::addressof(real_child), parent, std::unique_ptr(AllocateTyped(AllocationType_BuildDirContext, m_dir_entry.name, strlen(m_dir_entry.name)))); AMS_ABORT_UNLESS(real_child != nullptr); child_dirs[cur_child_dir_ind++] = real_child; AMS_ABORT_UNLESS(cur_child_dir_ind <= num_child_dirs); } else /* if (m_dir_entry.type == FsDirEntryType_File) */ { - this->AddFile(parent, std::make_unique(m_dir_entry.name, strlen(m_dir_entry.name), m_dir_entry.file_size, 0, m_cur_source_type)); + this->AddFile(parent, std::unique_ptr(AllocateTyped(AllocationType_BuildFileContext, m_dir_entry.name, strlen(m_dir_entry.name), m_dir_entry.file_size, 0, m_cur_source_type))); } } } @@ -403,7 +425,7 @@ namespace ams::mitm::fs { while (cur_file_offset != EmptyEntry) { const FileEntry *cur_file = file_table.GetEntry(cur_file_offset); - this->AddFile(parent, std::make_unique(cur_file->name, cur_file->name_size, cur_file->size, cur_file->offset, m_cur_source_type)); + this->AddFile(parent, std::unique_ptr(AllocateTyped(AllocationType_BuildFileContext, cur_file->name, cur_file->name_size, cur_file->size, cur_file->offset, m_cur_source_type))); cur_file_offset = cur_file->sibling; } @@ -415,7 +437,7 @@ namespace ams::mitm::fs { { const DirectoryEntry *cur_child = dir_table.GetEntry(cur_child_offset); - this->AddDirectory(std::addressof(real_child), parent, std::make_unique(cur_child->name, cur_child->name_size)); + this->AddDirectory(std::addressof(real_child), parent, std::unique_ptr(AllocateTyped(AllocationType_BuildDirContext, cur_child->name, cur_child->name_size))); AMS_ABORT_UNLESS(real_child != nullptr); next_child_offset = cur_child->sibling; @@ -438,7 +460,7 @@ namespace ams::mitm::fs { /* If there is no romfs folder on the SD, don't bother continuing. */ { FsDir dir; - if (R_FAILED(mitm::fs::OpenAtmosphereRomfsDirectory(std::addressof(dir), m_program_id, m_root->path.get(), OpenDirectoryMode_Directory, std::addressof(sd_filesystem)))) { + if (R_FAILED(mitm::fs::OpenAtmosphereRomfsDirectory(std::addressof(dir), m_program_id, m_root->path, OpenDirectoryMode_Directory, std::addressof(sd_filesystem)))) { return; } fsDirClose(std::addressof(dir)); @@ -461,7 +483,7 @@ namespace ams::mitm::fs { this->VisitDirectory(m_root, 0x0, dir_table, file_table); } - void Builder::Build(std::vector *out_infos) { + void Builder::Build(SourceInfoVector *out_infos) { /* Clear output. */ out_infos->clear(); @@ -477,7 +499,7 @@ namespace ams::mitm::fs { m_file_hash_table_size = sizeof(u32) * num_file_hash_table_entries; /* Allocate metadata, make pointers. */ - Header *header = reinterpret_cast
(std::malloc(sizeof(Header))); + Header *header = reinterpret_cast
(AllocateTracked(AllocationType_Memory, sizeof(Header))); std::memset(header, 0x00, sizeof(*header)); /* Open metadata file. */ @@ -552,13 +574,13 @@ namespace ams::mitm::fs { /* Set all files' hash value = hash index. */ for (const auto &it : m_files) { BuildFileContext *cur_file = it.get(); - cur_file->hash_value = CalculatePathHash(cur_file->parent->entry_offset, cur_file->path.get(), 0, cur_file->path_len) % num_file_hash_table_entries; + cur_file->hash_value = CalculatePathHash(cur_file->parent->entry_offset, cur_file->path, 0, cur_file->path_len) % num_file_hash_table_entries; } /* Set all directories' hash value = hash index. */ for (const auto &it : m_directories) { BuildDirectoryContext *cur_dir = it.get(); - cur_dir->hash_value = CalculatePathHash(cur_dir == m_root ? 0 : cur_dir->parent->entry_offset, cur_dir->path.get(), 0, cur_dir->path_len) % num_dir_hash_table_entries; + cur_dir->hash_value = CalculatePathHash(cur_dir == m_root ? 0 : cur_dir->parent->entry_offset, cur_dir->path, 0, cur_dir->path_len) % num_dir_hash_table_entries; } /* Write hash tables. */ @@ -661,7 +683,7 @@ namespace ams::mitm::fs { const u32 name_size = cur_file->path_len; cur_entry->name_size = name_size; if (name_size) { - std::memcpy(cur_entry->name, cur_file->path.get(), name_size); + std::memcpy(cur_entry->name, cur_file->path, name_size); for (size_t i = name_size; i < util::AlignUp(name_size, 4); i++) { cur_entry->name[i] = 0; } @@ -688,9 +710,10 @@ namespace ams::mitm::fs { AMS_ABORT_UNLESS(path_needed_size <= sizeof(full_path)); cur_file->GetPath(full_path); - cur_file->path.reset(); + FreeTracked(AllocationType_FileName, cur_file->path, cur_file->path_len + 1); + cur_file->path = nullptr; - char *new_path = new char[path_needed_size]; + char *new_path = static_cast(AllocateTracked(AllocationType_FullPath, path_needed_size)); std::memcpy(new_path, full_path, path_needed_size); out_infos->emplace_back(cur_file->offset + FilePartitionOffset, cur_file->size, cur_file->source_type, new_path); } @@ -719,7 +742,7 @@ namespace ams::mitm::fs { const u32 name_size = cur_dir->path_len; cur_entry->name_size = name_size; if (name_size) { - std::memcpy(cur_entry->name, cur_dir->path.get(), name_size); + std::memcpy(cur_entry->name, cur_dir->path, name_size); for (size_t i = name_size; i < util::AlignUp(name_size, 4); i++) { cur_entry->name[i] = 0; } diff --git a/stratosphere/ams_mitm/source/fs_mitm/fsmitm_romfs.hpp b/stratosphere/ams_mitm/source/fs_mitm/fsmitm_romfs.hpp index 98e74c430..69d61686d 100644 --- a/stratosphere/ams_mitm/source/fs_mitm/fsmitm_romfs.hpp +++ b/stratosphere/ams_mitm/source/fs_mitm/fsmitm_romfs.hpp @@ -27,6 +27,52 @@ namespace ams::mitm::fs::romfs { Memory, }; + enum AllocationType { + AllocationType_FileName, + AllocationType_DirName, + AllocationType_FullPath, + AllocationType_SourceInfo, + AllocationType_BuildFileContext, + AllocationType_BuildDirContext, + AllocationType_TableCache, + AllocationType_DirPointerArray, + AllocationType_DirContextSet, + AllocationType_FileContextSet, + AllocationType_Memory, + + AllocationType_Count, + }; + + void *AllocateTracked(AllocationType type, size_t size); + void FreeTracked(AllocationType type, void *p, size_t size); + + template + T *AllocateTyped(AllocationType type, Args &&... args) { + void *mem = AllocateTracked(type, sizeof(T)); + return std::construct_at(static_cast(mem), std::forward(args)...); + } + + template + class TrackedAllocator { + public: + using value_type = T; + + template + struct rebind { + using other = TrackedAllocator; + }; + public: + TrackedAllocator() = default; + + T *allocate(size_t n) { + return static_cast(AllocateTracked(AllocType, sizeof(T) * n)); + } + + void deallocate(T *p, size_t n) { + FreeTracked(AllocType, p, sizeof(T) * n); + } + }; + struct SourceInfo { s64 virtual_offset; s64 size; @@ -89,10 +135,10 @@ namespace ams::mitm::fs::romfs { delete this->metadata_source_info.file; break; case DataSourceType::LooseSdFile: - delete[] this->loose_source_info.path; + FreeTracked(AllocationType_FullPath, this->loose_source_info.path, std::strlen(this->loose_source_info.path) + 1); break; case DataSourceType::Memory: - std::free(static_cast(this->memory_source_info.data)); + FreeTracked(AllocationType_Memory, this->memory_source_info.data, this->size); break; AMS_UNREACHABLE_DEFAULT_CASE(); } @@ -113,7 +159,7 @@ namespace ams::mitm::fs::romfs { NON_COPYABLE(BuildDirectoryContext); NON_MOVEABLE(BuildDirectoryContext); - std::unique_ptr path; + char *path; union { BuildDirectoryContext *parent; }; @@ -139,16 +185,28 @@ namespace ams::mitm::fs::romfs { struct RootTag{}; BuildDirectoryContext(RootTag) : parent(nullptr), child(nullptr), sibling(nullptr), file(nullptr), path_len(0), entry_offset(0), hash_value(0xFFFFFFFF) { - this->path = std::make_unique(1); + this->path = static_cast(AllocateTracked(AllocationType_DirName, 1)); + this->path[0] = '\x00'; } BuildDirectoryContext(const char *entry_name, size_t entry_name_len) : parent(nullptr), child(nullptr), sibling(nullptr), file(nullptr), entry_offset(0) { this->path_len = entry_name_len; - this->path = std::unique_ptr(new char[this->path_len + 1]); - std::memcpy(this->path.get(), entry_name, entry_name_len); + this->path = static_cast(AllocateTracked(AllocationType_DirName, this->path_len + 1)); + std::memcpy(this->path, entry_name, entry_name_len); this->path[this->path_len] = '\x00'; } + ~BuildDirectoryContext() { + if (this->path != nullptr) { + FreeTracked(AllocationType_DirName, this->path, this->path_len + 1); + this->path = nullptr; + } + } + + void operator delete(void *p) { + FreeTracked(AllocationType_BuildDirContext, p, sizeof(BuildDirectoryContext)); + } + size_t GetPathLength() const { if (this->parent == nullptr) { return 0; @@ -165,7 +223,7 @@ namespace ams::mitm::fs::romfs { const size_t parent_len = this->parent->GetPath(dst); dst[parent_len] = '/'; - std::memcpy(dst + parent_len + 1, this->path.get(), this->path_len); + std::memcpy(dst + parent_len + 1, this->path, this->path_len); dst[parent_len + 1 + this->path_len] = '\x00'; return parent_len + 1 + this->path_len; } @@ -187,7 +245,7 @@ namespace ams::mitm::fs::romfs { NON_COPYABLE(BuildFileContext); NON_MOVEABLE(BuildFileContext); - std::unique_ptr path; + char *path; BuildDirectoryContext *parent; union { BuildFileContext *sibling; @@ -203,11 +261,22 @@ namespace ams::mitm::fs::romfs { BuildFileContext(const char *entry_name, size_t entry_name_len, s64 sz, s64 o_o, DataSourceType type) : parent(nullptr), sibling(nullptr), offset(0), size(sz), orig_offset(o_o), entry_offset(0), hash_value(0xFFFFFFFF), source_type(type) { this->path_len = entry_name_len; - this->path = std::unique_ptr(new char[this->path_len + 1]); - std::memcpy(this->path.get(), entry_name, entry_name_len); + this->path = static_cast(AllocateTracked(AllocationType_FileName, this->path_len + 1)); + std::memcpy(this->path, entry_name, entry_name_len); this->path[this->path_len] = 0; } + ~BuildFileContext() { + if (this->path != nullptr) { + FreeTracked(AllocationType_FileName, this->path, this->path_len + 1); + this->path = nullptr; + } + } + + void operator delete(void *p) { + FreeTracked(AllocationType_BuildFileContext, p, sizeof(BuildFileContext)); + } + size_t GetPathLength() const { if (this->parent == nullptr) { return 0; @@ -224,7 +293,7 @@ namespace ams::mitm::fs::romfs { const size_t parent_len = this->parent->GetPath(dst); dst[parent_len] = '/'; - std::memcpy(dst + parent_len + 1, this->path.get(), this->path_len); + std::memcpy(dst + parent_len + 1, this->path, this->path_len); dst[parent_len + 1 + this->path_len] = '\x00'; return parent_len + 1 + this->path_len; } @@ -248,6 +317,8 @@ namespace ams::mitm::fs::romfs { class Builder { NON_COPYABLE(Builder); NON_MOVEABLE(Builder); + public: + using SourceInfoVector = std::vector>; private: template struct Comparator { @@ -270,13 +341,13 @@ namespace ams::mitm::fs::romfs { } }; - template - using ContextSet = std::set, Comparator>; + template + using ContextSet = std::set, Comparator, TrackedAllocator>>; private: ncm::ProgramId m_program_id; BuildDirectoryContext *m_root; - ContextSet m_directories; - ContextSet m_files; + ContextSet m_directories; + ContextSet m_files; size_t m_num_dirs; size_t m_num_files; size_t m_dir_table_size; @@ -299,7 +370,7 @@ namespace ams::mitm::fs::romfs { void AddSdFiles(); void AddStorageFiles(ams::fs::IStorage *storage, DataSourceType source_type); - void Build(std::vector *out_infos); + void Build(SourceInfoVector *out_infos); }; }