diff --git a/docs/changelog.txt b/docs/changelog.txt index da5a2d4d7c..2ad0f48ae3 100644 --- a/docs/changelog.txt +++ b/docs/changelog.txt @@ -61,6 +61,7 @@ Template for new versions: ## New Features - `autodump`: new ``undestroy`` option reverts pending item destruction while the game is still paused - `stocks`: the overlay's ``collapse all`` hotkey now toggles, expanding all categories again when everything is collapsed +- `getplants`: new ``--dry-run`` option counts affected plants without designating, and new ``--brewable``/``--edible``/``--oil``/``--cloth``/``--dye`` trait filters restrict the selection by plant use ## Fixes - Fix broken weather lookup in ``World::ReadCurrentWeather`` diff --git a/docs/plugins/getplants.rst b/docs/plugins/getplants.rst index 82616b2358..3ba5a6e386 100644 --- a/docs/plugins/getplants.rst +++ b/docs/plugins/getplants.rst @@ -11,8 +11,9 @@ names. Usage ----- -``getplants [-t|-s|-f]`` - List valid tree/shrub ids, optionally restricted to the specified type. +``getplants [-t|-s|-f] []`` + List valid tree/shrub ids, optionally restricted to the specified type and + traits. ``getplants [ ...] []`` Designate trees/shrubs of the specified types for chopping/gathering. @@ -25,6 +26,10 @@ Examples Gather all plants on the map that yield seeds for farming. ``getplants NETHER_CAP -n 10`` Designate 10 nether cap trees for chopping. +``getplants --brewable -a`` + Designate all shrubs that can be brewed into alcohol. +``getplants --brewable`` + List all valid brewable plant IDs. Options ------- @@ -46,3 +51,22 @@ Options Verbose: Lists the number of (un)designations per plant. ``-n `` Number: Designate up to the specified number of plants of each species. +``-d``, ``--dry-run`` + Dry run: report how many designations would be (un)set without changing + anything. +``--brewable`` + Restrict the selection to plants that can be brewed into alcohol. +``--edible`` + Restrict the selection to plants with parts that can be eaten raw. +``--oil`` + Restrict the selection to plants with seeds, nuts, or fruits that can be + pressed into oil. +``--cloth`` + Restrict the selection to plants that yield thread for cloth. +``--dye`` + Restrict the selection to plants that yield dyes. + +Trait filters can be combined, in which case a plant that matches any of the +specified traits qualifies. They apply to the whole selection, so +``getplants --brewable -x QUARRY_BUSH`` designates every brewable plant except +quarry bushes. diff --git a/plugins/getplants.cpp b/plugins/getplants.cpp index 16fe5019d9..3cbb3381c1 100644 --- a/plugins/getplants.cpp +++ b/plugins/getplants.cpp @@ -50,6 +50,77 @@ enum class selectability { Unselected }; +// Trait categories for the --brewable/--edible/--oil/--cloth/--dye filters. +// Brewable and oil-bearing materials are identified by the reaction products +// that the vanilla brewing and pressing reactions require (DRINK_MAT and +// PRESS_LIQUID_MAT respectively). +enum plant_trait { + trait_none = 0, + trait_brewable = 1 << 0, + trait_edible = 1 << 1, + trait_oil = 1 << 2, + trait_cloth = 1 << 3, + trait_dye = 1 << 4, +}; + +static bool materialHasProduct(const df::material *mat, const std::string &name) { + if (!mat) + return false; + for (const auto &id : mat->reaction_product.id) { + if (*id == name) + return true; + } + return false; +} + +static df::material *plantMatDef(const df::plant_raw *plant, int16_t def) { + const DFHack::MaterialInfo mi(plant->material_defs.type[def], plant->material_defs.idx[def]); + return mi.isValid() ? mi.material : nullptr; +} + +static bool growthMatchesTraits(const df::plant_growth *growth, unsigned traits) { + const DFHack::MaterialInfo mat(growth->mat_type, growth->mat_index); + if (!mat.isValid()) + return false; + return ((traits & trait_brewable) && materialHasProduct(mat.material, "DRINK_MAT")) || + ((traits & trait_oil) && materialHasProduct(mat.material, "PRESS_LIQUID_MAT")) || + ((traits & trait_edible) && mat.material->flags.is_set(material_flags::EDIBLE_RAW)) || + ((traits & trait_dye) && mat.material->flags.is_set(material_flags::IS_DYE)); +} + +// Does the plant have a part matching any of the given trait categories? +static bool plantMatchesTraits(const df::plant_raw *plant, unsigned traits) { + if ((traits & trait_cloth) && plant->flags.is_set(plant_raw_flags::THREAD)) + return true; + + df::material *basic = plantMatDef(plant, plant_material_def::basic_mat); + if (basic && + (((traits & trait_brewable) && materialHasProduct(basic, "DRINK_MAT")) || + ((traits & trait_edible) && basic->flags.is_set(material_flags::EDIBLE_RAW)) || + ((traits & trait_dye) && basic->flags.is_set(material_flags::IS_DYE)))) + return true; + + // Seeds and nuts are milled into paste when their material yields a press liquid + if ((traits & trait_oil) && + materialHasProduct(plantMatDef(plant, plant_material_def::seed), "PRESS_LIQUID_MAT")) + return true; + + if (traits & trait_dye) { + for (int16_t def : {plant_material_def::mill, plant_material_def::extract_vial, + plant_material_def::extract_barrel, plant_material_def::extract_still_vial}) { + df::material *mat = plantMatDef(plant, def); + if (mat && mat->flags.is_set(material_flags::IS_DYE)) + return true; + } + } + + for (auto *growth : plant->growths) { + if (growthMatchesTraits(growth, traits)) + return true; + } + return false; +} + // Determination of whether seeds can be collected is somewhat messy: // - Growths of type SEEDS are collected only if they are edible either raw or cooked. // - Growths of type PLANT_GROWTH are collected provided the STOCKPILE_PLANT_GROWTH @@ -247,13 +318,15 @@ bool picked(const df::plant* plant, int32_t growth_subtype, int32_t growth_densi return false; } -bool designate(color_ostream& out, const df::plant* plant, bool farming) { +bool designate(color_ostream& out, const df::plant* plant, bool farming, bool dry_run) { TRACE(log, out).print("Attempting to designate {} at ({}, {}, {})\n", world->raws.plants.all[plant->material]->id, plant->pos.x, plant->pos.y, plant->pos.z); + auto mark = dry_run ? Designations::canMarkPlant : Designations::markPlant; + if (!farming) { bool istree = (tileMaterial(Maps::getTileBlock(plant->pos)->tiletype[plant->pos.x % 16][plant->pos.y % 16]) == tiletype_material::TREE); if (istree) - return Designations::markPlant(plant); + return mark(plant); } df::plant_raw* plant_raw = world->raws.plants.all[plant->material]; @@ -261,7 +334,7 @@ bool designate(color_ostream& out, const df::plant* plant, bool farming) { if (basic_mat.material->flags.is_set(material_flags::EDIBLE_RAW) || basic_mat.material->flags.is_set(material_flags::EDIBLE_COOKED)) { - return Designations::markPlant(plant); + return mark(plant); } if (plant_raw->flags.is_set(plant_raw_flags::THREAD) || @@ -270,14 +343,14 @@ bool designate(color_ostream& out, const df::plant* plant, bool farming) { plant_raw->flags.is_set(plant_raw_flags::EXTRACT_BARREL) || plant_raw->flags.is_set(plant_raw_flags::EXTRACT_STILL_VIAL)) { if (!farming) { - return Designations::markPlant(plant); + return mark(plant); } } if (basic_mat.material->reaction_product.id.size() > 0 || basic_mat.material->reaction_class.size() > 0) { if (!farming) { - return Designations::markPlant(plant); + return mark(plant); } } @@ -315,7 +388,7 @@ bool designate(color_ostream& out, const df::plant* plant, bool farming) { if ((!farming || seedSource) && ripe(plant->pos.x, plant->pos.y, plant->pos.z, plant_raw->growths[i]->timing_1, plant_raw->growths[i]->timing_2) && !picked(plant, i, plant_raw->growths[i]->density)) - return Designations::markPlant(plant); + return mark(plant); } return false; @@ -326,7 +399,8 @@ command_result df_getplants(color_ostream& out, vector & parameters) { std::vector plantSelections; std::vector collectionCount; set plantNames; - bool deselect = false, exclude = false, treesonly = false, shrubsonly = false, all = false, verbose = false, farming = false; + bool deselect = false, exclude = false, treesonly = false, shrubsonly = false, all = false, verbose = false, farming = false, dry_run = false; + unsigned traits = trait_none; size_t maxCount = 999999; int count = 0; @@ -357,6 +431,18 @@ command_result df_getplants(color_ostream& out, vector & parameters) { verbose = true; else if (parameters[i] == "-f") farming = true; + else if (parameters[i] == "-d" || parameters[i] == "--dry-run") + dry_run = true; + else if (parameters[i] == "--brewable") + traits |= trait_brewable; + else if (parameters[i] == "--edible") + traits |= trait_edible; + else if (parameters[i] == "--oil") + traits |= trait_oil; + else if (parameters[i] == "--cloth") + traits |= trait_cloth; + else if (parameters[i] == "--dye") + traits |= trait_dye; else if (parameters[i] == "-n") { if (parameters.size() > i + 1) { maxCount = atoi(parameters[i + 1].c_str()); @@ -393,14 +479,30 @@ command_result df_getplants(color_ostream& out, vector & parameters) { return CR_WRONG_USAGE; } + std::vector traitMatches(world->raws.plants.all.size(), false); + if (traits != trait_none) { + for (size_t i = 0; i < world->raws.plants.all.size(); i++) + traitMatches[i] = plantMatchesTraits(world->raws.plants.all[i], traits); + } + for (size_t i = 0; i < world->raws.plants.all.size(); i++) { df::plant_raw* plant = world->raws.plants.all[i]; if (all) { plantSelections[i] = selectablePlant(out, plant, farming); + if (traits != trait_none && !traitMatches[i]) + plantSelections[i] = selectability::Nonselectable; } else if (plantNames.find(plant->id) != plantNames.end()) { plantNames.erase(plant->id); plantSelections[i] = selectablePlant(out, plant, farming); + if (traits != trait_none && !traitMatches[i] && + (plantSelections[i] == selectability::Selectable || + plantSelections[i] == selectability::OutOfSeason)) + { + out.printerr("{} does not match the specified trait filters\n", plant->id); + plantSelections[i] = selectability::Nonselectable; + continue; + } switch (plantSelections[i]) { case selectability::Grass: out.printerr("{} is a grass and cannot be gathered\n", plant->id); @@ -447,6 +549,8 @@ command_result df_getplants(color_ostream& out, vector & parameters) { out.print("Valid plant IDs:\n"); for (size_t i = 0; i < world->raws.plants.all.size(); i++) { df::plant_raw* plant = world->raws.plants.all[i]; + if (traits != trait_none && !traitMatches[i]) + continue; switch (selectablePlant(out, plant, farming)) { case selectability::Grass: case selectability::Nonselectable: @@ -506,6 +610,8 @@ command_result df_getplants(color_ostream& out, vector & parameters) { if (!exclude) continue; } + if (traits != trait_none && !traitMatches[mat]) + continue; df::tiletype tt = cur->tiletype[x][y]; df::tiletype_material tile_mat = tileMaterial(tt); if ((treesonly || tt != tiletype::Shrub) && ENUM_ATTR(plant_type, is_shrub, plant->type)) @@ -516,12 +622,12 @@ command_result df_getplants(color_ostream& out, vector & parameters) { continue; if (collectionCount[mat] >= maxCount) continue; - if (deselect && Designations::unmarkPlant(plant)) + if (deselect && (dry_run ? Designations::canUnmarkPlant(plant) : Designations::unmarkPlant(plant))) { collectionCount[mat]++; ++count; } - if (!deselect && designate(out, plant, farming)) + if (!deselect && designate(out, plant, farming, dry_run)) { DEBUG(log, out).print("Designated {} at ({}, {}, {})\n", world->raws.plants.all[mat]->id, plant->pos.x, plant->pos.y, plant->pos.z); collectionCount[mat]++; @@ -531,11 +637,11 @@ command_result df_getplants(color_ostream& out, vector & parameters) { if (count && verbose) { for (size_t i = 0; i < plantSelections.size(); i++) { if (collectionCount[i] > 0) - out.print("Updated {} {} designations.\n", collectionCount[i], world->raws.plants.all[i]->id); + out.print("{} {} {} designations.\n", dry_run ? "Would update" : "Updated", collectionCount[i], world->raws.plants.all[i]->id); } out.print("\n"); } - out.print("Updated {} plant designations.\n", count); + out.print("{} {} plant designations.\n", dry_run ? "Would update" : "Updated", count); return CR_OK; } diff --git a/test/plugins/getplants.lua b/test/plugins/getplants.lua new file mode 100644 index 0000000000..95012e7769 --- /dev/null +++ b/test/plugins/getplants.lua @@ -0,0 +1,123 @@ +config.target = 'getplants' +config.mode = 'fortress' + +local function find_raw(id) + for _, raw in ipairs(df.global.world.raws.plants.all) do + if raw.id == id then return raw end + end +end + +local function run(...) + return dfhack.run_command_silent('getplants', ...) +end + +-- count shrubs carrying a gather designation (designation tile == plant tile +-- for shrubs, unlike trees) +local function count_marked(mat_idx) + local n = 0 + for _, p in ipairs(df.global.world.plants.all) do + if (mat_idx == nil or p.material == mat_idx) and + df.plant_type.attrs[p.type].is_shrub + then + local blk = dfhack.maps.getTileBlock(p.pos) + if blk and blk.designation[p.pos.x % 16][p.pos.y % 16].dig == + df.tile_dig_designation.Default + then + n = n + 1 + end + end + end + return n +end + +local function mat_has_product(mat, product) + if not mat then return false end + for _, id in ipairs(mat.reaction_product.id) do + if id.value == product then return true end + end + return false +end + +-- independent reimplementation of the brewable check used by --brewable +local function raw_is_brewable(raw) + local mi = dfhack.matinfo.decode( + raw.material_defs.type[df.plant_material_def.basic_mat], + raw.material_defs.idx[df.plant_material_def.basic_mat]) + if mat_has_product(mi.material, 'DRINK_MAT') then return true end + for _, g in ipairs(raw.growths) do + local gm = dfhack.matinfo.decode(g.mat_type, g.mat_index) + if mat_has_product(gm.material, 'DRINK_MAT') then return true end + end + return false +end + +function test.dry_run() + local before = count_marked() + local out = run('-s', '-a', '-n', '2', '-d') + local n_dry = tonumber(out:match('Would update (%d+) plant designations')) + expect.ne(nil, n_dry) + -- nothing may have been designated + expect.eq(before, count_marked()) + + -- a real run must update exactly the number the dry run reported + local out_real = run('-s', '-a', '-n', '2') + local n_real = tonumber(out_real:match('Updated (%d+) plant designations')) + expect.eq(n_dry, n_real) + expect.eq(before + n_real, count_marked()) + + -- dry-run clear reports the currently marked count without clearing + local out_cdry = run('-s', '-a', '-c', '-d') + local n_cdry = tonumber(out_cdry:match('Would update (%d+) plant designations')) + expect.eq(count_marked(), n_cdry) + expect.eq(before + n_real, count_marked()) + + local out_clear = run('-s', '-a', '-c') + local n_clear = tonumber(out_clear:match('Updated (%d+) plant designations')) + expect.eq(n_cdry, n_clear) + expect.eq(0, count_marked()) +end + +function test.trait_filters() + -- find a brewable and a non-brewable non-tree non-grass species + local brewable_id, other_id + for _, raw in ipairs(df.global.world.raws.plants.all) do + if not raw.flags.TREE and not raw.flags.GRASS then + if raw_is_brewable(raw) then + brewable_id = brewable_id or raw.id + else + other_id = other_id or raw.id + end + end + end + expect.ne(nil, brewable_id) + expect.ne(nil, other_id) + + -- a matching species is accepted by the filter + local out = run(brewable_id, '--brewable', '-d') + expect.str_find('Would update', out) + + -- a non-matching species is rejected with a specific message + local out2 = run(other_id, '--brewable', '-d') + expect.str_find('does not match the specified trait filters', out2) + + -- every species in the filtered ID listing must be brewable + local listed = run('--brewable') + local listed_count = 0 + for id in listed:gmatch('%* %(%a+%) ([%w_%-]+)') do + listed_count = listed_count + 1 + local raw = find_raw(id) + expect.true_(raw and raw_is_brewable(raw), id) + end + expect.true_(listed_count > 0) +end + +function test.cloth_listing() + local listed = run('--cloth') + local listed_count = 0 + for id in listed:gmatch('%* %(%a+%) ([%w_%-]+)') do + listed_count = listed_count + 1 + local raw = find_raw(id) + expect.true_(raw and raw.flags.THREAD, id) + end + expect.true_(listed_count > 0) +end