feat(shell/lean): add option '-M' to limit amount of memory consumed

src/shell/lean.cpp

@@ -15,6 +15,7 @@ Author: Leonardo de Moura
 #include "util/debug.h"
 #include "util/sstream.h"
 #include "util/interrupt.h"
+#include "util/memory.h"
 #include "util/script_state.h"
 #include "util/thread.h"
 #include "util/thread_script_state.h"
@@ -103,6 +104,10 @@ static void display_help(std::ostream & out) {
     std::cout << "  --to_axiom -X     discard proofs of all theorems after checking them, i.e.,\n";
     std::cout << "                    theorems become axioms after checking\n";
     std::cout << "  --quiet -q        do not print verbose messages\n";
+#if defined(LEAN_TRACK_MEMORY)
+    std::cout << "  --memory=num -M   maximum amount of memory that should be used by Lean ";
+    std::cout << "                    (in megabytes)\n";
+#endif
 #if defined(LEAN_MULTI_THREAD)
     std::cout << "  --server          start Lean in 'server' mode\n";
     std::cout << "  --threads=num -j  number of threads used to process lean files\n";
@@ -143,6 +148,7 @@ static struct option g_long_options[] = {
     {"githash",      no_argument,       0, 'g'},
     {"output",       required_argument, 0, 'o'},
     {"cpp",          required_argument, 0, 'C'},
+    {"memory",       required_argument, 0, 'M'},
     {"trust",        required_argument, 0, 't'},
     {"Trust",        no_argument,       0, 'T'},
     {"discard",      no_argument,       0, 'r'},
@@ -162,14 +168,20 @@ static struct option g_long_options[] = {
     {0, 0, 0, 0}
 };
 
-#define BASIC_OPT_STR "HRXTFC:dD:qrlupgvhk:012t:012o:c:i:"
+#define OPT_STR "HRXTFC:dD:qrlupgvhk:012t:012o:c:i:"
+
+#if defined(LEAN_TRACK_MEMORY)
+#define OPT_STR2 OPT_STR "M:012"
+#else
+#define OPT_STR2 OPT_STR
+#endif
 
 #if defined(LEAN_USE_BOOST) && defined(LEAN_MULTI_THREAD)
-static char const * g_opt_str = BASIC_OPT_STR "Sj:012s:012";
+static char const * g_opt_str = OPT_STR2 "Sj:012s:012";
 #elif !defined(LEAN_USE_BOOST) && defined(LEAN_MULTI_THREAD)
-static char const * g_opt_str = BASIC_OPT_STR "Sj:012";
+static char const * g_opt_str = OPT_STR2 "Sj:012";
 #else
-static char const * g_opt_str = BASIC_OPT_STR;
+static char const * g_opt_str = OPT_STR2;
 #endif
 
 class simple_pos_info_provider : public pos_info_provider {
@@ -391,6 +403,9 @@ int main(int argc, char ** argv) {
         case 'i':
             index_name = optarg;
             gen_index  = true;
+        case 'M':
+            lean::set_max_memory(atoi(optarg)*1024*1024);
+            break;
         case 't':
             trust_lvl = atoi(optarg);
             break;

src/util/exception.cpp

@@ -46,6 +46,14 @@ char const * stack_space_exception::what() const noexcept {
     return buffer.c_str();
 }
 
+char const * memory_exception::what() const noexcept {
+    std::string & buffer = get_g_buffer();
+    std::ostringstream s;
+    s << "excessive memory consumption detected at '" << m_component_name << "' (potential solution: increase memory consumption thresold)";
+    buffer = s.str();
+    return buffer.c_str();
+}
+
 constexpr char const * exception_mt = "exception_mt";
 exception const & to_exception(lua_State * L, int i) {
     return *(*static_cast<exception**>(luaL_checkudata(L, i, exception_mt)));

src/util/exception.h

@@ -26,6 +26,7 @@ public:
     virtual exception * clone() const { return new exception(m_msg); }
     virtual void rethrow() const { throw *this; }
 };
+
 /** \brief Exception produced by a Lean parser. */
 class parser_exception : public exception {
 protected:
@@ -55,6 +56,7 @@ public:
     virtual exception * clone() const { return new interrupted(); }
     virtual void rethrow() const { throw *this; }
 };
+
 class stack_space_exception : public exception {
     std::string m_component_name;
 public:
@@ -63,6 +65,16 @@ public:
     virtual exception * clone() const { return new stack_space_exception(m_component_name.c_str()); }
     virtual void rethrow() const { throw *this; }
 };
+
+class memory_exception : public exception {
+    std::string m_component_name;
+public:
+    memory_exception(char const * component_name):m_component_name(component_name) {}
+    virtual char const * what() const noexcept;
+    virtual exception * clone() const { return new memory_exception(m_component_name.c_str()); }
+    virtual void rethrow() const { throw *this; }
+};
+
 int push_exception(lua_State * L, exception const & e);
 exception const & to_exception(lua_State * L, int i);
 bool is_exception(lua_State * L, int i);

src/util/interrupt.cpp

@@ -7,6 +7,7 @@ Author: Leonardo de Moura
 #include "util/thread.h"
 #include "util/interrupt.h"
 #include "util/exception.h"
+#include "util/memory.h"
 
 namespace lean {
 MK_THREAD_LOCAL_GET(atomic_bool, get_g_interrupt, false);
@@ -30,6 +31,12 @@ void check_interrupted() {
     }
 }
 
+void check_system(char const * component_name) {
+    check_stack(component_name);
+    check_memory(component_name);
+    check_interrupted();
+}
+
 void sleep_for(unsigned ms, unsigned step_ms) {
     if (step_ms == 0)
         step_ms = 1;

src/util/interrupt.h

@@ -30,7 +30,10 @@ bool interrupt_requested();
 */
 void check_interrupted();
 
-inline void check_system(char const * component_name) { check_stack(component_name); check_interrupted(); }
+/**
+   \brief Check system resources: stack, memory, interrupt flag.
+*/
+void check_system(char const * component_name);
 
 constexpr unsigned g_small_sleep = 50;
 

src/util/memory.cpp

@@ -6,14 +6,24 @@ Author: Leonardo de Moura
 */
 #include <new>
 #include <cstdlib>
+#include <iostream>
+#include "util/exception.h"
 
 #if !defined(LEAN_TRACK_MEMORY)
-
 namespace lean {
 size_t     get_allocated_memory() {
     return 0;
 }
 
+void set_max_memory(size_t) {
+    throw exception("Lean was compiled without memory consumption tracking "
+                    "(possible solution: compile using LEAN_TRACK_MEMORY)");
+}
+
+void check_memory(char const * component_name) {
+    // do nothing when LEAN_TRACK_MEMORY is not defined
+}
+
 void * malloc(size_t sz, bool use_ex)  {
     void * r = ::malloc(sz);
     if (r || sz == 0)
@@ -108,9 +118,21 @@ public:
 };
 
 // TODO(Leo): use explicit initialization?
-static alloc_info                          g_global_memory;
+static alloc_info g_global_memory;
+static size_t     g_max_memory = std::numeric_limits<size_t>::max();
 
-size_t     get_allocated_memory() { return g_global_memory.size(); }
+void set_max_memory(size_t max) {
+    g_max_memory = max;
+}
+
+size_t get_allocated_memory() {
+    return g_global_memory.size();
+}
+
+void check_memory(char const * component_name) {
+    if (get_allocated_memory() > g_max_memory)
+        throw memory_exception(component_name);
+}
 
 void * malloc(size_t sz, bool use_ex)  {
     void * r = malloc_core(sz);

src/util/memory.h

@@ -8,6 +8,8 @@ Author: Leonardo de Moura
 #include <cstdlib>
 
 namespace lean {
+void set_max_memory(size_t max);
+void check_memory(char const * component_name);
 size_t get_allocated_memory();
 void * malloc(size_t sz);
 void * realloc(void * ptr, size_t sz);