dfhack/plugins/tweak/tweaks/military-assign.h

100 lines
3.8 KiB
C++

using namespace DFHack;
using namespace df::enums;
using Screen::Pen;
struct military_assign_hook : df::viewscreen_layer_militaryst {
typedef df::viewscreen_layer_militaryst interpose_base;
inline bool inPositionsMode() {
return page == Positions && !(in_create_squad || in_new_squad);
}
DEFINE_VMETHOD_INTERPOSE(void, feed, (set<df::interface_key> *input))
{
if (inPositionsMode() && !layer_objects[0]->active)
{
auto pos_list = layer_objects[1];
auto plist = layer_objects[2];
auto &cand = positions.candidates;
// Save the candidate list and cursors
std::vector<df::unit*> copy = cand;
int cursor = plist->getListCursor();
int pos_cursor = pos_list->getListCursor();
INTERPOSE_NEXT(feed)(input);
if (inPositionsMode() && !layer_objects[0]->active)
{
bool is_select = input->count(interface_key::SELECT);
// Resort the candidate list and restore cursor
// on add to squad OR scroll in the position list.
if (!plist->active || is_select)
{
// Since we don't know the actual sorting order, preserve
// the ordering of the items in the list before keypress.
// This does the right thing even if the list was sorted
// with sort-units.
std::set<df::unit*> prev, next;
prev.insert(copy.begin(), copy.end());
next.insert(cand.begin(), cand.end());
std::vector<df::unit*> out;
// (old-before-cursor) (new) |cursor| (old-after-cursor)
for (int i = 0; i < cursor && i < (int)copy.size(); i++)
if (next.count(copy[i])) out.push_back(copy[i]);
for (size_t i = 0; i < cand.size(); i++)
if (!prev.count(cand[i])) out.push_back(cand[i]);
int new_cursor = out.size();
for (int i = cursor; i < (int)copy.size(); i++)
if (next.count(copy[i])) out.push_back(copy[i]);
cand.swap(out);
plist->setListLength(cand.size());
if (new_cursor < (int)cand.size())
plist->setListCursor(new_cursor);
}
// Preserve the position list index on remove from squad
if (pos_list->active && is_select)
pos_list->setListCursor(pos_cursor);
}
}
else
INTERPOSE_NEXT(feed)(input);
}
DEFINE_VMETHOD_INTERPOSE(void, render, ())
{
INTERPOSE_NEXT(render)();
if (inPositionsMode())
{
auto plist = layer_objects[2];
int x1 = plist->getX1(), y1 = plist->getY1();
int x2 = plist->getX2(), y2 = plist->getY2();
int i1 = plist->getFirstVisible(), i2 = plist->getLastVisible();
int si = plist->getListCursor();
for (int y = y1, i = i1; i <= i2 && y <= y2; i++, y++)
{
auto unit = vector_get(positions.candidates, i);
if (!unit || unit->military.squad_id < 0)
continue;
for (int x = x1; x <= x2; x++)
{
Pen cur_tile = Screen::readTile(x, y);
if (!cur_tile.valid()) continue;
cur_tile.fg = (i == si) ? COLOR_BROWN : COLOR_GREEN;
Screen::paintTile(cur_tile, x, y);
}
}
}
}
};
IMPLEMENT_VMETHOD_INTERPOSE(military_assign_hook, feed);
IMPLEMENT_VMETHOD_INTERPOSE(military_assign_hook, render);