From d74f385a68b41dc180c98fc953ce7441527752f6 Mon Sep 17 00:00:00 2001
From: YZBruh <yagizzengin73@gmail.com>
Date: Wed, 13 Aug 2025 09:42:37 +0300
Subject: [PATCH] pmt: dont include new memoryTestFunction

---
 .idea/editor.xml                              |  7 ++--
 CMakeLists.txt                                |  2 +-
 include/PartitionManager/PartitionManager.hpp |  2 +-
 manager.sh                                    |  2 +-
 src/FunctionManager.cpp                       | 17 ++++++----
 src/PartitionManager.cpp                      |  9 ++---
 src/functions/BackupFunction.cpp              | 13 +++----
 src/functions/EraseFunction.cpp               | 11 +++---
 src/functions/FlashFunction.cpp               | 19 ++++++-----
 src/functions/InfoFunction.cpp                |  8 ++---
 src/functions/MemoryTestFunction.cpp          | 21 +++++++-----
 src/functions/PartitionSizeFunction.cpp       | 10 +++---
 src/functions/TypeFunction.cpp                |  8 ++---
 src/functions/functions.hpp                   | 34 ++++++++++++++-----
 14 files changed, 96 insertions(+), 67 deletions(-)

diff --git a/.idea/editor.xml b/.idea/editor.xml
index c8af898..b797d9c 100644
--- a/.idea/editor.xml
+++ b/.idea/editor.xml
@@ -209,7 +209,7 @@
     <option name="/Default/CodeInspection/Highlighting/InspectionSeverities/=CppTooWideScopeInitStatement/@EntryIndexedValue" value="HINT" type="string" />
     <option name="/Default/CodeInspection/Highlighting/InspectionSeverities/=CppTypeAliasNeverUsed/@EntryIndexedValue" value="WARNING" type="string" />
     <option name="/Default/CodeInspection/Highlighting/InspectionSeverities/=CppUninitializedDependentBaseClass/@EntryIndexedValue" value="WARNING" type="string" />
-    <option name="/Default/CodeInspection/Highlighting/InspectionSeverities/=CppUninitializedNonStaticDataMember/@EntryIndexedValue" value="WARNING" type="string" />
+    <option name="/Default/CodeInspection/Highlighting/InspectionSeverities/=CppUninitializedNonStaticDataMember/@EntryIndexedValue" value="DO_NOT_SHOW" type="string" />
     <option name="/Default/CodeInspection/Highlighting/InspectionSeverities/=CppUnionMemberOfReferenceType/@EntryIndexedValue" value="WARNING" type="string" />
     <option name="/Default/CodeInspection/Highlighting/InspectionSeverities/=CppUnmatchedPragmaEndRegionDirective/@EntryIndexedValue" value="WARNING" type="string" />
     <option name="/Default/CodeInspection/Highlighting/InspectionSeverities/=CppUnmatchedPragmaRegionDirective/@EntryIndexedValue" value="WARNING" type="string" />
@@ -268,8 +268,7 @@
     <option name="/Default/CodeStyle/CodeFormatting/CppFormatting/INDENT_CASE_FROM_SWITCH/@EntryValue" value="true" type="bool" />
     <option name="/Default/CodeStyle/CodeFormatting/CppFormatting/INDENT_CLASS_MEMBERS_FROM_ACCESS_SPECIFIERS/@EntryValue" value="true" type="bool" />
     <option name="/Default/CodeStyle/CodeFormatting/CppFormatting/INDENT_COMMENT/@EntryValue" value="true" type="bool" />
