From bb07d84314f9fcabaf3ec2a5b02b69298232ed9d Mon Sep 17 00:00:00 2001
From: hdpoliveira <35645457+hdpoliveira@users.noreply.github.com>
Date: Mon, 4 May 2020 19:43:25 -0300
Subject: [PATCH] Implement Train class with iterators (#11486)

---
 src/openrct2/ride/Vehicle.cpp | 90 ++++++++++++++++++++++++++++++-----
 1 file changed, 79 insertions(+), 11 deletions(-)

diff --git a/src/openrct2/ride/Vehicle.cpp b/src/openrct2/ride/Vehicle.cpp
index 01a283ac7e..b1e2f57336 100644
--- a/src/openrct2/ride/Vehicle.cpp
+++ b/src/openrct2/ride/Vehicle.cpp
@@ -812,18 +812,87 @@ void Vehicle::Invalidate()
     invalidate_sprite_2(this);
 }
 
-static int32_t get_train_mass(const Vehicle* first_vehicle)
+namespace
+{
+    template<typename T> class TrainIterator;
+    template<typename T> class Train
+    {
+    public:
+        explicit Train(T* vehicle)
+            : FirstCar(vehicle)
+        {
+            assert(FirstCar->IsHead());
+        }
+        int32_t Mass();
+
+        friend class TrainIterator<T>;
+        using iterator = TrainIterator<T>;
+        iterator begin()
+        {
+            return iterator{ FirstCar };
+        }
+        iterator end()
+        {
+            return iterator{};
+        }
+
+    private:
+        T* FirstCar;
+    };
+    template<typename T> class TrainIterator
+    {
+    public:
+        using iterator = TrainIterator;
+        using iterator_category = std::forward_iterator_tag;
+        using value_type = T;
+        using pointer = T*;
+        using reference = T&;
+
+        TrainIterator() = default;
+        explicit TrainIterator(T* vehicle)
+            : Current(vehicle)
+        {
+        }
+        reference operator*()
+        {
+            return *Current;
+        }
+        iterator& operator++()
+        {
+            if (NextVehicleId != SPRITE_INDEX_NULL)
+            {
+                Current = GET_VEHICLE(NextVehicleId);
+                NextVehicleId = Current->next_vehicle_on_train;
+            }
+            else
+            {
+                Current = nullptr;
+            }
+            return *this;
+        }
+        iterator operator++(int)
+        {
+            iterator temp = *this;
+            ++*this;
+            return temp;
+        }
+        bool operator!=(const iterator& other)
+        {
+            return Current != other.Current;
+        }
+
+    private:
+        T* Current = nullptr;
+        uint16_t NextVehicleId = SPRITE_INDEX_NULL;
+    };
+} // namespace
+
+template<typename T> int32_t Train<T>::Mass()
 {
     int32_t totalMass = 0;
-
-    for (const Vehicle* vehicle = first_vehicle; vehicle != nullptr;)
+    for (const auto& vehicle : *this)
     {
-        totalMass += vehicle->mass;
-
-        if (vehicle->next_vehicle_on_train == SPRITE_INDEX_NULL)
-            break;
-
-        vehicle = GET_VEHICLE(vehicle->next_vehicle_on_train);
+        totalMass += vehicle.mass;
     }
 
     return totalMass;
@@ -881,8 +950,7 @@ bool Vehicle::SoundCanPlay() const
  */
 uint16_t Vehicle::GetSoundPriority() const
 {
-    int32_t trainMass = get_train_mass(this);
-    int32_t result = trainMass + (std::abs(velocity) >> 13);
+    int32_t result = Train(this).Mass() + (std::abs(velocity) >> 13);
     rct_vehicle_sound* vehicle_sound = &gVehicleSoundList[0];
 
     while (vehicle_sound->id != sprite_index)