Fix re-routing of unrelated cargo when removing invalidated link graph flow

src/cargopacket.cpp

@@ -460,15 +460,20 @@ void VehicleCargoList::ShiftCargo(Taction action)
  *                 will be kept and the loop will be aborted.
  *                 The second method parameter can be appended to, to prepend items to the packet list
  * @param action Action instance to be applied.
+ * @param filter Cargo packet filter.
  */
-template<class Taction>
-void VehicleCargoList::ShiftCargoWithFrontInsert(Taction action)
+template<class Taction, class Tfilter>
+void VehicleCargoList::ShiftCargoWithFrontInsert(Taction action, Tfilter filter)
 {
 	std::vector<CargoPacket *> packets_to_front_insert;
 
 	Iterator it(this->packets.begin());
 	while (it != this->packets.end() && action.MaxMove() > 0) {
 		CargoPacket *cp = *it;
+		if (!filter(cp)) {
+			++it;
+			continue;
+		}
 		if (action(cp, packets_to_front_insert)) {
 			it = this->packets.erase(it);
 		} else {
@@ -858,7 +863,23 @@ uint VehicleCargoList::Truncate(uint max_move)
 uint VehicleCargoList::Reroute(uint max_move, VehicleCargoList *dest, StationID avoid, StationID avoid2, const GoodsEntry *ge)
 {
 	max_move = std::min(this->action_counts[MTA_TRANSFER], max_move);
-	this->ShiftCargoWithFrontInsert(VehicleCargoReroute(this, dest, max_move, avoid, avoid2, ge));
+	this->ShiftCargoWithFrontInsert(VehicleCargoReroute(this, dest, max_move, avoid, avoid2, ge), [](CargoPacket *cp) { return true; });
+	return max_move;
+}
+
+/**
+ * Routes packets with station "avoid" as next hop to a different place, for a specific source station.
+ * @param max_move Maximum amount of cargo to move.
+ * @param dest List to prepend the cargo to.
+ * @param source Source station.
+ * @param avoid Station to exclude from routing and current next hop of packets to reroute.
+ * @param avoid2 Additional station to exclude from routing.
+ * @param ge GoodsEntry to get the routing info from.
+ */
+uint VehicleCargoList::RerouteFromSource(uint max_move, VehicleCargoList *dest, StationID source, StationID avoid, StationID avoid2, const GoodsEntry *ge)
+{
+	max_move = std::min(this->action_counts[MTA_TRANSFER], max_move);
+	this->ShiftCargoWithFrontInsert(VehicleCargoReroute(this, dest, max_move, avoid, avoid2, ge), [source](CargoPacket *cp) { return cp->SourceStation() == source; });
 	return max_move;
 }
 
@@ -946,6 +967,67 @@ uint StationCargoList::ShiftCargo(Taction action, StationIDStack next, bool incl
 	return max_move - action.MaxMove();
 }
 
+/**
+ * Shifts cargo from the front of the packet list for a specific station
+ * from a specific source station and applies some action to it.
+ * @tparam Taction Action class or function to be used. It should define
+ *                 "bool operator()(CargoPacket *)". If true is returned the
+ *                 cargo packet will be removed from the list. Otherwise it
+ *                 will be kept and the loop will be aborted.
+ * @param action Action instance to be applied.
+ * @param source Source station.
+ * @param next Next hop the cargo wants to visit.
+ * @return True if all packets with the given next hop have been removed,
+ *         False otherwise.
+ */
+template <class Taction>
+bool StationCargoList::ShiftCargoFromSource(Taction &action, StationID source, StationID next)
+{
+	for (Iterator it = this->packets.lower_bound(next); it != this->packets.end() && it.GetKey() == next;) {
+		if (action.MaxMove() == 0) return false;
+		CargoPacket *cp = *it;
+		if (cp->SourceStation() != source) {
+			++it;
+			continue;
+		}
+		if (action(cp)) {
+			it = this->packets.erase(it);
+		} else {
+			return false;
+		}
+	}
+	return true;
+}
+
+/**
+ * Shifts cargo from the front of the packet list for a specific station
+ * and optional also from the list for "any station", then applies some action
+ * to it, for a specific source station.
+ * @tparam Taction Action class or function to be used. It should define
+ *                 "bool operator()(CargoPacket *)". If true is returned the
+ *                 cargo packet will be removed from the list. Otherwise it
+ *                 will be kept and the loop will be aborted.
+ * @param action Action instance to be applied.
+ * @param source Source station.
+ * @param next Next hop the cargo wants to visit.
+ * @param include_invalid If cargo from the INVALID_STATION list should be
+ *                        used if necessary.
+ * @return Amount of cargo actually moved.
+ */
+template <class Taction>
+uint StationCargoList::ShiftCargoFromSource(Taction action, StationID source, StationIDStack next, bool include_invalid)
+{
+	uint max_move = action.MaxMove();
+	while (!next.IsEmpty()) {
+		this->ShiftCargoFromSource(action, source, next.Pop());
+		if (action.MaxMove() == 0) break;
+	}
+	if (include_invalid && action.MaxMove() > 0) {
+		this->ShiftCargoFromSource(action, source, INVALID_STATION);
+	}
+	return max_move - action.MaxMove();
+}
+
 uint StationCargoList::AvailableViaCount(StationID next) const
 {
 	uint count = 0;
@@ -1059,6 +1141,19 @@ uint StationCargoList::Reroute(uint max_move, StationCargoList *dest, StationID
 	return this->ShiftCargo(StationCargoReroute(this, dest, max_move, avoid, avoid2, ge), avoid, false);
 }
 
+/**
+ * Routes packets with station "avoid" as next hop to a different place.
+ * @param max_move Maximum amount of cargo to move.
+ * @param dest List to append the cargo to.
+ * @param avoid Station to exclude from routing and current next hop of packets to reroute.
+ * @param avoid2 Additional station to exclude from routing.
+ * @param ge GoodsEntry to get the routing info from.
+ */
+uint StationCargoList::RerouteFromSource(uint max_move, StationCargoList *dest, StationID source, StationID avoid, StationID avoid2, const GoodsEntry *ge)
+{
+	return this->ShiftCargoFromSource(StationCargoReroute(this, dest, max_move, avoid, avoid2, ge), source, avoid, false);
+}
+
 /*
  * We have to instantiate everything we want to be usable.
  */

src/cargopacket.h

@@ -298,8 +298,8 @@ protected:
 	template<class Taction>
 	void ShiftCargo(Taction action);
 
-	template<class Taction>
-	void ShiftCargoWithFrontInsert(Taction action);
+	template<class Taction, class Tfilter>
+	void ShiftCargoWithFrontInsert(Taction action, Tfilter filter);
 
 	template<class Taction>
 	void PopCargo(Taction action);
@@ -467,6 +467,7 @@ public:
 	uint Shift(uint max_move, VehicleCargoList *dest);
 	uint Truncate(uint max_move = UINT_MAX);
 	uint Reroute(uint max_move, VehicleCargoList *dest, StationID avoid, StationID avoid2, const GoodsEntry *ge);
+	uint RerouteFromSource(uint max_move, VehicleCargoList *dest, StationID source, StationID avoid, StationID avoid2, const GoodsEntry *ge);
 
 	/**
 	 * Are the two CargoPackets mergeable in the context of
@@ -521,6 +522,12 @@ public:
 	template<class Taction>
 	uint ShiftCargo(Taction action, StationIDStack next, bool include_invalid);
 
+	template <class Taction>
+	bool ShiftCargoFromSource(Taction &action, StationID source, StationID next);
+
+	template<class Taction>
+	uint ShiftCargoFromSource(Taction action, StationID source, StationIDStack next, bool include_invalid);
+
 	void Append(CargoPacket *cp, StationID next);
 
 	/**
@@ -585,6 +592,7 @@ public:
 	uint Load(uint max_move, VehicleCargoList *dest, TileIndex load_place, StationIDStack next);
 	uint Truncate(uint max_move = UINT_MAX, StationCargoAmountMap *cargo_per_source = nullptr);
 	uint Reroute(uint max_move, StationCargoList *dest, StationID avoid, StationID avoid2, const GoodsEntry *ge);
+	uint RerouteFromSource(uint max_move, StationCargoList *dest, StationID source, StationID avoid, StationID avoid2, const GoodsEntry *ge);
 
 	void AfterLoadIncreaseReservationCount(uint count)
 	{

src/linkgraph/linkgraphjob.cpp

@@ -151,11 +151,19 @@ void LinkGraphJob::FinaliseJob()
 		for (FlowStatMap::iterator it(ge.flows.begin()); it != ge.flows.end();) {
 			FlowStatMap::iterator new_it = flows.find(it->GetOrigin());
 			if (new_it == flows.end()) {
-				bool should_erase = true;
 				if (_settings_game.linkgraph.GetDistributionType(this->Cargo()) != DT_MANUAL) {
-					should_erase = it->Invalidate();
-				}
-				if (should_erase) {
+					if (it->Invalidate()) {
+						FlowStat shares(INVALID_STATION, INVALID_STATION, 1);
+						it->SwapShares(shares);
+						it = ge.flows.erase(it);
+						for (FlowStat::const_iterator shares_it(shares.begin());
+								shares_it != shares.end(); ++shares_it) {
+							RerouteCargoFromSource(st, this->Cargo(), it->GetOrigin(), shares_it->second, st->index);
+						}
+					} else {
+						++it;
+					}
+				} else {
 					FlowStat shares(INVALID_STATION, INVALID_STATION, 1);
 					it->SwapShares(shares);
 					it = ge.flows.erase(it);
@@ -163,8 +171,6 @@ void LinkGraphJob::FinaliseJob()
 							shares_it != shares.end(); ++shares_it) {
 						RerouteCargo(st, this->Cargo(), shares_it->second, st->index);
 					}
-				} else {
-					++it;
 				}
 			} else {
 				it->SwapShares(*new_it);

src/station_cmd.cpp

@@ -4094,6 +4094,31 @@ void RerouteCargo(Station *st, CargoID c, StationID avoid, StationID avoid2)
 	}
 }
 
+/**
+ * Reroute cargo of type c from source at station st or in any vehicles unloading there.
+ * Make sure the cargo's new next hop is neither "avoid" nor "avoid2".
+ * @param st Station to be rerouted at.
+ * @param c Type of cargo.
+ * @param source Source station.
+ * @param avoid Original next hop of cargo, avoid this.
+ * @param avoid2 Another station to be avoided when rerouting.
+ */
+void RerouteCargoFromSource(Station *st, CargoID c, StationID source, StationID avoid, StationID avoid2)
+{
+	GoodsEntry &ge = st->goods[c];
+
+	/* Reroute cargo in station. */
+	ge.cargo.RerouteFromSource(UINT_MAX, &ge.cargo, source, avoid, avoid2, &ge);
+
+	/* Reroute cargo staged to be transferred. */
+	for (Vehicle *v : st->loading_vehicles) {
+		for (; v != nullptr; v = v->Next()) {
+			if (v->cargo_type != c) continue;
+			v->cargo.RerouteFromSource(UINT_MAX, &v->cargo, source, avoid, avoid2, &ge);
+		}
+	}
+}
+
 btree::btree_set<VehicleID> _delete_stale_links_vehicle_cache;
 
 void ClearDeleteStaleLinksVehicleCache()

src/station_func.h

@@ -55,6 +55,7 @@ bool SplitGroundSpriteForOverlay(const TileInfo *ti, SpriteID *ground, RailTrack
 void IncreaseStats(Station *st, const Vehicle *v, StationID next_station_id);
 void IncreaseStats(Station *st, CargoID cargo, StationID next_station_id, uint capacity, uint usage, EdgeUpdateMode mode);
 void RerouteCargo(Station *st, CargoID c, StationID avoid, StationID avoid2);
+void RerouteCargoFromSource(Station *st, CargoID c, StationID source, StationID avoid, StationID avoid2);
 
 /**
  * Calculates the maintenance cost of a number of station tiles.