4 files changed, 230 insertions, 69 deletions

diff --git a/src/articulated_vehicles.cpp b/src/articulated_vehicles.cpp
index 74b21a3b5..f959da77e 100644
--- a/src/articulated_vehicles.cpp
+++ b/src/articulated_vehicles.cpp
@@ -35,19 +35,78 @@ uint CountArticulatedParts(EngineID engine_type, bool purchase_window)
 }
 
 
+/**
+ * Returns the default (non-refitted) capacity of a specific EngineID.
+ * @param engine the EngineID of iterest
+ * @param type the type of the engine
+ * @param cargo_type returns the default cargo type, if needed
+ * @return capacity
+ */
+static inline uint16 GetVehicleDefaultCapacity(EngineID engine, VehicleType type, CargoID *cargo_type)
+{
+	switch (type) {
+		case VEH_TRAIN: {
+			const RailVehicleInfo *rvi = RailVehInfo(engine);
+			if (cargo_type != NULL) *cargo_type = rvi->cargo_type;
+			return GetEngineProperty(engine, 0x14, rvi->capacity) + (rvi->railveh_type == RAILVEH_MULTIHEAD ? rvi->capacity : 0);
+		}
+
+		case VEH_ROAD: {
+			const RoadVehicleInfo *rvi = RoadVehInfo(engine);
+			if (cargo_type != NULL) *cargo_type = rvi->cargo_type;
+			return GetEngineProperty(engine, 0x0F, rvi->capacity);
+		}
+
+		case VEH_SHIP: {
+			const ShipVehicleInfo *svi = ShipVehInfo(engine);
+			if (cargo_type != NULL) *cargo_type = svi->cargo_type;
+			return GetEngineProperty(engine, 0x0D, svi->capacity);
+		}
+
+		case VEH_AIRCRAFT: {
+			const AircraftVehicleInfo *avi = AircraftVehInfo(engine);
+			if (cargo_type != NULL) *cargo_type = CT_PASSENGERS;
+			return avi->passenger_capacity;
+		}
+
+		default: NOT_REACHED();
+	}
+
+}
+
+/**
+ * Returns all cargos a vehicle can carry.
+ * @param engine the EngineID of iterest
+ * @param type the type of the engine
+ * @param include_initial_cargo_type if true the default cargo type of the vehicle is included; if false only the refit_mask
+ * @return bit set of CargoIDs
+ */
+static inline uint32 GetAvailableVehicleCargoTypes(EngineID engine, VehicleType type, bool include_initial_cargo_type)
+{
+	uint32 cargos = 0;
+	CargoID initial_cargo_type;
+
+	if (GetVehicleDefaultCapacity(engine, type, &initial_cargo_type) > 0) {
+		if (type != VEH_SHIP || ShipVehInfo(engine)->refittable) {
+			const EngineInfo *ei = EngInfo(engine);
+			cargos = ei->refit_mask;
+		}
+		if (include_initial_cargo_type && initial_cargo_type < NUM_CARGO) SetBit(cargos, initial_cargo_type);
+	}
+
+	return cargos;
+}
+
 uint16 *GetCapacityOfArticulatedParts(EngineID engine, VehicleType type)
 {
 	static uint16 capacity[NUM_CARGO];
 	memset(capacity, 0, sizeof(capacity));
 
-	if (type == VEH_TRAIN) {
-		const RailVehicleInfo *rvi = RailVehInfo(engine);
-		capacity[rvi->cargo_type] = GetEngineProperty(engine, 0x14, rvi->capacity);
-		if (rvi->railveh_type == RAILVEH_MULTIHEAD) capacity[rvi->cargo_type] += rvi->capacity;
-	} else if (type == VEH_ROAD) {
-		const RoadVehicleInfo *rvi = RoadVehInfo(engine);
-		capacity[rvi->cargo_type] = GetEngineProperty(engine, 0x0F, rvi->capacity);
-	}
+	CargoID cargo_type;
+	uint16 cargo_capacity = GetVehicleDefaultCapacity(engine, type, &cargo_type);
+	if (cargo_type < NUM_CARGO) capacity[cargo_type] = cargo_capacity;
+
+	if (type != VEH_TRAIN && type != VEH_ROAD) return capacity;
 
 	if (!HasBit(EngInfo(engine)->callbackmask, CBM_VEHICLE_ARTIC_ENGINE)) return capacity;
 
@@ -57,18 +116,111 @@ uint16 *GetCapacityOfArticulatedParts(EngineID engine, VehicleType type)
 
 		EngineID artic_engine = GetNewEngineID(GetEngineGRF(engine), type, GB(callback, 0, 7));
 
-		if (type == VEH_TRAIN) {
-			const RailVehicleInfo *rvi = RailVehInfo(artic_engine);
-			capacity[rvi->cargo_type] += GetEngineProperty(artic_engine, 0x14, rvi->capacity);
-		} else if (type == VEH_ROAD) {
-			const RoadVehicleInfo *rvi = RoadVehInfo(artic_engine);
-			capacity[rvi->cargo_type] += GetEngineProperty(artic_engine, 0x0F, rvi->capacity);
-		}
+		cargo_capacity = GetVehicleDefaultCapacity(artic_engine, type, &cargo_type);
+		if (cargo_type < NUM_CARGO) capacity[cargo_type] += cargo_capacity;
 	}
 
 	return capacity;
 }
 
