Ensure all images are freed, fix leak

src/drawing/Image.cpp

@@ -37,6 +37,7 @@ struct ImageList
 
 static bool                 _initialised = false;
 static std::list<ImageList> _freeLists;
+static uint32               _allocatedImageCount;
 
 #ifdef DEBUG
 static std::list<ImageList> _allocatedLists;
@@ -52,6 +53,23 @@ static bool AllocatedListContains(uint32 baseImageId, uint32 count)
         });
     return contains;
 }
+
+static bool AllocatedListRemove(uint32 baseImageId, uint32 count)
+{
+    auto foundItem = std::find_if(
+        _allocatedLists.begin(),
+        _allocatedLists.end(),
+        [baseImageId, count](const ImageList &imageList) -> bool
+        {
+            return imageList.BaseId == baseImageId && imageList.Count == count;
+        });
+    if (foundItem != _allocatedLists.end())
+    {
+        _allocatedLists.erase(foundItem);
+        return true;
+    }
+    return false;
+}
 #endif
 
 static uint32 AllocateImageList(uint32 count)
@@ -66,6 +84,7 @@ static uint32 AllocateImageList(uint32 count)
 #ifdef DEBUG
         _allocatedLists.clear();
 #endif
+        _allocatedImageCount = 0;
     }
 
     for (auto it = _freeLists.begin(); it != _freeLists.end(); it++)
@@ -84,6 +103,7 @@ static uint32 AllocateImageList(uint32 count)
 #ifdef DEBUG
             _allocatedLists.push_back({ imageList.BaseId, count });
 #endif
+            _allocatedImageCount += count;
             return imageList.BaseId;
         }
     }
@@ -96,10 +116,11 @@ static void FreeImageList(uint32 baseImageId, uint32 count)
     Guard::Assert(baseImageId >= BASE_IMAGE_ID, GUARD_LINE);
 
 #ifdef DEBUG
-    bool contains = AllocatedListContains(baseImageId, count);
+    bool contains = AllocatedListRemove(baseImageId, count);
     Guard::Assert(contains, GUARD_LINE);
 #endif
-    
+    _allocatedImageCount -= count;
+
     for (auto it = _freeLists.begin(); it != _freeLists.end(); it++)
     {
         if (it->BaseId + it->Count == baseImageId) 
@@ -114,8 +135,7 @@ static void FreeImageList(uint32 baseImageId, uint32 count)
             return;
         }
     }
-    
+
-    // TODO validate that this was an allocated list
     _freeLists.push_back({ baseImageId, count });
 }
 
@@ -143,7 +163,7 @@ extern "C"
 
     void gfx_object_free_images(uint32 baseImageId, uint32 count)
     {
-        if (baseImageId != 0)
+        if (baseImageId != INVALID_IMAGE_ID)
         {
             // Zero the G1 elements so we don't have invalid pointers
             // and data lying about
@@ -157,4 +177,16 @@ extern "C"
             FreeImageList(baseImageId, count);
         }
     }
+
+    void gfx_object_check_all_images_freed()
+    {
+        if (_allocatedImageCount != 0)
+        {
+#if DEBUG
+            Guard::Assert(_allocatedImageCount == 0, "%u images were not freed", _allocatedImageCount);
+#else
+            Console::Error::WriteLine("%u images were not freed", _allocatedImageCount);
+#endif
+        }
+    }
 }

src/drawing/drawing.h

@@ -189,6 +189,7 @@ void gfx_unload_g2();
 rct_g1_element* gfx_get_g1_element(int image_id);
 uint32 gfx_object_allocate_images(const rct_g1_element * images, uint32 count);
 void gfx_object_free_images(uint32 baseImageId, uint32 count);
+void gfx_object_check_all_images_freed();
 void sub_68371D();
 void FASTCALL gfx_rle_sprite_to_buffer(const uint8* RESTRICT source_bits_pointer, uint8* RESTRICT dest_bits_pointer, const uint8* RESTRICT palette_pointer, const rct_drawpixelinfo * RESTRICT dpi, int image_type, int source_y_start, int height, int source_x_start, int width);
 void FASTCALL gfx_draw_sprite(rct_drawpixelinfo *dpi, int image_id, int x, int y, uint32 tertiary_colour);

src/object/WaterObject.cpp

@@ -52,6 +52,7 @@ void WaterObject::Load()
 
 void WaterObject::Unload()
 {
+    gfx_object_free_images(_legacyType.image_id, GetImageTable()->GetCount());
     language_free_object_string(_legacyType.string_idx);
 
     _legacyType.string_idx = 0;

src/rct2.c

@@ -140,6 +140,7 @@ void rct2_quit()
 void rct2_dispose()
 {
 	object_manager_unload_all_objects();
+	gfx_object_check_all_images_freed();
 	gfx_unload_g2();
 	gfx_unload_g1();
 }