From a1b60bcc858f4bed8946a7212b51c3f2716303c6 Mon Sep 17 00:00:00 2001
From: Jonathan G Rennison <j.g.rennison@gmail.com>
Date: Mon, 12 Oct 2015 20:06:59 +0100
Subject: [PATCH] Fix build and refit cost estimation with artic/multihead
 vehicles.

Trailing vehicle part refit costs were not being added to the cost
estimate.
---
 src/articulated_vehicles.cpp | 33 +++++++++++++++++++++++++++++++++
 src/articulated_vehicles.h   |  2 ++
 src/vehicle_cmd.cpp          | 10 ++++++++++
 3 files changed, 45 insertions(+)

diff --git a/src/articulated_vehicles.cpp b/src/articulated_vehicles.cpp
index 193c5ac4a8..c0c2563d8c 100644
--- a/src/articulated_vehicles.cpp
+++ b/src/articulated_vehicles.cpp
@@ -16,6 +16,7 @@
 #include "engine_func.h"
 #include "company_func.h"
 #include "newgrf.h"
+#include <vector>
 
 #include "table/strings.h"
 
@@ -97,6 +98,38 @@ uint CountArticulatedParts(EngineID engine_type, bool purchase_window)
 	return i - 1;
 }
 
+/**
+ * Count the number of articulated parts of an engine.
+ * @param engine_type The engine to get the number of parts of.
+ * @param purchase_window Whether we are in the scope of the purchase window or not, i.e. whether we cannot allocate vehicles.
+ * @param ids [Out] The list of engine IDs.
+ */
+void GetArticulatedPartsEngineIDs(EngineID engine_type, bool purchase_window, std::vector<EngineID> &ids)
+{
+	ids.clear();
+	if (!HasBit(EngInfo(engine_type)->callback_mask, CBM_VEHICLE_ARTIC_ENGINE)) return;
+
+	/* If we can't allocate a vehicle now, we can't allocate it in the command
+	 * either, so it doesn't matter how many articulated parts there are. */
+	if (!Vehicle::CanAllocateItem()) return;
+
+	Vehicle *v = NULL;
+	if (!purchase_window) {
+		v = new Vehicle();
+		v->engine_type = engine_type;
+		v->owner = _current_company;
+	}
+
+	uint i;
+	for (i = 1; i < MAX_ARTICULATED_PARTS; i++) {
+		EngineID id = GetNextArticulatedPart(i, engine_type, v);
+		if (id == INVALID_ENGINE) break;
+		ids.push_back(id);
+	}
+
+	delete v;
+}
+
 
 /**
  * Returns the default (non-refitted) capacity of a specific EngineID.
diff --git a/src/articulated_vehicles.h b/src/articulated_vehicles.h
index 46b4da3744..e1256807d8 100644
--- a/src/articulated_vehicles.h
+++ b/src/articulated_vehicles.h
@@ -14,8 +14,10 @@
 
 #include "vehicle_type.h"
 #include "engine_type.h"
+#include <vector>
 
 uint CountArticulatedParts(EngineID engine_type, bool purchase_window);
+void GetArticulatedPartsEngineIDs(EngineID engine_type, bool purchase_window, std::vector<EngineID> &ids);
 CargoArray GetCapacityOfArticulatedParts(EngineID engine);
 void AddArticulatedParts(Vehicle *first);
 void GetArticulatedRefitMasks(EngineID engine, bool include_initial_cargo_type, uint32 *union_mask, uint32 *intersection_mask);
diff --git a/src/vehicle_cmd.cpp b/src/vehicle_cmd.cpp
index 6d5596ef66..1590647c0c 100644
--- a/src/vehicle_cmd.cpp
+++ b/src/vehicle_cmd.cpp
@@ -159,6 +159,16 @@ CommandCost CmdBuildVehicle(TileIndex tile, DoCommandFlag flags, uint32 p1, uint
 			} else {
 				bool auto_refit_allowed = false;
 				value.AddCost(GetRefitCost(NULL, eid, cargo, 0, &auto_refit_allowed));
+				if (type == VEH_TRAIN || type == VEH_ROAD) {
+					std::vector<EngineID> engine_ids;
+					GetArticulatedPartsEngineIDs(eid, false, engine_ids);
+					for (size_t i = 0; i < engine_ids.size(); i++) {
+						value.AddCost(GetRefitCost(NULL, engine_ids[i], cargo, 0, &auto_refit_allowed));
+					}
+				}
+				if (type == VEH_TRAIN && e->u.rail.railveh_type == RAILVEH_MULTIHEAD) {
+					value.AddCost(GetRefitCost(NULL, eid, cargo, 0, &auto_refit_allowed));
+				}
 			}
 		}
 	}