-    <option name="/Default/CodeStyle/CodeFormatting/CppFormatting/INDENT_SIZE/@EntryValue" value="4" type="int" />
-    <option name="/Default/CodeStyle/CodeFormatting/CppFormatting/INDENT_STYLE/@EntryValue" value="Tab" type="string" />
+    <option name="/Default/CodeStyle/CodeFormatting/CppFormatting/INDENT_SIZE/@EntryValue" value="2" type="long" />
     <option name="/Default/CodeStyle/CodeFormatting/CppFormatting/INITIALIZER_BRACES/@EntryValue" value="END_OF_LINE_NO_SPACE" type="string" />
     <option name="/Default/CodeStyle/CodeFormatting/CppFormatting/INT_ALIGN_EQ/@EntryValue" value="false" type="bool" />
     <option name="/Default/CodeStyle/CodeFormatting/CppFormatting/INVOCABLE_DECLARATION_BRACES/@EntryValue" value="END_OF_LINE" type="string" />
@@ -327,7 +326,7 @@
     <option name="/Default/CodeStyle/CodeFormatting/CppFormatting/SPACE_WITHIN_TEMPLATE_ARGS/@EntryValue" value="false" type="bool" />
     <option name="/Default/CodeStyle/CodeFormatting/CppFormatting/SPACE_WITHIN_TEMPLATE_PARAMS/@EntryValue" value="false" type="bool" />
     <option name="/Default/CodeStyle/CodeFormatting/CppFormatting/SPECIAL_ELSE_IF_TREATMENT/@EntryValue" value="true" type="bool" />
-    <option name="/Default/CodeStyle/CodeFormatting/CppFormatting/TAB_WIDTH/@EntryValue" value="4" type="int" />
+    <option name="/Default/CodeStyle/CodeFormatting/CppFormatting/TAB_WIDTH/@EntryValue" value="2" type="long" />
     <option name="/Default/CodeStyle/CodeFormatting/CppFormatting/TYPE_DECLARATION_BRACES/@EntryValue" value="END_OF_LINE" type="string" />
     <option name="/Default/CodeStyle/CodeFormatting/CppFormatting/WRAP_AFTER_BINARY_OPSIGN/@EntryValue" value="true" type="bool" />
     <option name="/Default/CodeStyle/CodeFormatting/CppFormatting/WRAP_AFTER_DECLARATION_LPAR/@EntryValue" value="false" type="bool" />
diff --git a/CMakeLists.txt b/CMakeLists.txt
index c420c86..b272648 100755
--- a/CMakeLists.txt
+++ b/CMakeLists.txt
@@ -16,7 +16,7 @@
 
 # Project info
 cmake_minimum_required(VERSION 3.10)
-project(pmt VERSION 1.0.0)
+project(pmt VERSION 1.1.0)
 
 # Set compiler flags
 add_compile_options(-Wall -Werror -Wno-deprecated-declarations)
diff --git a/include/PartitionManager/PartitionManager.hpp b/include/PartitionManager/PartitionManager.hpp
index d091b50..7575931 100644
--- a/include/PartitionManager/PartitionManager.hpp
+++ b/include/PartitionManager/PartitionManager.hpp
@@ -111,7 +111,7 @@ void processCommandLine(std::vector<std::string> &vec1,
                         bool checkForBadUsage = false);
 
 // Setting ups buffer size
-void setupBufferSize(int &size, const std::string &partition);
+void setupBufferSize(uint64_t &size, const std::string &entry);
 
 std::string getLibVersion();
 
diff --git a/manager.sh b/manager.sh
index 51d5656..0efde4e 100644
--- a/manager.sh
+++ b/manager.sh
@@ -29,7 +29,7 @@ checks()
 	echo "Updating repositories and checking required packages..."
 	pkg update &>/dev/null
 	[ ! -f $PREFIX/bin/unzip ] && pkg install -y unzip
-    [ ! -f $PREFIX/bin/wget ] && pkg install -y wget
+  [ ! -f $PREFIX/bin/wget ] && pkg install -y wget
 }
 
 select_variant()
diff --git a/src/FunctionManager.cpp b/src/FunctionManager.cpp
index 638f660..17692c0 100644
--- a/src/FunctionManager.cpp
+++ b/src/FunctionManager.cpp
@@ -39,12 +39,17 @@ std::vector<std::string> splitIfHasDelim(const std::string &s, const char delim,
   return vec;
 }
 
