Print allocations on --builddir

Standardize how we exit `main` through the `leave` function so that we can make sure to print out our memory stats on every path. To make sure we call `leave` we can decorate `main` as `[[noreturn]]`.
elliotgoodrich · Nov 5, 2024 · bc2eaf5 · bc2eaf5
1 parent 8384534
commit bc2eaf5
Show file tree

Hide file tree

Showing 3 changed files with 65 additions and 45 deletions.
diff --git a/CMakeLists.txt b/CMakeLists.txt
@@ -196,17 +196,28 @@ set_tests_properties(
 # --memory
 if(WIN32)
     add_test(
-        NAME trimja.snapshot.--memory
+        NAME trimja.--memory
         WORKING_DIRECTORY ${CMAKE_CURRENT_SOURCE_DIR}/tests/passthrough/
         COMMAND trimja
         --memory=3
-        --expected expected.ninja
         --affected changed.txt
     )
     set_tests_properties(
-        trimja.snapshot.--memory
+        trimja.--memory
         PROPERTIES FIXTURES_REQUIRED trimja.snapshot.passthrough.fixture
     )
+
+    add_test(
+        NAME trimja.--memory2
+        WORKING_DIRECTORY ${CMAKE_CURRENT_SOURCE_DIR}/tests/builddir/
+        COMMAND trimja
+        --memory=3
+        --builddir
+    )
+    set_tests_properties(
+        trimja.--memory2
+        PROPERTIES FIXTURES_REQUIRED trimja.snapshot.builddir.fixture
+    )
 endif()
 
 # --builddir

diff --git a/README.md b/README.md
@@ -53,6 +53,7 @@ Options:
   -w, --write               overwrite input ninja build file
   --explain                 print why each part of the build file was kept
   --builddir                print the $builddir variable relative to the cwd
+  --memory=N                print memory stats and top N allocating functions
   -h, --help                print help
   -v, --version             print trimja version
 

diff --git a/src/trimja.m.cpp b/src/trimja.m.cpp
@@ -73,7 +73,7 @@ work as expected.
   --builddir                print the $builddir variable relative to the cwd)HELP"
 #if WIN32
     R"HELP(
-  --memory[=N]              print memory stats and top N allocating functions)HELP"
+  --memory=N                print memory stats and top N allocating functions)HELP"
 #endif
     R"HELP(
   -h, --help                print help
@@ -105,13 +105,24 @@ static const option g_longOptions[] = {
     {"affected", required_argument, nullptr, 'a'},
     {"version", no_argument, nullptr, 'v'},
     {"write", no_argument, nullptr, 'w'},
-    {"memory", OPTIONAL_ARG, nullptr, 'm'},
+    {"memory", required_argument, nullptr, 'm'},
     {},
 };
 
-}  // namespace
+std::size_t topAllocatingStacks = 0;
+bool instrumentMemory = false;
+
+[[noreturn]] void leave(int rc) {
+  if (instrumentMemory) {
+    trimja::AllocationProfiler::print(std::cerr, topAllocatingStacks);
+    std::cerr.flush();
+  }
+  std::_Exit(rc);
+};
 
