Take cargo refit into account in build vehicle window capacity sort modes

src/articulated_vehicles.cpp

@@ -8,6 +8,7 @@
 /** @file articulated_vehicles.cpp Implementation of articulated vehicles. */
 
 #include "stdafx.h"
+#include "articulated_vehicles.h"
 #include "core/bitmath_func.hpp"
 #include "core/random_func.hpp"
 #include "train.h"
@@ -141,15 +142,23 @@ void GetArticulatedPartsEngineIDs(EngineID engine_type, bool purchase_window, st
  * Returns the default (non-refitted) capacity of a specific EngineID.
  * @param engine the EngineID of interest
  * @param cargo_type returns the default cargo type, if needed
+ * @param attempt_refit cargo ID to attempt to use
  * @return capacity
  */
-static inline uint16_t GetVehicleDefaultCapacity(EngineID engine, CargoID *cargo_type)
+static inline uint16_t GetVehicleDefaultCapacity(EngineID engine, CargoID *cargo_type, CargoID attempt_refit = INVALID_CARGO)
 {
 	const Engine *e = Engine::Get(engine);
-	CargoID cargo = (e->CanCarryCargo() ? e->GetDefaultCargoType() : INVALID_CARGO);
+	CargoID cargo = INVALID_CARGO;
+	if (e->CanCarryCargo()) {
+		if (attempt_refit != INVALID_CARGO && HasBit(e->info.refit_mask, attempt_refit)) {
+			cargo = attempt_refit;
+		} else {
+			cargo = e->GetDefaultCargoType();
+		}
+	}
 	if (cargo_type != nullptr) *cargo_type = cargo;
 	if (cargo == INVALID_CARGO) return 0;
-	return e->GetDisplayDefaultCapacity();
+	return e->GetDisplayDefaultCapacity(nullptr, cargo);
 }
 
 /**
@@ -175,16 +184,20 @@ static inline CargoTypes GetAvailableVehicleCargoTypes(EngineID engine, bool inc
 /**
  * Get the capacity of the parts of a given engine.
  * @param engine The engine to get the capacities from.
+ * @param attempt_refit Attempt to get capacity when refitting to this cargo.
  * @return The cargo capacities.
  */
-CargoArray GetCapacityOfArticulatedParts(EngineID engine)
+CargoArray GetCapacityOfArticulatedParts(EngineID engine, CargoID attempt_refit)
 {
 	CargoArray capacity{};
 	const Engine *e = Engine::Get(engine);
 
-	CargoID cargo_type;
+	auto get_engine_cargo = [&capacity, attempt_refit](EngineID eng) {
-	uint16_t cargo_capacity = GetVehicleDefaultCapacity(engine, &cargo_type);
+		CargoID cargo_type;
-	if (cargo_type < NUM_CARGO) capacity[cargo_type] = cargo_capacity;
+		uint16_t cargo_capacity = GetVehicleDefaultCapacity(eng, &cargo_type, attempt_refit);
+		if (cargo_type < NUM_CARGO) capacity[cargo_type] += cargo_capacity;
+	};
+	get_engine_cargo(engine);
 
 	if (!e->IsArticulatedCallbackVehicleType()) return capacity;
 
@@ -194,8 +207,7 @@ CargoArray GetCapacityOfArticulatedParts(EngineID engine)
 		EngineID artic_engine = GetNextArticulatedPart(i, engine);
 		if (artic_engine == INVALID_ENGINE) break;
 
-		cargo_capacity = GetVehicleDefaultCapacity(artic_engine, &cargo_type);
+		get_engine_cargo(artic_engine);
-		if (cargo_type < NUM_CARGO) capacity[cargo_type] += cargo_capacity;
 	}
 
 	return capacity;

src/articulated_vehicles.h

@@ -16,7 +16,7 @@
 
 uint CountArticulatedParts(EngineID engine_type, bool purchase_window);
 void GetArticulatedPartsEngineIDs(EngineID engine_type, bool purchase_window, std::vector<EngineID> &ids);
-CargoArray GetCapacityOfArticulatedParts(EngineID engine);
+CargoArray GetCapacityOfArticulatedParts(EngineID engine, CargoID attempt_refit = INVALID_CARGO);
 CargoTypes GetCargoTypesOfArticulatedParts(EngineID engine);
 void AddArticulatedParts(Vehicle *first);
 void GetArticulatedRefitMasks(EngineID engine, bool include_initial_cargo_type, CargoTypes *union_mask, CargoTypes *intersection_mask);

src/build_vehicle_gui.cpp

@@ -36,6 +36,7 @@
 #include "querystring_gui.h"
 #include "stringfilter_type.h"
 #include "hotkeys.h"
+#include "3rdparty/cpp-btree/btree_map.h"
 
 #include "widgets/build_vehicle_widget.h"
 
@@ -205,6 +206,29 @@ bool _engine_sort_show_hidden_locos     = false;                        ///< Las
 bool _engine_sort_show_hidden_wagons    = false;                        ///< Last set 'show hidden wagons' setting.
 static CargoID _engine_sort_last_cargo_criteria[] = {CargoFilterCriteria::CF_ANY, CargoFilterCriteria::CF_ANY, CargoFilterCriteria::CF_ANY, CargoFilterCriteria::CF_ANY}; ///< Last set filter criteria, for each vehicle type.
 
+struct BuildVehicleWindowBase;
+
+struct EngineCapacityCache {
+	const BuildVehicleWindowBase *parent = nullptr;
+	CargoID current_cargo = INVALID_CARGO;
+	btree::btree_map<EngineID, uint> capacities;
+
+	void UpdateCargoFilter(const BuildVehicleWindowBase *parent, CargoID cargo_filter_criteria)
+	{
+		this->parent = parent;
+
+		if (cargo_filter_criteria >= NUM_CARGO) cargo_filter_criteria = INVALID_CARGO;
+
+		if (cargo_filter_criteria != this->current_cargo) {
+			this->current_cargo = cargo_filter_criteria;
+			this->capacities.clear();
+		}
+	}
+
+	uint GetArticulatedCapacity(EngineID eng, bool dual_headed = false);
+};
+static EngineCapacityCache *_engine_sort_capacity_cache = nullptr;
+
 static byte _last_sort_criteria_loco      = 0;
 static bool _last_sort_order_loco         = false;
 static CargoID _last_filter_criteria_loco = CargoFilterCriteria::CF_ANY;
@@ -449,8 +473,8 @@ static bool TrainEngineCapacitySorter(const GUIEngineListItem &a, const GUIEngin
 	const RailVehicleInfo *rvi_a = RailVehInfo(a.engine_id);
 	const RailVehicleInfo *rvi_b = RailVehInfo(b.engine_id);
 
-	int va = GetTotalCapacityOfArticulatedParts(a.engine_id) * (rvi_a->railveh_type == RAILVEH_MULTIHEAD ? 2 : 1);
+	int va = _engine_sort_capacity_cache->GetArticulatedCapacity(a.engine_id, rvi_a->railveh_type == RAILVEH_MULTIHEAD);
-	int vb = GetTotalCapacityOfArticulatedParts(b.engine_id) * (rvi_b->railveh_type == RAILVEH_MULTIHEAD ? 2 : 1);
+	int vb = _engine_sort_capacity_cache->GetArticulatedCapacity(b.engine_id, rvi_b->railveh_type == RAILVEH_MULTIHEAD);
 	int r = va - vb;
 
 	/* Use EngineID to sort instead since we want consistent sorting */
@@ -469,8 +493,8 @@ static bool TrainEngineCapacityVsRunningCostSorter(const GUIEngineListItem &a, c
 	const RailVehicleInfo *rvi_a = RailVehInfo(a.engine_id);
 	const RailVehicleInfo *rvi_b = RailVehInfo(b.engine_id);
 
-	uint va = GetTotalCapacityOfArticulatedParts(a.engine_id) * (rvi_a->railveh_type == RAILVEH_MULTIHEAD ? 2 : 1);
+	uint va = _engine_sort_capacity_cache->GetArticulatedCapacity(a.engine_id, rvi_a->railveh_type == RAILVEH_MULTIHEAD);
-	uint vb = GetTotalCapacityOfArticulatedParts(b.engine_id) * (rvi_b->railveh_type == RAILVEH_MULTIHEAD ? 2 : 1);
+	uint vb = _engine_sort_capacity_cache->GetArticulatedCapacity(b.engine_id, rvi_b->railveh_type == RAILVEH_MULTIHEAD);
 
 	return GenericEngineValueVsRunningCostSorter(a, va, b, vb);
 }
@@ -502,8 +526,8 @@ static bool TrainEnginesThenWagonsSorter(const GUIEngineListItem &a, const GUIEn
  */
 static bool RoadVehEngineCapacitySorter(const GUIEngineListItem &a, const GUIEngineListItem &b)
 {
-	int va = GetTotalCapacityOfArticulatedParts(a.engine_id);
+	int va = _engine_sort_capacity_cache->GetArticulatedCapacity(a.engine_id);
-	int vb = GetTotalCapacityOfArticulatedParts(b.engine_id);
+	int vb = _engine_sort_capacity_cache->GetArticulatedCapacity(b.engine_id);
 	int r = va - vb;
 
 	/* Use EngineID to sort instead since we want consistent sorting */
@@ -519,7 +543,9 @@ static bool RoadVehEngineCapacitySorter(const GUIEngineListItem &a, const GUIEng
  */
 static bool RoadVehEngineCapacityVsRunningCostSorter(const GUIEngineListItem &a, const GUIEngineListItem &b)
 {
-	return GenericEngineValueVsRunningCostSorter(a, GetTotalCapacityOfArticulatedParts(a.engine_id), b, GetTotalCapacityOfArticulatedParts(b.engine_id));
+	int capacity_a = _engine_sort_capacity_cache->GetArticulatedCapacity(a.engine_id);
+	int capacity_b = _engine_sort_capacity_cache->GetArticulatedCapacity(b.engine_id);
+	return GenericEngineValueVsRunningCostSorter(a, capacity_a, b, capacity_b);
 }
 
 /* Ship vehicle sorting functions */
@@ -532,8 +558,8 @@ static bool RoadVehEngineCapacityVsRunningCostSorter(const GUIEngineListItem &a,
  */
 static bool ShipEngineCapacitySorter(const GUIEngineListItem &a, const GUIEngineListItem &b)
 {
-	int va = GetTotalCapacityOfArticulatedParts(a.engine_id);
+	int va = _engine_sort_capacity_cache->GetArticulatedCapacity(a.engine_id);
-	int vb = GetTotalCapacityOfArticulatedParts(b.engine_id);
+	int vb = _engine_sort_capacity_cache->GetArticulatedCapacity(b.engine_id);
 	int r = va - vb;
 
 	/* Use EngineID to sort instead since we want consistent sorting */
@@ -549,7 +575,9 @@ static bool ShipEngineCapacitySorter(const GUIEngineListItem &a, const GUIEngine
  */
 static bool ShipEngineCapacityVsRunningCostSorter(const GUIEngineListItem &a, const GUIEngineListItem &b)
 {
-	return GenericEngineValueVsRunningCostSorter(a, GetTotalCapacityOfArticulatedParts(a.engine_id), b, GetTotalCapacityOfArticulatedParts(b.engine_id));
+	int capacity_a = _engine_sort_capacity_cache->GetArticulatedCapacity(a.engine_id);
+	int capacity_b = _engine_sort_capacity_cache->GetArticulatedCapacity(b.engine_id);
+	return GenericEngineValueVsRunningCostSorter(a, capacity_a, b, capacity_b);
 }
 
 /* Aircraft sorting functions */
@@ -1475,6 +1503,24 @@ struct BuildVehicleWindowBase : Window {
 	}
 };
 
+uint EngineCapacityCache::GetArticulatedCapacity(EngineID eng, bool dual_headed)
+{
+	auto iter = this->capacities.insert({ eng, 0 });
+	if (iter.second) {
+		/* New cache entry */
+		const Engine *e = Engine::Get(eng);
+		if (this->current_cargo != INVALID_CARGO && this->current_cargo != e->GetDefaultCargoType() && HasBit(e->info.callback_mask, CBM_VEHICLE_REFIT_CAPACITY) && e->refit_capacity_values == nullptr) {
+			/* Expensive path simulating vehicle construction is required to determine capacity */
+			TestedEngineDetails te{};
+			this->parent->FillTestedEngineCapacity(eng, this->current_cargo, te);
+			iter.first->second = te.all_capacities.GetSum<uint>();
+		} else {
+			iter.first->second = GetTotalCapacityOfArticulatedParts(eng, this->current_cargo) * (dual_headed ? 2 : 1);
+		}
+	}
+	return iter.first->second;
+}
+
 /** GUI for building vehicles. */
 struct BuildVehicleWindow : BuildVehicleWindowBase {
 	union {
@@ -1491,6 +1537,7 @@ struct BuildVehicleWindow : BuildVehicleWindowBase {
 	int details_height;                         ///< Minimal needed height of the details panels, in text lines (found so far).
 	Scrollbar *vscroll;
 	TestedEngineDetails te;                     ///< Tested cost and capacity after refit.
+	EngineCapacityCache capacity_cache;         ///< Engine capacity cache.
 
 	StringFilter string_filter;                 ///< Filter for vehicle name
 	QueryString vehicle_editbox;                ///< Filter editbox
@@ -1727,6 +1774,10 @@ struct BuildVehicleWindow : BuildVehicleWindowBase {
 		/* invalidate cached values for name sorter - engine names could change */
 		_last_engine[0] = _last_engine[1] = INVALID_ENGINE;
 
+		/* setup engine capacity cache */
+		this->capacity_cache.UpdateCargoFilter(this, this->cargo_filter_criteria);
+		_engine_sort_capacity_cache = &(this->capacity_cache);
+
 		/* make engines first, and then wagons, sorted by selected sort_criteria */
 		_engine_sort_direction = false;
 		EngList_Sort(list, TrainEnginesThenWagonsSorter);
@@ -1864,6 +1915,10 @@ struct BuildVehicleWindow : BuildVehicleWindowBase {
 			}
 		}
 
+		/* setup engine capacity cache */
+		this->capacity_cache.UpdateCargoFilter(this, this->cargo_filter_criteria);
+		_engine_sort_capacity_cache = &(this->capacity_cache);
+
 		_engine_sort_direction = this->descending_sort_order;
 		EngList_Sort(this->eng_list, _engine_sort_functions[this->vehicle_type][this->sort_criteria]);
 
@@ -2303,6 +2358,7 @@ struct BuildVehicleWindowTrainAdvanced final : BuildVehicleWindowBase {
 		bool show_hidden;                              ///< State of the 'show hidden' button.
 		int details_height;                            ///< Minimal needed height of the details panels (found so far).
 		TestedEngineDetails te;                        ///< Tested cost and capacity after refit.
+		EngineCapacityCache capacity_cache;            ///< Engine capacity cache.
 		StringFilter string_filter;                    ///< Filter for vehicle name
 		QueryString vehicle_editbox { MAX_LENGTH_VEHICLE_NAME_CHARS * MAX_CHAR_LENGTH, MAX_LENGTH_VEHICLE_NAME_CHARS }; ///< Filter editbox
 	};
@@ -2607,6 +2663,10 @@ struct BuildVehicleWindowTrainAdvanced final : BuildVehicleWindowBase {
 		/* invalidate cached values for name sorter - engine names could change */
 		_last_engine[0] = _last_engine[1] = INVALID_ENGINE;
 
+		/* setup engine capacity cache */
+		state.capacity_cache.UpdateCargoFilter(this, state.cargo_filter_criteria);
+		_engine_sort_capacity_cache = &(state.capacity_cache);
+
 		/* Sort */
 		_engine_sort_direction = state.descending_sort_order;
 		EngList_Sort(list, sorters[state.sort_criteria]);

src/engine.cpp

@@ -223,9 +223,10 @@ bool Engine::CanPossiblyCarryCargo() const
  * For aircraft the main capacity is determined. Mail might be present as well.
  * @param v Vehicle of interest; nullptr in purchase list
  * @param mail_capacity returns secondary cargo (mail) capacity of aircraft
+ * @param attempt_refit cargo ID to attempt to use, when v is nullptr
  * @return Capacity
  */
-uint Engine::DetermineCapacity(const Vehicle *v, uint16_t *mail_capacity) const
+uint Engine::DetermineCapacity(const Vehicle *v, uint16_t *mail_capacity, CargoID attempt_refit) const
 {
 	assert(v == nullptr || this->index == v->engine_type);
 	if (mail_capacity != nullptr) *mail_capacity = 0;
@@ -234,7 +235,16 @@ uint Engine::DetermineCapacity(const Vehicle *v, uint16_t *mail_capacity) const
 
 	bool new_multipliers = HasBit(this->info.misc_flags, EF_NO_DEFAULT_CARGO_MULTIPLIER);
 	CargoID default_cargo = this->GetDefaultCargoType();
-	CargoID cargo_type = (v != nullptr) ? v->cargo_type : default_cargo;
+	CargoID cargo_type;
+	if (v != nullptr) {
+		cargo_type = v->cargo_type;
+	} else {
+		if (attempt_refit != INVALID_CARGO && HasBit(this->info.refit_mask, attempt_refit)) {
+			cargo_type = attempt_refit;
+		} else {
+			cargo_type = default_cargo;
+		}
+	}
 
 	if (mail_capacity != nullptr && this->type == VEH_AIRCRAFT && IsCargoInClass(cargo_type, CC_PASSENGERS)) {
 		*mail_capacity = GetEngineProperty(this->index, PROP_AIRCRAFT_MAIL_CAPACITY, this->u.air.mail_capacity, v);

src/engine_base.h

@@ -112,7 +112,7 @@ struct Engine : EnginePool::PoolItem<&_engine_pool> {
 		return this->info.cargo_type;
 	}
 
-	uint DetermineCapacity(const Vehicle *v, uint16_t *mail_capacity = nullptr) const;
+	uint DetermineCapacity(const Vehicle *v, uint16_t *mail_capacity = nullptr, CargoID attempt_refit = INVALID_CARGO) const;
 
 	bool CanCarryCargo() const;
 	bool CanPossiblyCarryCargo() const;
@@ -125,12 +125,13 @@ struct Engine : EnginePool::PoolItem<&_engine_pool> {
 	 * For articulated engines use GetCapacityOfArticulatedParts
 	 *
 	 * @param mail_capacity returns secondary cargo (mail) capacity of aircraft
+	 * @param attempt_refit cargo ID to attempt to use
 	 * @return The default capacity
 	 * @see GetDefaultCargoType
 	 */
-	uint GetDisplayDefaultCapacity(uint16_t *mail_capacity = nullptr) const
+	uint GetDisplayDefaultCapacity(uint16_t *mail_capacity = nullptr, CargoID attempt_refit = INVALID_CARGO) const
 	{
-		return this->DetermineCapacity(nullptr, mail_capacity);
+		return this->DetermineCapacity(nullptr, mail_capacity, attempt_refit);
 	}
 
 	Money GetRunningCost() const;

src/engine_func.h

@@ -29,6 +29,6 @@ void SetYearEngineAgingStops();
 void CalcEngineReliability(Engine *e, bool new_month);
 void StartupOneEngine(Engine *e, const CalTime::YearMonthDay &aging_ymd, const CalTime::YearMonthDay &expire_stop_ymd, uint32_t seed, CalTime::Date no_introduce_after_date);
 
-uint GetTotalCapacityOfArticulatedParts(EngineID engine);
+uint GetTotalCapacityOfArticulatedParts(EngineID engine, CargoID attempt_refit = INVALID_CARGO);
 
 #endif /* ENGINE_FUNC_H */

src/engine_gui.cpp

@@ -159,11 +159,12 @@ void ShowEnginePreviewWindow(EngineID engine)
 /**
  * Get the capacity of an engine with articulated parts.
  * @param engine The engine to get the capacity of.
+ * @param attempt_refit Attempt to get capacity when refitting to this cargo.
  * @return The capacity.
  */
-uint GetTotalCapacityOfArticulatedParts(EngineID engine)
+uint GetTotalCapacityOfArticulatedParts(EngineID engine, CargoID attempt_refit)
 {
-	CargoArray cap = GetCapacityOfArticulatedParts(engine);
+	CargoArray cap = GetCapacityOfArticulatedParts(engine, attempt_refit);
 	return cap.GetSum<uint>();
 }