-void setupBufferSize(int &size, const std::string &partition) {
-  if (Variables->PartMap->sizeOf(partition) % size != 0) {
-    print("%sWARNING%s: Specified buffer size is invalid! Using 1 byte as "
-          "buffer size.",
-          YELLOW, STYLE_RESET);
-    size = 1;
+void setupBufferSize(uint64_t &size, const std::string &entry) {
+  if (Variables->PartMap->hasPartition(entry) && Variables->PartMap->sizeOf(entry) % size != 0) {
+    println("%sWARNING%s: Specified buffer size is invalid for %s! Using different buffer size for %s.",
+            YELLOW, STYLE_RESET, entry.data(), entry.data());
+    size = Variables->PartMap->sizeOf(entry) % 4096 == 0 ? 4096 : 1;
+  } else if (Helper::fileIsExists(entry)) {
+    if (Helper::fileSize(entry) % size != 0) {
+      println("%sWARNING%s: Specified buffer size is invalid for %s! using different buffer size for %s.",
+            YELLOW, STYLE_RESET, entry.data(), entry.data());
+      size = Helper::fileSize(entry) % 4096 == 0 ? 4096 : 1;
+    }
   }
 }
 
diff --git a/src/PartitionManager.cpp b/src/PartitionManager.cpp
index feea9ca..5124392 100644
--- a/src/PartitionManager.cpp
+++ b/src/PartitionManager.cpp
@@ -53,8 +53,8 @@ int Main(int argc, char **argv) {
 
     AppMain.fallthrough(true);
     AppMain.set_help_all_flag("--help-all", "Print full help message and exit");
-    AppMain.footer("Partition Manager Tool is written by YZBruh\nThis project "
-                   "licensed under "
+    AppMain.footer("Partition Manager Tool is written by YZBruh\n"
+                   "This project licensed under "
                    "Apache 2.0 license\nReport "
                    "bugs to https://github.com/ShawkTeam/pmt-renovated/issues");
     AppMain
@@ -62,7 +62,7 @@ int Main(int argc, char **argv) {
                     "Set partition search path")
         ->check([&](const std::string &val) {
           if (val.find("/block") == std::string::npos)
-            throw CLI::ValidationError(
+            return std::string(
                 "Partition search path is unexpected! Couldn't find "
                 "'block' in input path!");
           return std::string();
@@ -98,6 +98,7 @@ int Main(int argc, char **argv) {
                                  AppMain);
     FuncManager.registerFunction(std::make_unique<typeFunction>(), AppMain);
     FuncManager.registerFunction(std::make_unique<rebootFunction>(), AppMain);
+    // FuncManager.registerFunction(std::make_unique<memoryTestFunction>(), AppMain);
 
     CLI11_PARSE(AppMain, argc, argv);
 
@@ -115,7 +116,7 @@ int Main(int argc, char **argv) {
                   "(--search-path)");
 
     if (!Helper::hasSuperUser()) {
-      if (!(FuncManager.isUsed("rebootFunction") && Helper::hasAdbPermissions()))
+      if (!((FuncManager.isUsed("rebootFunction") || FuncManager.isUsed("memoryTestFunction")) && Helper::hasAdbPermissions()))
         throw Error(
             "Partition Manager Tool is requires super-user privileges!\n");
     }
diff --git a/src/functions/BackupFunction.cpp b/src/functions/BackupFunction.cpp
index d25845e..e5d0e1a 100644
--- a/src/functions/BackupFunction.cpp
+++ b/src/functions/BackupFunction.cpp
@@ -28,7 +28,7 @@
 
 namespace PartitionManager {
 pair backupFunction::runAsync(const std::string &partitionName,
-                              const std::string &outputName, int bufferSize) {
+                              const std::string &outputName, const uint64_t bufferSize) {
   if (!Variables->PartMap->hasPartition(partitionName))
     return {format("Couldn't find partition: %s", partitionName.data()), false};
 
@@ -53,7 +53,6 @@ pair backupFunction::runAsync(const std::string &partitionName,
                    outputName.data()),
             false};
 
-  setupBufferSize(bufferSize, partitionName);
   LOGN(BFUN, INFO) << "Using buffer size (for back upping " << partitionName
                    << "): " << bufferSize << std::endl;
 
@@ -76,7 +75,7 @@ pair backupFunction::runAsync(const std::string &partitionName,
 
   LOGN(BFUN, INFO) << "Writing partition " << partitionName
                    << " to file: " << outputName << std::endl;
-  auto *buffer = new char[bufferSize];
+  auto *buffer = new(std::nothrow) char[bufferSize];
   collector.delAfterProgress(buffer);
   memset(buffer, 0x00, bufferSize);
 
@@ -94,7 +93,7 @@ pair backupFunction::runAsync(const std::string &partitionName,
                         << outputName
                         << ". Access problems maybe occur in non-root mode"
                         << std::endl;
-  if (!Helper::changeMode(outputName, 0660))
+  if (!Helper::changeMode(outputName, 0664))
     LOGN(BFUN, WARNING) << "Failed to change mode of output file as 660: "
                         << outputName
                         << ". Access problems maybe occur in non-root mode"
@@ -117,7 +116,7 @@ bool backupFunction::init(CLI::App &_app) {
       ->check(CLI::ExistingDirectory);
   cmd->add_option(
       "-b,--buffer-size", bufferSize,
-      "Buffer size for reading partition(s) and writing to file(s)");
+      "Buffer size for reading partition(s) and writing to file(s)")->transform(CLI::AsSizeValue(false))->default_val("4KB");
 
   return true;
 }
@@ -131,13 +130,15 @@ bool backupFunction::run() {
 
   std::vector<std::future<pair>> futures;
   for (size_t i = 0; i < partitions.size(); i++) {
+    uint64_t buf = bufferSize;
     std::string partitionName = partitions[i];
     std::string outputName =
         outputNames.empty() ? partitionName + ".img" : outputNames[i];
     if (!outputDirectory.empty()) outputName.insert(0, outputDirectory + '/');
 
+    setupBufferSize(buf, partitionName);
     futures.push_back(std::async(std::launch::async, runAsync, partitionName,
-                                 outputName, bufferSize));
+                                 outputName, buf));
     LOGN(BFUN, INFO) << "Created thread backup upping " << partitionName
                      << std::endl;
   }
diff --git a/src/functions/EraseFunction.cpp b/src/functions/EraseFunction.cpp
index fc37cd5..04e7f04 100644
--- a/src/functions/EraseFunction.cpp
+++ b/src/functions/EraseFunction.cpp
@@ -25,7 +25,7 @@ Copyright 2025 Yağız Zengin
 #define EFUN "eraseFunction"
 
 namespace PartitionManager {
-pair eraseFunction::runAsync(const std::string &partitionName, int bufferSize) {
+pair eraseFunction::runAsync(const std::string &partitionName, const uint64_t bufferSize) {
   if (!Variables->PartMap->hasPartition(partitionName))
     return {format("Couldn't find partition: %s", partitionName.data()), false};
 
@@ -42,7 +42,6 @@ pair eraseFunction::runAsync(const std::string &partitionName, int bufferSize) {
           false};
   }
 
-  setupBufferSize(bufferSize, partitionName);
   LOGN(EFUN, INFO) << "Using buffer size: " << bufferSize;
 
   // Automatically close file descriptors and delete allocated memories (arrays)
@@ -63,7 +62,7 @@ pair eraseFunction::runAsync(const std::string &partitionName, int bufferSize) {
 
   LOGN(EFUN, INFO) << "Writing zero bytes to partition: " << partitionName
                    << std::endl;
-  auto *buffer = new char[bufferSize];
+  auto *buffer = new(std::nothrow) char[bufferSize];
   collector.delAfterProgress(buffer);
   memset(buffer, 0x00, bufferSize);
 
@@ -94,15 +93,17 @@ bool eraseFunction::init(CLI::App &_app) {
       ->required()
       ->delimiter(',');
   cmd->add_option("-b,--buffer-size", bufferSize,
-                  "Buffer size for writing zero bytes to partition(s)");
+                  "Buffer size for writing zero bytes to partition(s)")->transform(CLI::AsSizeValue(false))->default_val("4KB");
   return true;
 }
 
 bool eraseFunction::run() {
   std::vector<std::future<pair>> futures;
   for (const auto &partitionName : partitions) {
+    uint64_t buf = bufferSize;
+    setupBufferSize(buf, partitionName);
     futures.push_back(
-        std::async(std::launch::async, runAsync, partitionName, bufferSize));
+        std::async(std::launch::async, runAsync, partitionName, buf));
     LOGN(EFUN, INFO) << "Created thread for writing zero bytes to "
                      << partitionName << std::endl;
   }
diff --git a/src/functions/FlashFunction.cpp b/src/functions/FlashFunction.cpp
index 07998f3..79cb25d 100644
--- a/src/functions/FlashFunction.cpp
+++ b/src/functions/FlashFunction.cpp
@@ -26,7 +26,7 @@ Copyright 2025 Yağız Zengin
 
 namespace PartitionManager {
 pair flashFunction::runAsync(const std::string &partitionName,
-                             const std::string &imageName, int bufferSize) {
+                             const std::string &imageName, const uint64_t bufferSize) {
   if (!Helper::fileIsExists(imageName))
     return {format("Couldn't find image file: %s", imageName.data()), false};
   if (!Variables->PartMap->hasPartition(partitionName))
@@ -52,8 +52,7 @@ pair flashFunction::runAsync(const std::string &partitionName,
           false};
   }
 
-  setupBufferSize(bufferSize, imageName);
-  LOGN(FFUN, INFO) << "Using buffer size: " << bufferSize;
+  LOGN(FFUN, INFO) << "Using buffer size: " << bufferSize << std::endl;
 
   // Automatically close file descriptors and delete allocated memories (arrays)
   Helper::garbageCollector collector;
@@ -74,7 +73,7 @@ pair flashFunction::runAsync(const std::string &partitionName,
 
   LOGN(FFUN, INFO) << "Writing image " << imageName
                    << " to partition: " << partitionName << std::endl;
-  auto *buffer = new char[bufferSize];
+  auto *buffer = new(std::nothrow) char[bufferSize];
   collector.delAfterProgress(buffer);
   memset(buffer, 0x00, bufferSize);
 
@@ -101,7 +100,7 @@ bool flashFunction::init(CLI::App &_app) {
       ->required();
   cmd->add_option(
       "-b,--buffer-size", bufferSize,
-      "Buffer size for reading image(s) and writing to partition(s)");
+      "Buffer size for reading image(s) and writing to partition(s)")->transform(CLI::AsSizeValue(false))->default_val("4KB");
   cmd->add_option("-I,--image-directory", imageDirectory,
                   "Directory to find image(s) and flash to partition(s)");
 
@@ -115,13 +114,17 @@ bool flashFunction::run() {
     throw CLI::ValidationError(
         "You must provide an image file(s) as long as the partition name(s)");
 
+  for (size_t i = 0; i < partitions.size(); i++) {
+    if (!imageDirectory.empty()) imageNames[i].insert(0, imageDirectory + '/');
+  }
+
   std::vector<std::future<pair>> futures;
   for (size_t i = 0; i < partitions.size(); i++) {
-    std::string imageName = imageNames[i];
-    if (!imageDirectory.empty()) imageName.insert(0, imageDirectory + '/');
+    uint64_t buf = bufferSize;
 
+    setupBufferSize(buf, imageNames[i]);
     futures.push_back(std::async(std::launch::async, runAsync, partitions[i],
-                                 imageName, bufferSize));
+                                 imageNames[i], bufferSize));
     LOGN(FFUN, INFO) << "Created thread for flashing image to " << partitions[i]
                      << std::endl;
   }
diff --git a/src/functions/InfoFunction.cpp b/src/functions/InfoFunction.cpp
index 5e2f175..6841f09 100644
--- a/src/functions/InfoFunction.cpp
+++ b/src/functions/InfoFunction.cpp
@@ -36,13 +36,13 @@ bool infoFunction::init(CLI::App &_app) {
       ->delimiter(',');
   cmd->add_flag("-J,--json", jsonFormat,
                 "Print info(s) as JSON body. The body of each partition will "
-                "be written separately");
+                "be written separately")->default_val(false);
   cmd->add_option("--json-partition-name", jNamePartition,
-                  "Speficy partition name element for JSON body");
+                  "Speficy partition name element for JSON body")->default_val("name");
   cmd->add_option("--json-size-name", jNameSize,
-                  "Speficy size element name for JSON body");
+                  "Speficy size element name for JSON body")->default_val("size");
   cmd->add_option("--json-logical-name", jNameLogical,
-                  "Speficy logical element name for JSON body");
+                  "Speficy logical element name for JSON body")->default_val("isLogical");
   return true;
 }
 
diff --git a/src/functions/MemoryTestFunction.cpp b/src/functions/MemoryTestFunction.cpp
index 6dcc568..828fa8a 100644
--- a/src/functions/MemoryTestFunction.cpp
+++ b/src/functions/MemoryTestFunction.cpp
@@ -15,9 +15,10 @@ Copyright 2025 Yağız Zengin
 */
 
 #include <chrono>
+#include <random>
 #include <fcntl.h>
-#include <string.h>
-#include <stdlib.h>
+#include <cstring>
+#include <cstdlib>
 #include <unistd.h>
 #include "functions.hpp"
 #include <PartitionManager/PartitionManager.hpp>
@@ -31,10 +32,10 @@ bool memoryTestFunction::init(CLI::App &_app) {
   cmd = _app.add_subcommand("memtest", "Test your write/read speed of device.");
   cmd->add_option("testDirectory", testPath, "Path to test directory")->default_val("/data/local/tmp")->check([&](const std::string &val) {
     if (val != "/data/local/tmp" && !Helper::directoryIsExists(val))
-      return std::string("Couldn't find directory: " + val);
+      return std::string("Couldn't find directory: " + val + ", no root? Try executing in ADB shell.");
     return std::string();
   });
-  cmd->add_option("-b,--buffer-size", bufferSize, "Buffer size for reading partition(s) and writing to file(s)")->transform(CLI::AsSizeValue(false))->default_val("4KB");
+  cmd->add_option("-b,--buffer-size", bufferSize, "Buffer size for reading partition(s) and writing to file(s)")->transform(CLI::AsSizeValue(false))->default_val("4MB");
   cmd->add_option("-s,--file-size", testFileSize, "File size of test file")->transform(CLI::AsSizeValue(false))->default_val("1GB");
   cmd->add_flag("--no-write-test", doNotWriteTest, "Don't write test data to disk")->default_val(false);
   cmd->add_flag("--no-read-test", doNotReadTest, "Don't read test data from disk")->default_val(false);
@@ -53,8 +54,10 @@ bool memoryTestFunction::run() {
   LOGN(MTFUN, INFO) << "Generating random data for testing" << std::endl;
   auto *buffer = new(std::nothrow) char[bufferSize];
   collector.delAfterProgress(buffer);
-  srand(time(nullptr));
-  for (size_t i = 0; i < bufferSize; i++) buffer[i] = rand() % 256;
+  std::mt19937 rng(std::random_device{}());
+  std::uniform_int_distribution<int> dist(0, 255);
+
+  for (size_t i = 0; i < bufferSize; i++) buffer[i] = static_cast<char>(dist(rng));
 
   if (!doNotWriteTest) {
     const int wfd = Helper::openAndAddToCloseList(test, collector, O_WRONLY | O_CREAT | O_TRUNC | O_SYNC, 0644);
@@ -71,7 +74,7 @@ bool memoryTestFunction::run() {
     const auto endWrite = std::chrono::high_resolution_clock::now();
 
     const double writeTime = std::chrono::duration<double>(endWrite - startWrite).count();
-    println("Write speed: %f MB/s", (testFileSize / (1024.0 * 1024.0)) / writeTime);
+    println("Write speed: %f MB/s", (static_cast<double>(testFileSize) / (1024.0 * 1024.0)) / writeTime);
     LOGN(MTFUN, INFO) << "Write test done!" << std::endl;
   }
 
@@ -85,7 +88,7 @@ bool memoryTestFunction::run() {
     const auto endRead = std::chrono::high_resolution_clock::now();
 
     const double read_time = std::chrono::duration<double>(endRead - startRead).count();
-    println("Read speed: %f MB/s", (testFileSize / (1024.0 * 1024.0)) / read_time);
+    println("Read speed: %f MB/s", (static_cast<double>(testFileSize) / (1024.0 * 1024.0)) / read_time);
     LOGN(MTFUN, INFO) << "Read test done!" << std::endl;
   }
 
@@ -97,4 +100,4 @@ bool memoryTestFunction::isUsed() const { return cmd->parsed(); }
 
 const char *memoryTestFunction::name() const { return MTFUN; }
 
-} // namespace PartitionManager
\ No newline at end of file
+} // namespace PartitionManager
diff --git a/src/functions/PartitionSizeFunction.cpp b/src/functions/PartitionSizeFunction.cpp
index 1cb0e41..ce55fe6 100644
--- a/src/functions/PartitionSizeFunction.cpp
+++ b/src/functions/PartitionSizeFunction.cpp
@@ -36,16 +36,16 @@ bool partitionSizeFunction::init(CLI::App &_app) {
       ->required()
       ->delimiter(',');
   cmd->add_flag("--as-byte", asByte,
-                "Tell input size of partition list as byte.");
+                "Tell input size of partition list as byte.")->default_val(false);
   cmd->add_flag("--as-kilobyte", asKiloBytes,
-                "Tell input size of partition list as kilobyte.");
+                "Tell input size of partition list as kilobyte.")->default_val(false);
   cmd->add_flag("--as-megabyte", asMega,
-                "Tell input size of partition list as megabyte.");
+                "Tell input size of partition list as megabyte.")->default_val(false);
   cmd->add_flag("--as-gigabyte", asGiga,
-                "Tell input size of partition list as gigabyte.");
+                "Tell input size of partition list as gigabyte.")->default_val(false);
   cmd->add_flag("--only-size", onlySize,
                 "Tell input size of partition list as not printing multiple "
-                "and partition name.");
+                "and partition name.")->default_val(false);
   return true;
 }
 
diff --git a/src/functions/TypeFunction.cpp b/src/functions/TypeFunction.cpp
index 0734099..6a08148 100644
--- a/src/functions/TypeFunction.cpp
+++ b/src/functions/TypeFunction.cpp
@@ -27,11 +27,11 @@ bool typeFunction::init(CLI::App &_app) {
       ->required()
       ->delimiter(',');
   cmd->add_option("-b,--buffer-size", bufferSize,
-                  "Buffer size for max seek depth");
+                  "Buffer size for max seek depth")->transform(CLI::AsSizeValue(false))->default_val("4KB");
   cmd->add_flag("--only-check-android-magics", onlyCheckAndroidMagics,
-                "Only check Android magic values.");
+                "Only check Android magic values.")->default_val(false);
   cmd->add_flag("--only-check-filesystem-magics", onlyCheckFileSystemMagics,
-                "Only check filesystem magic values.");
+                "Only check filesystem magic values.")->default_val(false);
   return true;
 }
 