+/**
+ * Ors the refit_masks of all articulated parts.
+ * Note: Vehicles with a default capacity of zero are ignored.
+ * @param engine the first part
+ * @param type the vehicle type
+ * @param include_initial_cargo_type if true the default cargo type of the vehicle is included; if false only the refit_mask
+ * @return bit mask of CargoIDs which are a refit option for at least one articulated part
+ */
+uint32 GetUnionOfArticulatedRefitMasks(EngineID engine, VehicleType type, bool include_initial_cargo_type)
+{
+	uint32 cargos = GetAvailableVehicleCargoTypes(engine, type, include_initial_cargo_type);
+
+	if (type != VEH_TRAIN && type != VEH_ROAD) return cargos;
+
+	if (!HasBit(EngInfo(engine)->callbackmask, CBM_VEHICLE_ARTIC_ENGINE)) return cargos;
+
+	for (uint i = 1; i < MAX_ARTICULATED_PARTS; i++) {
+		uint16 callback = GetVehicleCallback(CBID_VEHICLE_ARTIC_ENGINE, i, 0, engine, NULL);
+		if (callback == CALLBACK_FAILED || GB(callback, 0, 8) == 0xFF) break;
+
+		EngineID artic_engine = GetNewEngineID(GetEngineGRF(engine), type, GB(callback, 0, 7));
+		cargos |= GetAvailableVehicleCargoTypes(artic_engine, type, include_initial_cargo_type);
+	}
+
+	return cargos;
+}
+
+/**
+ * Ands the refit_masks of all articulated parts.
+ * Note: Vehicles with a default capacity of zero are ignored.
+ * @param engine the first part
+ * @param type the vehicle type
+ * @param include_initial_cargo_type if true the default cargo type of the vehicle is included; if false only the refit_mask
+ * @return bit mask of CargoIDs which are a refit option for every articulated part (with default capacity > 0)
+ */
+uint32 GetIntersectionOfArticulatedRefitMasks(EngineID engine, VehicleType type, bool include_initial_cargo_type)
+{
+	uint32 cargos = UINT32_MAX;
+
+	uint32 veh_cargos = GetAvailableVehicleCargoTypes(engine, type, include_initial_cargo_type);
+	if (veh_cargos != 0) cargos &= veh_cargos;
+
+	if (type != VEH_TRAIN && type != VEH_ROAD) return cargos;
+
+	if (!HasBit(EngInfo(engine)->callbackmask, CBM_VEHICLE_ARTIC_ENGINE)) return cargos;
+
+	for (uint i = 1; i < MAX_ARTICULATED_PARTS; i++) {
+		uint16 callback = GetVehicleCallback(CBID_VEHICLE_ARTIC_ENGINE, i, 0, engine, NULL);
+		if (callback == CALLBACK_FAILED || GB(callback, 0, 8) == 0xFF) break;
+
+		EngineID artic_engine = GetNewEngineID(GetEngineGRF(engine), type, GB(callback, 0, 7));
+		veh_cargos = GetAvailableVehicleCargoTypes(artic_engine, type, include_initial_cargo_type);
+		if (veh_cargos != 0) cargos &= veh_cargos;
+	}
+
+	return cargos;
+}
+
+
+/**
+ * Tests if all parts of an articulated vehicle are refitted to the same cargo.
+ * Note: Vehicles not carrying anything are ignored
+ * @param v the first vehicle in the chain
+ * @param cargo_type returns the common CargoID if needed. (CT_INVALID if no part is carrying something or they are carrying different things)
+ * @return true if some parts are carrying different cargos, false if all parts are carrying the same (nothing is also the same)
+ */
+bool IsArticulatedVehicleCarryingDifferentCargos(const Vehicle *v, CargoID *cargo_type)
+{
+	CargoID first_cargo = CT_INVALID;
+
+	do {
+		if (v->cargo_cap > 0 && v->cargo_type != CT_INVALID) {
+			if (first_cargo == CT_INVALID) first_cargo = v->cargo_type;
+			if (first_cargo != v->cargo_type) {
+				if (cargo_type != NULL) *cargo_type = CT_INVALID;
+				return true;
+			}
+		}
+
+		switch (v->type) {
+			case VEH_TRAIN:
+				v = (EngineHasArticPart(v) ? GetNextArticPart(v) : NULL);
+				break;
+
+			case VEH_ROAD:
+				v = (RoadVehHasArticPart(v) ? v->Next() : NULL);
+				break;
+
+			default:
+				v = NULL;
+				break;
+		}
+	} while (v != NULL);
+
+	if (cargo_type != NULL) *cargo_type = first_cargo;
+	return false;
+}
+
 
 void AddArticulatedParts(Vehicle **vl, VehicleType type)
 {
diff --git a/src/articulated_vehicles.h b/src/articulated_vehicles.h
index abfe495b0..66f2e7f2f 100644
--- a/src/articulated_vehicles.h
+++ b/src/articulated_vehicles.h
@@ -11,5 +11,8 @@
 uint CountArticulatedParts(EngineID engine_type, bool purchase_window);
 uint16 *GetCapacityOfArticulatedParts(EngineID engine, VehicleType type);
 void AddArticulatedParts(Vehicle **vl, VehicleType type);
+uint32 GetUnionOfArticulatedRefitMasks(EngineID engine, VehicleType type, bool include_initial_cargo_type);
+uint32 GetIntersectionOfArticulatedRefitMasks(EngineID engine, VehicleType type, bool include_initial_cargo_type);
+bool IsArticulatedVehicleCarryingDifferentCargos(const Vehicle *v, CargoID *cargo_type);
 
 #endif /* ARTICULATED_VEHICLES_H */
diff --git a/src/autoreplace_cmd.cpp b/src/autoreplace_cmd.cpp
index 5d728c0c5..1a8ac4393 100644
--- a/src/autoreplace_cmd.cpp
+++ b/src/autoreplace_cmd.cpp
@@ -60,11 +60,21 @@ static void MoveVehicleCargo(Vehicle *dest, Vehicle *source)
 	if (dest->type == VEH_TRAIN) TrainConsistChanged(dest->First(), true);
 }
 
