dfhack/plugins/digFlood.cpp

#include "Core.h"
#include "DataDefs.h"
#include "Export.h"
#include "PluginManager.h"

#include "modules/EventManager.h"
#include "modules/MapCache.h"
#include "modules/Maps.h"

#include "df/coord.h"
#include "df/global_objects.h"
#include "df/job.h"
#include "df/map_block.h"
#include "df/tile_dig_designation.h"
#include "df/world.h"

#include <set>
#include <string>
#include <vector>

using namespace DFHack;
using namespace std;

command_result digFlood (color_ostream &out, std::vector <std::string> & parameters);

DFHACK_PLUGIN("digFlood");

void onDig(color_ostream& out, void* ptr);
void maybeExplore(color_ostream& out, MapExtras::MapCache& cache, df::coord pt, set<df::coord>& jobLocations);
EventManager::EventHandler digHandler(onDig, 0);

//bool enabled = false;
DFHACK_PLUGIN_IS_ENABLED(enabled);
bool digAll = false;
set<string> autodigMaterials;

DFhackCExport command_result plugin_enable(color_ostream& out, bool enable) {
    if (enabled == enable)
        return CR_OK;
    
    enabled = enable;
    if ( enabled ) {
        EventManager::registerListener(EventManager::EventType::JOB_COMPLETED, digHandler, plugin_self);
    } else {
        EventManager::unregisterAll(plugin_self);
    }
    return CR_OK;
}

DFhackCExport command_result plugin_init ( color_ostream &out, std::vector <PluginCommand> &commands)
{
    commands.push_back(PluginCommand(
        "digFlood", "Automatically dig out veins as you discover them.",
        digFlood, false,
        "Example:\n"
        "  digFlood 0\n"
        "    disable plugin\n"
        "  digFlood 1\n"
        "    enable plugin\n"
        "  digFlood 0 MICROCLINE COAL_BITUMINOUS 1\n"
        "    disable plugin and remove microcline and bituminous coal from being monitored, then re-enable plugin"
        "  digFlood 1 MICROCLINE 0 COAL_BITUMINOUS 1\n"
        "    do monitor microcline, don't monitor COAL_BITUMINOUS, then enable plugin\n"
        "  digFlood CLEAR\n"
        "    remove all inorganics from monitoring\n"
        "  digFlood digAll1\n"
        "    enable digAll mode: dig any vein, regardless of the monitor list\n"
        "  digFlood digAll0\n"
        "    disable digAll mode\n"
        "\n"
        "Note that while order matters, multiple commands can be sequenced in one line. It is recommended to alter your dfhack.init file so that you won't have to type in every mineral type you want to dig every time you start the game. Material names are case sensitive.\n"
    ));
    return CR_OK;
}

void onDig(color_ostream& out, void* ptr) {
    CoreSuspender bob;
    df::job* job = (df::job*)ptr;
    if ( job->completion_timer > 0 )
        return;
    
    if ( job->job_type != df::enums::job_type::Dig &&
         job->job_type != df::enums::job_type::CarveUpwardStaircase && 
         job->job_type != df::enums::job_type::CarveDownwardStaircase && 
         job->job_type != df::enums::job_type::CarveUpDownStaircase && 
         job->job_type != df::enums::job_type::CarveRamp && 
         job->job_type != df::enums::job_type::DigChannel )
        return;
    
    set<df::coord> jobLocations;
    for ( df::job_list_link* link = &df::global::world->job_list; link != NULL; link = link->next ) {
        if ( link->item == NULL )
            continue;
        
        if ( link->item->job_type != df::enums::job_type::Dig &&
             link->item->job_type != df::enums::job_type::CarveUpwardStaircase && 
             link->item->job_type != df::enums::job_type::CarveDownwardStaircase && 
             link->item->job_type != df::enums::job_type::CarveUpDownStaircase && 
             link->item->job_type != df::enums::job_type::CarveRamp && 
             link->item->job_type != df::enums::job_type::DigChannel )
            continue;
        
        jobLocations.insert(link->item->pos);
    }
    
    MapExtras::MapCache cache;
    df::coord pos = job->pos;
    for ( int16_t a = -1; a <= 1; a++ ) {
        for ( int16_t b = -1; b <= 1; b++ ) {
            maybeExplore(out, cache, df::coord(pos.x+a,pos.y+b,pos.z), jobLocations);
        }
    }
    cache.trash();
}

