propagate luacov debug hook through coroutines

library/Core.cpp

@@ -1186,7 +1186,13 @@ command_result Core::runCommand(color_ostream &con, const std::string &first_, v
                     force = true;
             }
             if (!Lua::Interrupt(force))
-                con.printerr("Failed to register hook - use 'kill-lua force' to force\n");
+            {
+                con.printerr(
+                    "Failed to register hook. This can happen if you have"
+                    " lua profiling or coverage monitoring enabled. Use"
+                    " 'kill-lua force' to force, but this will disable"
+                    " profiling and coverage monitoring.\n");
+            }
         }
         else if (builtin == "script")
         {

library/LuaTools.cpp

@@ -554,6 +554,17 @@ static int dfhack_error(lua_State *L)
     return lua_error(L);
 }
 
+// replaces os.exit with a thrown exception
+static int dfhack_exit_error(lua_State *L)
+{
+    int exit_code = luaL_checkint(L, 1);
+    lua_pop(L, 1);
+    lua_pushfstring(L,
+                    "prevented execution of os.exit(%d); please use error()"
+                    " or qerror() instead.", exit_code);
+    return dfhack_error(L);
+}
+
 static int dfhack_exception_tostring(lua_State *L)
 {
     luaL_checktype(L, 1, LUA_TTABLE);
@@ -775,7 +786,22 @@ static int dfhack_cowrap (lua_State *L) {
     return 1;
 }
 
+static void add_luacov_coroutine_wrapper(lua_State *L)
+{
+    color_ostream *out = Lua::GetOutput(L);
+    if (!Lua::PushModulePublic(*out, L, "luacov_helper", "with_luacov"))
+    {
+        out->printerr("Failed to load luacov_helper.with_luacov.\n");
+        return;
+    }
+    lua_swap(L);
+    if (!Lua::SafeCall(*out, L, 1, 1))
+        out->printerr("Failed to set luacov monitoring wrapper.\n");
+}
+
 lua_State *DFHack::Lua::NewCoroutine(lua_State *L) {
+    if (getenv("DFHACK_ENABLE_LUACOV"))
+        add_luacov_coroutine_wrapper(L);
     lua_State *NL = lua_newthread(L);
     lua_swap(L);
     lua_xmove(L, NL, 1); /* move function from L to NL */
@@ -1328,6 +1354,11 @@ static const luaL_Reg dfhack_coro_funcs[] = {
     { NULL, NULL }
 };
 
+static const luaL_Reg dfhack_os_funcs[] = {
+    { "exit", dfhack_exit_error },
+    { NULL, NULL }
+};
+
 /************
  *  Events  *
  ************/
@@ -1701,6 +1732,11 @@ lua_State *DFHack::Lua::Open(color_ostream &out, lua_State *state)
     luaL_setfuncs(state, dfhack_coro_funcs, 0);
     lua_pop(state, 1);
 
+    // replace some os functions
+    lua_getglobal(state, "os");
+    luaL_setfuncs(state, dfhack_os_funcs, 0);
+    lua_pop(state, 1);
+
     // split the global environment
     lua_newtable(state);
     lua_newtable(state);
@@ -1716,6 +1752,23 @@ lua_State *DFHack::Lua::Open(color_ostream &out, lua_State *state)
     if (IsCoreContext(state))
         Lua::Core::InitCoreContext();
 
+    if (getenv("DFHACK_ENABLE_LUACOV"))
+    {
+        // reads config from .luacov or uses defaults if file doesn't exist.
+        // note that luacov overrides the debug hook installed by
+        // interrupt_init() above.
+        if (PushModulePublic(out, state, "luacov_helper", "init") &&
+            SafeCall(out, state, 0, 0))
+        {
+            out.print("Initialized luacov coverage monitoring\n");
+        }
+        else
+        {
+            out.printerr("Failed to initialize luacov coverage monitoring\n");
+            // non-fatal error
+        }
+    }
+
     // load dfhack.lua
     if (!Require(out, state, "dfhack"))
     {

library/lua/dfhack.lua

@@ -51,6 +51,26 @@ if dfhack.is_core_context then
     SC_UNPAUSED = 8
 end
 
+-- redirect any output explicitly sent to io.stdout or io.stderr to the standard
+-- output sinks. note that although the functions take and pass ...,
+-- only the first string argument will be printed by dfhack.printerr().
+if not io_write_rerouted then
+    local io_file_methods = getmetatable(io.stderr).__index
+    local std_io_file_write = io_file_methods.write
+    io_file_methods.write = function(f, ...)
+            if f == io.stdout then
+                print(...)
+                return f
+            end
+            if f == io.stderr then
+                dfhack.printerr(...)
+                return f
+            end
+            return std_io_file_write(f, ...)
+        end
+    io_write_rerouted = true
+end
+
 -- Error handling
 
 safecall = dfhack.safecall

library/lua/luacov_helper.lua

@@ -0,0 +1,30 @@
+-- Luacov helper functions. Note that this is not a dfhack module since it can't
+-- depend on dfhack.lua.
+
+local runner = require('luacov.runner')
+
+print('runner.debug_hook:', runner.debug_hook)
+
+function init()
+    runner.init()
+    print('** initializing luacov')
+    print('**   set debug hook to:', debug.gethook())
+end
+
+-- Called by LuaTools.cpp to set the debug hook for new threads. We could do
+-- this in C++, but that's complicated and scary.
+function with_luacov(f)
+    print('** wrapping function', f)
+    return function(...)
+        print('** setting debug hook')
+        print('**   was:', debug.gethook())
+        debug.sethook(runner.debug_hook, "l")
+        print('**   is now:', debug.gethook())
+        print('**   running function:', f)
+        print('**   params:', ...)
+        return f(...)
+    end
+end
+
+
+return _ENV