-static bool VerifyAutoreplaceRefitForOrders(const Vehicle *v, const EngineID engine_type)
+
+/**
+ * Tests whether refit orders that applied to v will also apply to the new vehicle type
+ * @param v The vehicle to be replaced
+ * @param engine_type The type we want to replace with
+ * @return true iff all refit orders stay valid
+ */
+static bool VerifyAutoreplaceRefitForOrders(const Vehicle *v, EngineID engine_type)
 {
 	const Order *o;
 	const Vehicle *u;
 
+	uint32 union_refit_mask_a = GetUnionOfArticulatedRefitMasks(v->engine_type, v->type, false);
+	uint32 union_refit_mask_b = GetUnionOfArticulatedRefitMasks(engine_type, v->type, false);
+
 	if (v->type == VEH_TRAIN) {
 		u = v->First();
 	} else {
@@ -73,8 +83,10 @@ static bool VerifyAutoreplaceRefitForOrders(const Vehicle *v, const EngineID eng
 
 	FOR_VEHICLE_ORDERS(u, o) {
 		if (!o->IsRefit()) continue;
-		if (!CanRefitTo(v->engine_type, o->GetRefitCargo())) continue;
-		if (!CanRefitTo(engine_type, o->GetRefitCargo())) return false;
+		CargoID cargo_type = o->GetRefitCargo();
+
+		if (!HasBit(union_refit_mask_a, cargo_type)) continue;
+		if (!HasBit(union_refit_mask_b, cargo_type)) return false;
 	}
 
 	return true;
@@ -90,33 +102,48 @@ static bool VerifyAutoreplaceRefitForOrders(const Vehicle *v, const EngineID eng
  */
 static CargoID GetNewCargoTypeForReplace(Vehicle *v, EngineID engine_type)
 {
-	CargoID new_cargo_type = GetEngineCargoType(engine_type);
+	CargoID cargo_type;
 
-	if (new_cargo_type == CT_INVALID) return CT_NO_REFIT; // Don't try to refit an engine with no cargo capacity
+	if (GetUnionOfArticulatedRefitMasks(engine_type, v->type, true) == 0) return CT_NO_REFIT; // Don't try to refit an engine with no cargo capacity
 
-	if (v->cargo_cap != 0 && (v->cargo_type == new_cargo_type || CanRefitTo(engine_type, v->cargo_type))) {
-		if (VerifyAutoreplaceRefitForOrders(v, engine_type)) {
-			return v->cargo_type == new_cargo_type ? (CargoID)CT_NO_REFIT : v->cargo_type;
-		} else {
-			return CT_INVALID;
+	if (IsArticulatedVehicleCarryingDifferentCargos(v, &cargo_type)) return CT_INVALID; // We cannot refit to mixed cargos in an automated way
+
+	uint32 available_cargo_types = GetIntersectionOfArticulatedRefitMasks(engine_type, v->type, true);
+
+	if (cargo_type == CT_INVALID) {
+		if (v->type != VEH_TRAIN) return CT_NO_REFIT; // If the vehicle does not carry anything at all, every replacement is fine.
+
+		/* the old engine didn't have cargo capacity, but the new one does
+		 * now we will figure out what cargo the train is carrying and refit to fit this */
+
+		for (v = v->First(); v != NULL; v = v->Next()) {
+			if (v->cargo_cap == 0) continue;
+			/* Now we found a cargo type being carried on the train and we will see if it is possible to carry to this one */
+			if (HasBit(available_cargo_types, v->cargo_type)) {
+				/* Do we have to refit the vehicle, or is it already carrying the right cargo? */
+				uint16 *default_capacity = GetCapacityOfArticulatedParts(engine_type, v->type);
+				for (CargoID cid = 0; cid < NUM_CARGO; cid++) {
+					if (cid != cargo_type && default_capacity[cid] > 0) return cargo_type;
+				}
+
+				return CT_NO_REFIT;
+			}
 		}
-	}
-	if (v->type != VEH_TRAIN) return CT_INVALID; // We can't refit the vehicle to carry the cargo we want
 
-	/* Below this line it's safe to assume that the vehicle in question is a train */
+		return CT_NO_REFIT; // We failed to find a cargo type on the old vehicle and we will not refit the new one
+	} else {
+		if (!HasBit(available_cargo_types, cargo_type)) return CT_INVALID; // We can't refit the vehicle to carry the cargo we want
+
+		if (!VerifyAutoreplaceRefitForOrders(v, engine_type)) return CT_INVALID; // Some refit orders loose their effect
 
-	if (v->cargo_cap != 0) return CT_INVALID; // trying to replace a vehicle with cargo capacity into another one with incompatible cargo type
+		/* Do we have to refit the vehicle, or is it already carrying the right cargo? */
+		uint16 *default_capacity = GetCapacityOfArticulatedParts(engine_type, v->type);
+		for (CargoID cid = 0; cid < NUM_CARGO; cid++) {
+			if (cid != cargo_type && default_capacity[cid] > 0) return cargo_type;
+		}
 
-	/* the old engine didn't have cargo capacity, but the new one does
-	 * now we will figure out what cargo the train is carrying and refit to fit this */
-	v = v->First();
-	do {
-		if (v->cargo_cap == 0) continue;
-		/* Now we found a cargo type being carried on the train and we will see if it is possible to carry to this one */
-		if (v->cargo_type == new_cargo_type) return CT_NO_REFIT;
-		if (CanRefitTo(engine_type, v->cargo_type)) return v->cargo_type;
-	} while ((v = v->Next()) != NULL);
-	return CT_NO_REFIT; // We failed to find a cargo type on the old vehicle and we will not refit the new one
+		return CT_NO_REFIT;
+	}
 }
 
 /** Replaces a vehicle (used to be called autorenew)
@@ -344,6 +371,8 @@ static CommandCost WagonRemoval(Vehicle *v, uint16 old_total_length)
  */
 static EngineID GetNewEngineType(const Vehicle *v, const Player *p)
 {
+	assert(v->type != VEH_TRAIN || !IsArticulatedPart(v));
+
 	if (v->type == VEH_TRAIN && IsRearDualheaded(v)) {
 		/* we build the rear ends of multiheaded trains with the front ones */
 		return INVALID_ENGINE;
diff --git a/src/autoreplace_gui.cpp b/src/autoreplace_gui.cpp
index 8fbe26198..6b2bbb8d3 100644
--- a/src/autoreplace_gui.cpp
+++ b/src/autoreplace_gui.cpp
@@ -23,6 +23,7 @@
 #include "engine_base.h"
 #include "window_gui.h"
 #include "engine_gui.h"
+#include "articulated_vehicles.h"
 
 #include "table/sprites.h"
 #include "table/strings.h"
@@ -98,41 +99,17 @@ void AddRemoveEngineFromAutoreplaceAndBuildWindows(VehicleType type)
 	InvalidateWindowClassesData(WC_BUILD_VEHICLE); // The build windows needs updating as well
 }
 
-/** Get the default cargo type for an engine
- * @param engine the EngineID to get the cargo for
- * @return the cargo type carried by the engine (CT_INVALID if engine got no cargo capacity)
- */
-static CargoID EngineCargo(EngineID engine)
-{
-	if (engine == INVALID_ENGINE) return CT_INVALID; // surely INVALID_ENGINE can't carry anything but CT_INVALID
-
-	switch (GetEngine(engine)->type) {
-		default: NOT_REACHED();
-		case VEH_TRAIN:
-			if (RailVehInfo(engine)->capacity == 0) return CT_INVALID; // no capacity -> can't carry cargo
-			return RailVehInfo(engine)->cargo_type;
-		case VEH_ROAD:       return RoadVehInfo(engine)->cargo_type;
-		case VEH_SHIP:       return ShipVehInfo(engine)->cargo_type;
-		case VEH_AIRCRAFT:   return CT_PASSENGERS; // all planes are build with passengers by default
-	}
-}
-
 /** Figure out if two engines got at least one type of cargo in common (refitting if needed)
  * @param engine_a one of the EngineIDs
  * @param engine_b the other EngineID
+ * @param type the type of the engines
  * @return true if they can both carry the same type of cargo (or at least one of them got no capacity at all)
  */
-static bool EnginesGotCargoInCommon(EngineID engine_a, EngineID engine_b)
+static bool EnginesGotCargoInCommon(EngineID engine_a, EngineID engine_b, VehicleType type)
 {
-	CargoID a = EngineCargo(engine_a);
-	CargoID b = EngineCargo(engine_b);
-
-	 /* we should always be able to refit to/from locomotives without capacity
-	  * Because of that, CT_INVALID shoudl always return true */
-	if (a == CT_INVALID || b == CT_INVALID || a == b) return true; // they carry no ro the same type by default
-	if (EngInfo(engine_a)->refit_mask & EngInfo(engine_b)->refit_mask) return true; // both can refit to the same
-	if (CanRefitTo(engine_a, b) || CanRefitTo(engine_b, a)) return true; // one can refit to what the other one carries
-	return false;
+	uint32 available_cargos_a = GetUnionOfArticulatedRefitMasks(engine_a, type, true);
+	uint32 available_cargos_b = GetUnionOfArticulatedRefitMasks(engine_b, type, true);
+	return (available_cargos_a == 0 || available_cargos_b == 0 || (available_cargos_a & available_cargos_b) != 0);
 }
 
 /**
@@ -202,7 +179,7 @@ class ReplaceVehicleWindow : public Window {
 			} else {
 				/* This is for engines we can replace to and they should depend on what we selected to replace from */
 				if (!IsEngineBuildable(eid, type, _local_player)) continue; // we need to be able to build the engine
-				if (!EnginesGotCargoInCommon(eid, this->sel_engine[0])) continue; // the engines needs to be able to carry the same cargo
+				if (!EnginesGotCargoInCommon(eid, this->sel_engine[0], type)) continue; // the engines needs to be able to carry the same cargo
 
 				/* Road vehicles can't be replaced by trams and vice-versa */
 				if (type == VEH_ROAD && HasBit(EngInfo(this->sel_engine[0])->misc_flags, EF_ROAD_TRAM) != HasBit(e->info.misc_flags, EF_ROAD_TRAM)) continue;