From 1a02abf7b230582aba66ae3d687cd350a97bc694 Mon Sep 17 00:00:00 2001
From: Leonardo de Moura <leonardo@microsoft.com>
Date: Wed, 4 Dec 2013 09:57:33 -0800
Subject: [PATCH] feat(util/script_state): add a lua hook function that checks
 for the interrupt flag

This is a very convenient feature for interrupting non-terminating user scripts.
Before this commit, the user had to manually invoke check_interrupt() in potentially expensive loops. Now, this is not needed anymore.

Remark: we still have to check whether this trick works with LuaJIT or not.

Signed-off-by: Leonardo de Moura <leonardo@microsoft.com>
---
 src/shell/lean.cpp        | 28 ++++++++++++++++++----------
 src/util/script_state.cpp | 18 ++++++++++++++++++
 src/util/script_state.h   |  2 ++
 3 files changed, 38 insertions(+), 10 deletions(-)

diff --git a/src/shell/lean.cpp b/src/shell/lean.cpp
index 528ef4f32..d98ea6534 100644
--- a/src/shell/lean.cpp
+++ b/src/shell/lean.cpp
@@ -7,6 +7,7 @@ Author: Leonardo de Moura
 #include <iostream>
 #include <fstream>
 #include <signal.h>
+#include <cstdlib>
 #include <getopt.h>
 #include "util/stackinfo.h"
 #include "util/debug.h"
@@ -39,11 +40,15 @@ static void display_header(std::ostream & out) {
 static void display_help(std::ostream & out) {
     display_header(out);
     std::cout << "Input format:\n";
-    std::cout << "  --lean        use parser for Lean default input format for files with unknown extension (default)\n";
-    std::cout << "  --lua         use Lua parser for files with unknown extension\n";
+    std::cout << "  --lean         use parser for Lean default input format for files,\n";
+    std::cout << "                 with unknown extension (default)\n";
+    std::cout << "  --lua          use Lua parser for files with unknown extension\n";
     std::cout << "Miscellaneous:\n";
-    std::cout << "  --help -h     display this message\n";
-    std::cout << "  --version -v  display version number\n";
+    std::cout << "  --help -h      display this message\n";
+    std::cout << "  --version -v   display version number\n";
+    std::cout << "  --luahook=num  how often the Lua interpreter checks the interrupted flag,\n";
+    std::cout << "                 it is useful for interrupting non-terminating user scripts,\n";
+    std::cout << "                 0 means 'do not check'.\n";
 }
 
 static char const * get_file_extension(char const * fname) {
@@ -61,10 +66,11 @@ static char const * get_file_extension(char const * fname) {
 }
 
 static struct option g_long_options[] = {
-    {"version", no_argument,       0, 'v'},
-    {"help",    no_argument,       0, 'h'},
-    {"lean",    no_argument,       0, 'l'},
-    {"lua",     no_argument,       0, 'u'},
+    {"version",    no_argument,       0, 'v'},
+    {"help",       no_argument,       0, 'h'},
+    {"lean",       no_argument,       0, 'l'},
+    {"lua",        no_argument,       0, 'u'},
+    {"luahook",    required_argument, 0, 'c'},
     {0, 0, 0, 0}
 };
 
@@ -72,9 +78,8 @@ int main(int argc, char ** argv) {
     lean::save_stack_info();
     lean::register_modules();
     input_kind default_k = input_kind::Lean; // default
-    int optind     = 1;
     while (true) {
-        int c = getopt_long(argc, argv, "vh", g_long_options, &optind);
+        int c = getopt_long(argc, argv, "vhc:012", g_long_options, &optind);
         if (c == -1)
             break; // end of command line
         switch (c) {
@@ -90,6 +95,9 @@ int main(int argc, char ** argv) {
         case 'u':
             default_k = input_kind::Lua;
             break;
+        case 'c':
+            script_state::set_check_interrupt_freq(atoi(optarg));
+            break;
         default:
             std::cerr << "Unknown command line option\n";
             display_help(std::cerr);
diff --git a/src/util/script_state.cpp b/src/util/script_state.cpp
index 4a6a84da9..27bb1b7ba 100644
--- a/src/util/script_state.cpp
+++ b/src/util/script_state.cpp
@@ -36,6 +36,12 @@ void script_state::register_code(char const * code) {
     g_code.push_back(code);
 }
 
+static unsigned g_check_interrupt_freq = 1048576;
+
+void script_state::set_check_interrupt_freq(unsigned count) {
+    g_check_interrupt_freq = count;
+}
+
 void open_extra(lua_State * L);
 
 static char g_weak_ptr_key; // key for Lua registry (used at get_weak_ptr and save_weak_ptr)
@@ -59,6 +65,15 @@ struct script_state::imp {
         lua_settable(m_state, LUA_REGISTRYINDEX);
     }
 
+    static void check_interrupted_hook(lua_State * L, lua_Debug *) {
+        try {
+            check_interrupted();
+        } catch (interrupted & ex) {
+            push_exception(L, ex);
+            lua_error(L);
+        }
+    }
+
     imp() {
         // TODO(Leo) investigate why TCMALLOC + lua_realloc do not work together
         // #ifdef LEAN_USE_LUA_NEWSTATE
@@ -69,6 +84,9 @@ struct script_state::imp {
         #endif
         if (m_state == nullptr)
             throw exception("fail to create Lua interpreter");
+        if (g_check_interrupt_freq > 0) {
+            lua_sethook(m_state, check_interrupted_hook, LUA_MASKCOUNT, g_check_interrupt_freq);
+        }
         luaL_openlibs(m_state);
         open_exception(m_state);
         open_name(m_state);
diff --git a/src/util/script_state.h b/src/util/script_state.h
index 2fd84c998..e755ac968 100644
--- a/src/util/script_state.h
+++ b/src/util/script_state.h
@@ -24,6 +24,8 @@ private:
     lua_State * get_state();
     friend class data_channel;
 public:
+    static void set_check_interrupt_freq(unsigned count);
+
     typedef std::weak_ptr<imp> weak_ref;
 
     script_state();