From 7ccf1eaf8ce96b6f26f053da9d68fb276de974af Mon Sep 17 00:00:00 2001
From: Daniel Olson <daniel.olson@me.com>
Date: Thu, 12 Sep 2024 18:42:30 -0600
Subject: [PATCH] Changed the control flow in main to allow a try/catch of a
 bad_alloc exception.

---
 src/cli.hpp  |  2 +-
 src/main.cpp | 31 ++++++++++++++++++++++++++++++-
 2 files changed, 31 insertions(+), 2 deletions(-)

diff --git a/src/cli.hpp b/src/cli.hpp
index e82023e..ce43a3a 100644
--- a/src/cli.hpp
+++ b/src/cli.hpp
@@ -5,7 +5,7 @@
 #ifndef ULTRA_CLI_HPP
 #define ULTRA_CLI_HPP
 
-#define ULTRA_VERSION_STRING "1.0.0"
+#define ULTRA_VERSION_STRING "1.0.2"
 
 
 #include "../lib/CLI11.hpp"
diff --git a/src/main.cpp b/src/main.cpp
index 1f8c0d7..c46aa33 100644
--- a/src/main.cpp
+++ b/src/main.cpp
@@ -2,9 +2,10 @@
 #include "mask.hpp"
 #include "ultra.hpp"
 #include <string>
+#include <new>  // for std::bad_alloc
 
-int main(int argc, const char *argv[]) {
 
+int main_wrapper(int argc, const char * argv[]) {
   // Prepare settings
   Settings *settings = new Settings();
   settings->prepare_settings();
@@ -172,3 +173,31 @@ int main(int argc, const char *argv[]) {
 
   return 0;
 }
+
+
+int main(int argc, const char *argv[]) {
+
+  char *reserve_memory = (char *)malloc(65536);
+
+  try {
+    int r = main_wrapper(argc, argv);
+    return r;
+  }
+  catch (const std::bad_alloc& e) {
+    // This block is executed if memory allocation fails
+    free(reserve_memory); // May be necessary in order to print
+    std::cerr << "Memory allocation failed: " << e.what() << std::endl;
+    std::cerr << "Your model may be too large to fit in memory"  << std::endl;
+    std::cerr << "Try running: ultra --mem <your arguments> to see expected memory usage" << std::endl;
+  }
+  catch (const std::exception& e) {
+    // This block is executed for any other standard exceptions
+    std::cerr << "Standard exception caught: " << e.what() << std::endl;
+  }
+  catch (...) {
+    // This block catches any other non-standard exceptions
+    std::cerr << "Unknown exception caught!" << std::endl;
+  }
+
+  return -1;
+}