(svn r14753) -Fix (r1): after buying a company, one could have more vehicles with the same UnitID

3 files changed, 60 insertions, 60 deletions

diff --git a/src/economy.cpp b/src/economy.cpp
index f3039b46d..fb8e04928 100644
--- a/src/economy.cpp
+++ b/src/economy.cpp
@@ -363,25 +363,12 @@ void ChangeOwnershipOfCompanyItems(Owner old_owner, Owner new_owner)
 	}
 
 	{
-		int num_train = 0;
-		int num_road = 0;
-		int num_ship = 0;
-		int num_aircraft = 0;
-		Vehicle *v;
-
-		/*  Determine Ids for the new vehicles */
-		FOR_ALL_VEHICLES(v) {
-			if (v->owner == new_owner) {
-				switch (v->type) {
-					case VEH_TRAIN:    if (IsFrontEngine(v)) num_train++; break;
-					case VEH_ROAD:     if (IsRoadVehFront(v)) num_road++; break;
-					case VEH_SHIP:     num_ship++; break;
-					case VEH_AIRCRAFT: if (IsNormalAircraft(v)) num_aircraft++; break;
-					default: break;
-				}
-			}
-		}
+		FreeUnitIDGenerator unitidgen[] = {
+			FreeUnitIDGenerator(VEH_TRAIN, new_owner), FreeUnitIDGenerator(VEH_ROAD,     new_owner),
+			FreeUnitIDGenerator(VEH_SHIP,  new_owner), FreeUnitIDGenerator(VEH_AIRCRAFT, new_owner)
+		};
 
+		Vehicle *v;
 		FOR_ALL_VEHICLES(v) {
 			if (v->owner == old_owner && IsCompanyBuildableVehicleType(v->type)) {
 				if (new_owner == INVALID_OWNER) {
@@ -391,13 +378,7 @@ void ChangeOwnershipOfCompanyItems(Owner old_owner, Owner new_owner)
 					v->colormap = PAL_NONE;
 					v->group_id = DEFAULT_GROUP;
 					if (IsEngineCountable(v)) GetCompany(new_owner)->num_engines[v->engine_type]++;
-					switch (v->type) {
-						case VEH_TRAIN:    if (IsFrontEngine(v)) v->unitnumber = ++num_train; break;
-						case VEH_ROAD:     if (IsRoadVehFront(v)) v->unitnumber = ++num_road; break;
-						case VEH_SHIP:     v->unitnumber = ++num_ship; break;
-						case VEH_AIRCRAFT: if (IsNormalAircraft(v)) v->unitnumber = ++num_aircraft; break;
-						default: NOT_REACHED();
-					}
+					if (v->IsPrimaryVehicle()) v->unitnumber = unitidgen[v->type].NextID();
 				}
 			}
 		}
diff --git a/src/vehicle.cpp b/src/vehicle.cpp
index 51166d6c0..5bdbc228d 100644
--- a/src/vehicle.cpp
+++ b/src/vehicle.cpp
@@ -53,6 +53,7 @@
 #include "core/alloc_func.hpp"
 #include "core/smallmap_type.hpp"
 #include "vehiclelist.h"
+#include "core/mem_func.hpp"
 #include "depot_func.h"
 
 #include "table/sprites.h"
@@ -1867,51 +1868,47 @@ VehicleEnterTileStatus VehicleEnterTile(Vehicle *v, TileIndex tile, int x, int y
 	return _tile_type_procs[GetTileType(tile)]->vehicle_enter_tile_proc(v, tile, x, y);
 }
 
-UnitID GetFreeUnitNumber(VehicleType type)
+FreeUnitIDGenerator::FreeUnitIDGenerator(VehicleType type, CompanyID owner) : cache(NULL), maxid(0), curid(0)
 {
-	UnitID max = 0;
-	const Vehicle *u;
-	static bool *cache = NULL;
-	static UnitID gmax = 0;
-
-	switch (type) {
-		case VEH_TRAIN:    max = _settings_game.vehicle.max_trains; break;
-		case VEH_ROAD:     max = _settings_game.vehicle.max_roadveh; break;
-		case VEH_SHIP:     max = _settings_game.vehicle.max_ships; break;
-		case VEH_AIRCRAFT: max = _settings_game.vehicle.max_aircraft; break;
-		default: NOT_REACHED();
+	/* Find maximum */
+	const Vehicle *v;
+	FOR_ALL_VEHICLES(v) {
+		if (v->type == type && v->owner == owner) {
+			this->maxid = max<UnitID>(this->maxid, v->unitnumber);
+		}
 	}
 
-	if (max == 0) {
-		/* we can't build any of this kind of vehicle, so we just return 1 instead of looking for a free number
-		 * a max of 0 will cause the following code to write to a NULL pointer
-		 * We know that 1 is bigger than the max allowed vehicle number, so it's the same as returning something, that is too big
-		 */
-		return 1;
-	}
+	if (this->maxid == 0) return;
 
-	if (max > gmax) {
-		gmax = max;
-		free(cache);
-		cache = MallocT<bool>(max + 1);
-	}
+	this->maxid++; // so there is space for last item (with v->unitnumber == maxid)
+	this->maxid++; // this one will always be free (well, it will fail when there are 65535 units, so this overflows)
+
+	this->cache = MallocT<bool>(this->maxid);
 
-	/* Clear the cache */
-	memset(cache, 0, (max + 1) * sizeof(*cache));
+	MemSetT(this->cache, 0, this->maxid);
 
 	/* Fill the cache */
-	FOR_ALL_VEHICLES(u) {
-		if (u->type == type && u->owner == _current_company && u->unitnumber != 0 && u->unitnumber <= max)
-			cache[u->unitnumber] = true;
+	FOR_ALL_VEHICLES(v) {
+		if (v->type == type && v->owner == owner) {
+			this->cache[v->unitnumber] = true;
+		}
 	}
+}
 
-	/* Find the first unused unit number */
-	UnitID unit = 1;
-	for (; unit <= max; unit++) {
-		if (!cache[unit]) break;
-	}
+UnitID FreeUnitIDGenerator::NextID()
+{
+	if (this->maxid <= this->curid) return ++this->curid;
+
+	while (this->cache[++this->curid]) { } // it will stop, we reserved more space than needed
+
+	return this->curid;
+}
+
+UnitID GetFreeUnitNumber(VehicleType type)
+{
+	FreeUnitIDGenerator gen(type, _current_company);
 
-	return unit;
+	return gen.NextID();
 }
 
 
diff --git a/src/vehicle_base.h b/src/vehicle_base.h
index d051f79d9..15f2df1aa 100644
--- a/src/vehicle_base.h
+++ b/src/vehicle_base.h
@@ -647,6 +647,28 @@ static inline bool IsValidVehicleID(uint index)
 	return index < GetVehiclePoolSize() && GetVehicle(index)->IsValid();
 }
 
+
+/** Generates sequence of free UnitID numbers */
+struct FreeUnitIDGenerator {
+	bool *cache;  ///< array of occupied unit id numbers
+	UnitID maxid; ///< maximum ID at the moment of constructor call
+	UnitID curid; ///< last ID returned ; 0 if none
+
+	/** Initializes the structure. Vehicle unit numbers are supposed not to change after
+	 * struct initialization, except after each call to this->NextID() the returned value
+	 * is assigned to a vehicle.
+	 * @param type type of vehicle
+	 * @param owner owner of vehicles
+	 */
+	FreeUnitIDGenerator(VehicleType type, CompanyID owner);
+
+	/** Returns next free UnitID. Supposes the last returned value was assigned to a vehicle. */
+	UnitID NextID();
+
+	/** Releases allocated memory */
+	~FreeUnitIDGenerator() { free(this->cache); }
+};
+
 /* Returns order 'index' of a vehicle or NULL when it doesn't exists */
 static inline Order *GetVehicleOrder(const Vehicle *v, int index)
 {