-int main(int argc, char* argv[]) try {
+[[noreturn]] void mainImp(int argc, char* argv[]) try {
+  // Use a separate function for `main` that is declared `noreturn` so we make
+  // sure that we always call `leave` to avoid object destruction.
   using namespace trimja;
 
   struct StdIn {};
@@ -140,8 +151,6 @@ int main(int argc, char* argv[]) try {
   std::filesystem::path ninjaFile = "build.ninja";
   bool explain = false;
   bool builddir = false;
-  std::size_t topAllocatingStacks = 0;
-  bool instrumentMemory = false;
 
   int ch;
   while ((ch = getopt_long(argc, argv, "a:f:ho:vw", g_longOptions, nullptr)) !=
@@ -153,7 +162,7 @@ int main(int argc, char* argv[]) try {
         } else {
           std::cerr << "Cannot specify --affected when - was given"
                     << std::endl;
-          std::_Exit(EXIT_FAILURE);
+          leave(EXIT_FAILURE);
         }
         break;
       case 'b':
@@ -167,18 +176,16 @@ int main(int argc, char* argv[]) try {
         break;
       case 'h':
         std::cout << g_helpText << std::endl;
-        std::_Exit(EXIT_SUCCESS);
+        leave(EXIT_SUCCESS);
       case 'm':
-        if (optarg) {
-          const char* last = optarg + std::strlen(optarg);
-          auto [ptr, ec] = std::from_chars(optarg, last, topAllocatingStacks);
-          if (ec != std::errc{} || ptr != last) {
-            std::string msg;
-            msg = "'";
-            msg += optarg;
-            msg += "' is an invalid value for --memory!";
-            throw std::runtime_error{msg};
-          }
+        const char* last = optarg + std::strlen(optarg);
+        auto [ptr, ec] = std::from_chars(optarg, last, topAllocatingStacks);
+        if (ec != std::errc{} || ptr != last) {
+          std::string msg;
+          msg = "'";
+          msg += optarg;
+          msg += "' is an invalid value for --memory!";
+          throw std::runtime_error{msg};
         }
         instrumentMemory = true;
         AllocationProfiler::start();
@@ -189,33 +196,33 @@ int main(int argc, char* argv[]) try {
         } else if (std::get_if<Write>(&outputFile)) {
           std::cerr << "Cannot specify --output when --write was given"
                     << std::endl;
-          std::_Exit(EXIT_FAILURE);
+          leave(EXIT_FAILURE);
         } else if (std::get_if<Expected>(&outputFile)) {
           std::cerr << "Cannot specify --output when --expected was given"
                     << std::endl;
-          std::_Exit(EXIT_FAILURE);
+          leave(EXIT_FAILURE);
         } else {
           assert(false);
-          std::_Exit(EXIT_FAILURE);
+          leave(EXIT_FAILURE);
         }
         break;
       case 'v':
         std::cout << TRIMJA_VERSION << "" << std::endl;
-        std::_Exit(EXIT_SUCCESS);
+        leave(EXIT_SUCCESS);
       case 'w':
         if (std::get_if<Stdout>(&outputFile)) {
           outputFile.emplace<Write>();
         } else if (std::get_if<std::filesystem::path>(&outputFile)) {
           std::cerr << "Cannot specify --write when --output was given"
                     << std::endl;
-          std::_Exit(EXIT_FAILURE);
+          leave(EXIT_FAILURE);
         } else if (std::get_if<Expected>(&outputFile)) {
           std::cerr << "Cannot specify --write when --expected was given"
                     << std::endl;
-          std::_Exit(EXIT_FAILURE);
+          leave(EXIT_FAILURE);
         } else {
           assert(false);
-          std::_Exit(EXIT_FAILURE);
+          leave(EXIT_FAILURE);
         }
         break;
       case 'x':
@@ -225,22 +232,22 @@ int main(int argc, char* argv[]) try {
         } else if (std::get_if<std::filesystem::path>(&outputFile)) {
           std::cerr << "Cannot specify --expected when --output was given"
                     << std::endl;
-          std::_Exit(EXIT_FAILURE);
+          leave(EXIT_FAILURE);
         } else if (std::get_if<Write>(&outputFile)) {
           std::cerr << "Cannot specify --expected when --write was given"
                     << std::endl;
-          std::_Exit(EXIT_FAILURE);
+          leave(EXIT_FAILURE);
         } else {
           assert(false);
-          std::_Exit(EXIT_FAILURE);
+          leave(EXIT_FAILURE);
         }
         break;
       case '?':
         std::cerr << "Unknown option" << std::endl;
-        std::_Exit(EXIT_FAILURE);
+        leave(EXIT_FAILURE);
       default:
         std::cerr << "Unknown command line parsing error" << std::endl;
-        std::_Exit(EXIT_FAILURE);
+        leave(EXIT_FAILURE);
     }
   }
 
@@ -257,7 +264,7 @@ int main(int argc, char* argv[]) try {
   if (builddir) {
     std::cout << BuildDirUtil::builddir(ninjaFile, ninjaFileContents).string()
               << std::endl;
-    std::_Exit(EXIT_SUCCESS);
+    leave(EXIT_SUCCESS);
   }
 
   std::ifstream affectedFileStream;
@@ -268,7 +275,7 @@ int main(int argc, char* argv[]) try {
           std::cerr << "A list of affected files needs to be supplied with "
                        "either --affected [FILE] or - to read from stdin"
                     << std::endl;
-          std::_Exit(EXIT_FAILURE);
+          leave(EXIT_FAILURE);
         } else if constexpr (std::is_same_v<T, StdIn>) {
           return std::cin;
         } else if constexpr (std::is_same_v<T, std::filesystem::path>) {
@@ -296,13 +303,8 @@ int main(int argc, char* argv[]) try {
   TrimUtil::trim(output, ninjaFile, ninjaFileContents, affected, explain);
   output.flush();
 
-  if (instrumentMemory) {
-    AllocationProfiler::print(std::cout, topAllocatingStacks);
-    std::cout.flush();
-  }
-
   if (!expectedFile.has_value()) {
-    std::_Exit(EXIT_SUCCESS);
+    leave(EXIT_SUCCESS);
   }
 
   // Remove CR characters on Windows from the output so that we can cleanly
@@ -327,16 +329,22 @@ int main(int argc, char* argv[]) try {
               << actual << "---\n"
               << "expected (size " << expected.size() << "):\n"
               << expected << std::endl;
-    std::_Exit(EXIT_FAILURE);
+    leave(EXIT_FAILURE);
   } else {
     std::cout << "Files are equal!\n"
               << "actual:\n"
               << actual << "---\n"
               << "expected:\n"
               << expected << std::endl;
-    std::_Exit(EXIT_SUCCESS);
+    leave(EXIT_SUCCESS);
   }
 } catch (const std::exception& e) {
   std::cout << e.what() << std::endl;
-  std::_Exit(EXIT_FAILURE);
-}
+  exit(EXIT_FAILURE);
+}
+
+}  // namespace
+
+int main(int argc, char* argv[]) {
+  mainImp(argc, argv);
+}