@@ -52,7 +52,7 @@ bool typeFunction::run() {
     bool found = false;
     for (const auto &[magic, name] : magics) {
       if (PartitionMap::Extras::hasMagic(
-              magic, bufferSize,
+              magic, static_cast<ssize_t>(bufferSize),
               Helper::fileIsExists(content)
                   ? content
                   : Variables->PartMap->getRealPathOf(content))) {
diff --git a/src/functions/functions.hpp b/src/functions/functions.hpp
index b29aca4..ffd8ef1 100644
--- a/src/functions/functions.hpp
+++ b/src/functions/functions.hpp
@@ -29,7 +29,7 @@ class backupFunction final : public FunctionBase {
 private:
   std::vector<std::string> partitions, outputNames;
   std::string rawPartitions, rawOutputNames, outputDirectory;
-  int bufferSize = 4096;
+  uint64_t bufferSize = 0;
 
 public:
   CLI::App *cmd = nullptr;
@@ -37,7 +37,7 @@ public:
   bool init(CLI::App &_app) override;
   bool run() override;
   static pair runAsync(const std::string &partitionName,
-                       const std::string &outputName, int bufferSize);
+                       const std::string &outputName, uint64_t bufferSize);
 
   [[nodiscard]] bool isUsed() const override;
   [[nodiscard]] const char *name() const override;
@@ -48,7 +48,7 @@ class flashFunction final : public FunctionBase {
 private:
   std::vector<std::string> partitions, imageNames;
   std::string rawPartitions, rawImageNames, imageDirectory;
-  int bufferSize = 4096;
+  uint64_t bufferSize = 0;
 
 public:
   CLI::App *cmd = nullptr;
@@ -56,7 +56,7 @@ public:
   bool init(CLI::App &_app) override;
   bool run() override;
   static pair runAsync(const std::string &partitionName,
-                       const std::string &imageName, int bufferSize);
+                       const std::string &imageName, uint64_t bufferSize);
 
   [[nodiscard]] bool isUsed() const override;
   [[nodiscard]] const char *name() const override;
@@ -66,14 +66,14 @@ public:
 class eraseFunction final : public FunctionBase {
 private:
   std::vector<std::string> partitions;
-  int bufferSize = 4096;
+  uint64_t bufferSize = 0;
 
 public:
   CLI::App *cmd = nullptr;
 
   bool init(CLI::App &_app) override;
   bool run() override;
-  static pair runAsync(const std::string &partitionName, int bufferSize);
+  static pair runAsync(const std::string &partitionName, uint64_t bufferSize);
 
   [[nodiscard]] bool isUsed() const override;
   [[nodiscard]] const char *name() const override;
@@ -100,8 +100,7 @@ public:
 class infoFunction final : public FunctionBase {
 private:
   std::vector<std::string> partitions;
-  std::string jNamePartition = "name", jNameSize = "size",
-              jNameLogical = "isLogical";
+  std::string jNamePartition, jNameSize, jNameLogical;
   bool jsonFormat = false;
 
 public:
@@ -146,7 +145,7 @@ class typeFunction final : public FunctionBase {
 private:
   std::vector<std::string> contents;
   bool onlyCheckAndroidMagics = false, onlyCheckFileSystemMagics = false;
-  int bufferSize = 4096;
+  uint64_t bufferSize = 0;
 
 public:
   CLI::App *cmd = nullptr;
@@ -171,6 +170,23 @@ public:
   [[nodiscard]] bool isUsed() const override;
   [[nodiscard]] const char *name() const override;
 };
+
+class memoryTestFunction final : public FunctionBase {
+private:
+  uint64_t bufferSize = 0, testFileSize = 0;
+  std::string testPath;
+  bool doNotWriteTest = false, doNotReadTest = false;
+
+public:
+  CLI::App *cmd = nullptr;
+
+  bool init(CLI::App &_app) override;
+  bool run() override;
+
+  [[nodiscard]] bool isUsed() const override;
+  [[nodiscard]] const char *name() const override;
+};
+
 } // namespace PartitionManager
 
 #endif // #ifndef FUNCTIONS_HPP