From 9b63c451b1ce45c8d6423154635b84ca0c7ead85 Mon Sep 17 00:00:00 2001
From: lethosor <lethosor@gmail.com>
Date: Sat, 10 Jun 2017 21:03:42 -0400
Subject: [PATCH] Expose getUnitsInBox to Lua

---
 docs/Lua API.rst   |  7 +++++++
 library/LuaApi.cpp | 31 +++++++++++++++++++++++++++++++
 2 files changed, 38 insertions(+)

diff --git a/docs/Lua API.rst b/docs/Lua API.rst
index 96d459a88..0e242a815 100644
--- a/docs/Lua API.rst	
+++ b/docs/Lua API.rst	
@@ -1059,6 +1059,13 @@ Units module
 
   Returns true *x,y,z* of the unit, or *nil* if invalid; may be not equal to unit.pos if caged.
 
+* ``dfhack.getUnitsInBox(x1,y1,z1,x2,y2,z2[,filter])``
+
+  Returns a table of all units within the specified coordinates. If the ``filter``
+  argument is given, only units where ``filter(unit)`` returns true will be included.
+  Note that ``pos2xyz()`` cannot currently be used to convert coordinate objects to
+  the arguments required by this function.
+
 * ``dfhack.units.getGeneralRef(unit, type)``
 
   Searches for a general_ref with the given type.
diff --git a/library/LuaApi.cpp b/library/LuaApi.cpp
index cb57f8456..af6e03676 100644
--- a/library/LuaApi.cpp
+++ b/library/LuaApi.cpp
@@ -1638,9 +1638,40 @@ static int units_getNoblePositions(lua_State *state)
     return 1;
 }
 
+static int units_getUnitsInBox(lua_State *state)
+{
+    std::vector<df::unit*> units;
+    int x1 = luaL_checkint(state, 1);
+    int y1 = luaL_checkint(state, 2);
+    int z1 = luaL_checkint(state, 3);
+    int x2 = luaL_checkint(state, 4);
+    int y2 = luaL_checkint(state, 5);
+    int z2 = luaL_checkint(state, 6);
+
+    bool ok = Units::getUnitsInBox(units, x1, y1, z1, x2, y2, z2);
+
+    if (ok && !lua_isnone(state, 7))
+    {
+        luaL_checktype(state, 7, LUA_TFUNCTION);
+        units.erase(std::remove_if(units.begin(), units.end(), [&state](df::unit *unit) -> bool {
+            lua_dup(state); // copy function
+            Lua::PushDFObject(state, unit);
+            lua_call(state, 1, 1);
+            bool ret = lua_toboolean(state, -1);
+            lua_pop(state, 1); // remove return value
+            return !ret;
+        }), units.end());
+    }
+
+    Lua::PushVector(state, units);
+    lua_pushboolean(state, ok);
+    return 2;
+}
+
 static const luaL_Reg dfhack_units_funcs[] = {
     { "getPosition", units_getPosition },
     { "getNoblePositions", units_getNoblePositions },
+    { "getUnitsInBox", units_getUnitsInBox },
     { NULL, NULL }
 };