(svn r26586) -Fix [FS#5995]: Consider multiheaded trains in station refits (and clean up a bit)

1 files changed, 131 insertions, 51 deletions

diff --git a/src/economy.cpp b/src/economy.cpp
index 0d3902dfa..e36e43147 100644
--- a/src/economy.cpp
+++ b/src/economy.cpp
@@ -19,6 +19,7 @@
 #include "network/network_func.h"
 #include "ai/ai.hpp"
 #include "aircraft.h"
+#include "train.h"
 #include "newgrf_engine.h"
 #include "engine_base.h"
 #include "ground_vehicle.hpp"
@@ -1340,22 +1341,136 @@ static void ReserveConsist(Station *st, Vehicle *u, CargoArray *consist_capleft,
 }
 
 /**
- * Checks whether an articulated vehicle is empty.
- * @param v Vehicle
- * @return true if all parts are empty.
+ * Iterate the articulated parts of a vehicle, also considering the special cases of "normal"
+ * aircraft and double headed trains. Apply an action to each vehicle and immediately return false
+ * if that action does so. Otherwise return true.
+ * @tparam Taction Class of action to be applied. Must implement bool operator()([const] Vehicle *).
+ * @param v First articulated part.
+ * @param action Instance of Taction.
+ * @return false if any of the action invocations returned false, true otherwise.
  */
-static bool IsArticulatedVehicleEmpty(Vehicle *v)
+template<class Taction>
+bool IterateVehicleParts(Vehicle *v, Taction action)
 {
-	v = v->GetFirstEnginePart();
-
-	for (; v != NULL; v = v->HasArticulatedPart() ? v->GetNextArticulatedPart() : NULL) {
-		if (v->cargo.StoredCount() != 0) return false;
+	for (Vehicle *w = v; w != NULL;
+			w = w->HasArticulatedPart() ? w->GetNextArticulatedPart() : NULL) {
+		if (!action(w)) return false;
+		if (w->type == VEH_TRAIN) {
+			Train *train = Train::From(w);
+			if (train->IsMultiheaded() && !action(train->other_multiheaded_part)) return false;
+		}
 	}
-
+	if (v->type == VEH_AIRCRAFT && Aircraft::From(v)->IsNormalAircraft()) return action(v->Next());
 	return true;
 }
 
 /**
+ * Action to check if a vehicle has no stored cargo.
+ */
+struct IsEmptyAction
+{
+	/**
+	 * Checks if the vehicle has stored cargo.
+	 * @param v Vehicle to be checked.
+	 * @return true if v is either empty or has only reserved cargo, false otherwise.
+	 */
+	bool operator()(const Vehicle *v)
+	{
+		return v->cargo.StoredCount() == 0;
+	}
+};
+
+/**
+ * Refit preparation action.
+ */
+struct PrepareRefitAction
+{
+	CargoArray &consist_capleft; ///< Capacities left in the consist.
+	uint32 &refit_mask;          ///< Bitmask of possible refit cargoes.
+
+	/**
+	 * Create a refit preparation action.
+	 * @param consist_capleft Capacities left in consist, to be updated here.
+	 * @param refit_mask Refit mask to be constructed from refit information of vehicles.
+	 */
+	PrepareRefitAction(CargoArray &consist_capleft, uint32 &refit_mask) :
+		consist_capleft(consist_capleft), refit_mask(refit_mask) {}
+
+	/**
+	 * Prepares for refitting of a vehicle, subtracting its free capacity from consist_capleft and
+	 * adding the cargoes it can refit to to the refit mask.
+	 * @param v The vehicle to be refitted.
+	 * @return true.
+	 */
+	bool operator()(const Vehicle *v)
+	{
+		this->consist_capleft[v->cargo_type] -= v->cargo_cap - v->cargo.ReservedCount();
+		this->refit_mask |= EngInfo(v->engine_type)->refit_mask;
+		return true;
+	}
+};
+
+/**
+ * Action for returning reserved cargo.
+ */
+struct ReturnCargoAction
+{
+	Station *st;        ///< Station to give the returned cargo to.
+	StationID next_hop; ///< Next hop the cargo should be assigned to.
+
+	/**
+	 * Construct a cargo return action.
+	 * @param st Station to give the returned cargo to.
+	 * @param next_one Next hop the cargo should be assigned to.
+	 */
+	ReturnCargoAction(Station *st, StationID next_one) : st(st), next_hop(next_one) {}
+
+	/**
+	 * Return all reserved cargo from a vehicle.
+	 * @param v Vehicle to return cargo from.
+	 * @return true.
+	 */
+	bool operator()(Vehicle *v)
+	{
+		v->cargo.Return(UINT_MAX, &this->st->goods[v->cargo_type].cargo, this->next_hop);
+		return true;
+	}
+};
+
+/**
+ * Action for finalizing a refit.
+ */
+struct FinalizeRefitAction
+{
+	CargoArray &consist_capleft;  ///< Capacities left in the consist.
+	Station *st;                  ///< Station to reserve cargo from.
+	StationIDStack &next_station; ///< Next hops to reserve cargo for.
+
+	/**
+	 * Create a finalizing action.
+	 * @param consist_capleft Capacities left in the consist.
+	 * @param st Station to reserve cargo from.
+	 * @param next_station Next hops to reserve cargo for.
+	 */
+	FinalizeRefitAction(CargoArray &consist_capleft, Station *st, StationIDStack &next_station) :
+		consist_capleft(consist_capleft), st(st), next_station(next_station) {}
+
+	/**
+	 * Reserve cargo from the station and update the remaining consist capacities with the
+	 * vehicle's remaining free capacity.
+	 * @param v Vehicle to be finalized.
+	 * @return true.
+	 */
+	bool operator()(Vehicle *v)
+	{
+		this->st->goods[v->cargo_type].cargo.Reserve(v->cargo_cap - v->cargo.RemainingCount(),
+				&v->cargo, st->xy, next_station);
+		this->consist_capleft[v->cargo_type] += v->cargo_cap - v->cargo.RemainingCount();
+		return true;
+	}
+};
+
+/**
  * Refit a vehicle in a station.
  * @param v Vehicle to be refitted.
  * @param consist_capleft Added cargo capacities in the consist.
@@ -1365,32 +1480,15 @@ static bool IsArticulatedVehicleEmpty(Vehicle *v)
  */
 static void HandleStationRefit(Vehicle *v, CargoArray &consist_capleft, Station *st, StationIDStack next_station, CargoID new_cid)
 {
-	if (v->type == VEH_AIRCRAFT && (!Aircraft::From(v)->IsNormalAircraft() || v->Next()->cargo.StoredCount() > 0)) {
-		return;
-	}
-
-	bool is_normal_aircraft = (v->type == VEH_AIRCRAFT && Aircraft::From(v)->IsNormalAircraft());
 	Vehicle *v_start = v->GetFirstEnginePart();
-
-	/* Remove old capacity from consist capacity */
-	consist_capleft[v_start->cargo_type] -= (v_start->cargo_cap - v_start->cargo.ReservedCount());
-	for (Vehicle *w = v_start; w->HasArticulatedPart(); ) {
-		w = w->GetNextArticulatedPart();
-		consist_capleft[w->cargo_type] -= (w->cargo_cap - w->cargo.ReservedCount());
-	}
-	if (is_normal_aircraft) {
-		consist_capleft[v->Next()->cargo_type] -= (v->Next()->cargo_cap - v->Next()->cargo.ReservedCount());
-	}
+	if (!IterateVehicleParts(v_start, IsEmptyAction())) return;
 
 	Backup<CompanyByte> cur_company(_current_company, v->owner, FILE_LINE);
 
-	/* Check if all articulated parts are empty and collect refit mask. */
 	uint32 refit_mask = v->GetEngine()->info.refit_mask;
-	Vehicle *w = v_start;
-	while (w->HasArticulatedPart()) {
-		w = w->GetNextArticulatedPart();
-		refit_mask |= EngInfo(w->engine_type)->refit_mask;
-	}
+
+	/* Remove old capacity from consist capacity and collect refit mask. */
+	IterateVehicleParts(v_start, PrepareRefitAction(consist_capleft, refit_mask));
 
 	if (new_cid == CT_AUTO_REFIT) {
 		/* Get a refittable cargo type with waiting cargo for next_station or INVALID_STATION. */
@@ -1418,31 +1516,13 @@ static void HandleStationRefit(Vehicle *v, CargoArray &consist_capleft, Station
 
 	/* Refit if given a valid cargo. */
 	if (new_cid < NUM_CARGO && new_cid != v_start->cargo_type) {
-		StationID next_one = StationIDStack(next_station).Pop();
-		v_start->cargo.Return(UINT_MAX, &st->goods[v_start->cargo_type].cargo, next_one);
-		for (w = v_start; w->HasArticulatedPart();) {
-			w = w->GetNextArticulatedPart();
-			w->cargo.Return(UINT_MAX, &st->goods[w->cargo_type].cargo, next_one);
-		}
-		if (is_normal_aircraft) {
-			v->Next()->cargo.Return(UINT_MAX, &st->goods[v->Next()->cargo_type].cargo, next_one);
-		}
+		IterateVehicleParts(v_start, ReturnCargoAction(st, StationIDStack(next_station).Pop()));
 		CommandCost cost = DoCommand(v_start->tile, v_start->index, new_cid | 1U << 6 | 0xFF << 8 | 1U << 16, DC_EXEC, GetCmdRefitVeh(v_start)); // Auto-refit and only this vehicle including artic parts.
 		if (cost.Succeeded()) v->First()->profit_this_year -= cost.GetCost() << 8;
 	}
 
 	/* Add new capacity to consist capacity and reserve cargo */
-	w = v_start;
-	do {
-		st->goods[w->cargo_type].cargo.Reserve(w->cargo_cap - w->cargo.RemainingCount(), &w->cargo, st->xy, next_station);
-		consist_capleft[w->cargo_type] += w->cargo_cap - w->cargo.RemainingCount();
-		w = w->HasArticulatedPart() ? w->GetNextArticulatedPart() : NULL;
-	} while (w != NULL);
-	if (is_normal_aircraft) {
-		w = v->Next();
-		st->goods[w->cargo_type].cargo.Reserve(w->cargo_cap - w->cargo.RemainingCount(), &w->cargo, st->xy, next_station);
-		consist_capleft[w->cargo_type] += w->cargo_cap - w->cargo.RemainingCount();
-	}
+	IterateVehicleParts(v_start, FinalizeRefitAction(consist_capleft, st, next_station));
 
 	cur_company.Restore();
 }
@@ -1595,7 +1675,7 @@ static void LoadUnloadVehicle(Vehicle *front)
 		if (front->current_order.GetLoadType() & OLFB_NO_LOAD || HasBit(front->vehicle_flags, VF_STOP_LOADING)) continue;
 
 		/* This order has a refit, if this is the first vehicle part carrying cargo and the whole vehicle is empty, try refitting. */
-		if (front->current_order.IsRefit() && artic_part == 1 && IsArticulatedVehicleEmpty(v)) {
+		if (front->current_order.IsRefit() && artic_part == 1) {
 			HandleStationRefit(v, consist_capleft, st, next_station, front->current_order.GetRefitCargo());
 			ge = &st->goods[v->cargo_type];
 		}