void maybeExplore(color_ostream& out, MapExtras::MapCache& cache, df::coord pt, set<df::coord>& jobLocations) {
    if ( !Maps::isValidTilePos(pt) ) {
        return;
    }
    
    df::map_block* block = Maps::getTileBlock(pt);
    if (!block)
        return;
    
    if ( block->designation[pt.x&0xF][pt.y&0xF].bits.hidden )
        return;
    
    df::tiletype type = block->tiletype[pt.x&0xF][pt.y&0xF];
    if ( ENUM_ATTR(tiletype, material, type) != df::enums::tiletype_material::MINERAL )
        return;
    if ( ENUM_ATTR(tiletype, shape, type) != df::enums::tiletype_shape::WALL )
        return;
    
    if ( block->designation[pt.x&0xF][pt.y&0xF].bits.dig != df::enums::tile_dig_designation::No )
        return;
    
    uint32_t xMax,yMax,zMax;
    Maps::getSize(xMax,yMax,zMax);
    if ( pt.x == 0 || pt.y == 0 || pt.x+1 == xMax*16 || pt.y+1 == yMax*16 )
        return;
    if ( jobLocations.find(pt) != jobLocations.end() ) {
        return;
    }
    
    int16_t mat = cache.veinMaterialAt(pt);
    if ( mat == -1 )
        return;
    if ( !digAll ) {
        df::inorganic_raw* inorganic = df::global::world->raws.inorganics[mat];
        if ( autodigMaterials.find(inorganic->id) == autodigMaterials.end() ) {
            return;
        }
    }
    
    block->designation[pt.x&0xF][pt.y&0xF].bits.dig = df::enums::tile_dig_designation::Default;
    block->flags.bits.designated = true;
//    *df::global::process_dig  = true;
//    *df::global::process_jobs = true;
}

command_result digFlood (color_ostream &out, std::vector <std::string> & parameters)
{
    bool adding = true;
    set<string> toAdd, toRemove;
    for ( size_t a = 0; a < parameters.size(); a++ ) {
        int32_t i = (int32_t)strtol(parameters[a].c_str(), NULL, 0);
        if ( i == 0 && parameters[a] == "0" ) {
            plugin_enable(out, false);
            adding = false;
            continue;
        } else if ( i == 1 ) {
            plugin_enable(out, true);
            adding = true;
            continue;
        }
        
        if ( parameters[a] == "CLEAR" )
            autodigMaterials.clear();
        
        if ( parameters[a] == "digAll0" ) {
            digAll = false;
            continue;
        }
        if ( parameters[a] == "digAll1" ) {
            digAll = true;
            continue;
        }
        
        for ( size_t b = 0; b < df::global::world->raws.inorganics.size(); b++ ) {
            df::inorganic_raw* inorganic = df::global::world->raws.inorganics[b];
            if ( parameters[a] == inorganic->id ) {
                if ( adding )
                    toAdd.insert(parameters[a]);
                else
                    toRemove.insert(parameters[a]);
                goto loop;
            }
        }
        
        out.print("Could not find material \"%s\".\n", parameters[a].c_str());
        return CR_WRONG_USAGE;
        
        loop: continue;
    }
    
    autodigMaterials.insert(toAdd.begin(), toAdd.end());
    for ( auto a = toRemove.begin(); a != toRemove.end(); a++ )
        autodigMaterials.erase(*a);
    
    return CR_OK;
}