From a9bb11c145928d5ce9d0301842b889d18509afd4 Mon Sep 17 00:00:00 2001
From: lethosor <lethosor@gmail.com>
Date: Fri, 20 Nov 2020 17:57:54 -0500
Subject: [PATCH] Optimize Lua's internal.runCommand() when printing directly
 to the console

This also makes commands run with `run_command()` detect the console properly (notably used by `df2console()`)
---
 library/LuaApi.cpp     | 46 +++++++++++++++++++++++++++++-------------
 library/lua/dfhack.lua | 16 ++++-----------
 2 files changed, 36 insertions(+), 26 deletions(-)

diff --git a/library/LuaApi.cpp b/library/LuaApi.cpp
index 431ab3ae5..ed214397d 100644
--- a/library/LuaApi.cpp
+++ b/library/LuaApi.cpp
@@ -2822,13 +2822,25 @@ static int internal_diffscan(lua_State *L)
 
 static int internal_runCommand(lua_State *L)
 {
-    buffered_color_ostream out;
+    color_ostream *out = NULL;
+    std::unique_ptr<buffered_color_ostream> out_buffer;
     command_result res;
     if (lua_gettop(L) == 0)
     {
         lua_pushstring(L, "");
     }
     int type_1 = lua_type(L, 1);
+    bool use_console = lua_toboolean(L, 2);
+    if (use_console)
+    {
+        out = &Core::getInstance().getConsole();
+    }
+    else
+    {
+        out_buffer.reset(new buffered_color_ostream());
+        out = out_buffer.get();
+    }
+
     if (type_1 == LUA_TTABLE)
     {
         std::string command = "";
@@ -2843,13 +2855,13 @@ static int internal_runCommand(lua_State *L)
             lua_pop(L, 1);  // remove value, leave key
         }
         CoreSuspender suspend;
-        res = Core::getInstance().runCommand(out, command, args);
+        res = Core::getInstance().runCommand(*out, command, args);
     }
     else if (type_1 == LUA_TSTRING)
     {
         std::string command = lua_tostring(L, 1);
         CoreSuspender suspend;
-        res = Core::getInstance().runCommand(out, command);
+        res = Core::getInstance().runCommand(*out, command);
     }
     else
     {
@@ -2857,22 +2869,28 @@ static int internal_runCommand(lua_State *L)
         lua_pushfstring(L, "Expected table, got %s", lua_typename(L, type_1));
         return 2;
     }
-    auto fragments = out.fragments();
+
     lua_newtable(L);
     lua_pushinteger(L, (int)res);
     lua_setfield(L, -2, "status");
-    int i = 1;
-    for (auto iter = fragments.begin(); iter != fragments.end(); iter++, i++)
+
+    if (out_buffer)
     {
-        int color = iter->first;
-        std::string output = iter->second;
-        lua_createtable(L, 2, 0);
-        lua_pushinteger(L, color);
-        lua_rawseti(L, -2, 1);
-        lua_pushstring(L, output.c_str());
-        lua_rawseti(L, -2, 2);
-        lua_rawseti(L, -2, i);
+        auto fragments = out_buffer->fragments();
+        int i = 1;
+        for (auto iter = fragments.begin(); iter != fragments.end(); iter++, i++)
+        {
+            int color = iter->first;
+            std::string output = iter->second;
+            lua_createtable(L, 2, 0);
+            lua_pushinteger(L, color);
+            lua_rawseti(L, -2, 1);
+            lua_pushstring(L, output.c_str());
+            lua_rawseti(L, -2, 2);
+            lua_rawseti(L, -2, i);
+        }
     }
+
     lua_pushvalue(L, -1);
     return 1;
 }
diff --git a/library/lua/dfhack.lua b/library/lua/dfhack.lua
index 73fd651e2..26a23db55 100644
--- a/library/lua/dfhack.lua
+++ b/library/lua/dfhack.lua
@@ -735,8 +735,7 @@ function dfhack.script_help(script_name, extension)
     return help
 end
 
-local function _run_command(...)
-    args = {...}
+local function _run_command(args, use_console)
     if type(args[1]) == 'table' then
         command = args[1]
     elseif #args > 1 and type(args[2]) == 'table' then
@@ -750,11 +749,11 @@ local function _run_command(...)
     else
         error('Invalid arguments')
     end
-    return internal.runCommand(command)
+    return internal.runCommand(command, use_console)
 end
 
 function dfhack.run_command_silent(...)
-    local result = _run_command(...)
+    local result = _run_command({...})
     local output = ""
     for i, f in pairs(result) do
         if type(f) == 'table' then
@@ -765,14 +764,7 @@ function dfhack.run_command_silent(...)
 end
 
 function dfhack.run_command(...)
-    local result = _run_command(...)
-    for i, f in pairs(result) do
-        if type(f) == 'table' then
-            dfhack.color(f[1])
-            dfhack.print(f[2])
-        end
-    end
-    dfhack.color(COLOR_RESET)
+    local result = _run_command({...}, true)
     return result.status
 end