From 79050dfdfb26789fd6496e34f2255ccbcf4cb0e5 Mon Sep 17 00:00:00 2001
From: arch1t3cht <arch1t3cht@gmail.com>
Date: Sun, 21 May 2023 20:26:20 +0200
Subject: [PATCH] lua: Fix memory leak on aegisub.cancel()

---
 src/auto4_lua.h           |  4 ++++
 src/auto4_lua_assfile.cpp | 15 +++++++++++----
 2 files changed, 15 insertions(+), 4 deletions(-)

diff --git a/src/auto4_lua.h b/src/auto4_lua.h
index 46068173e..787ce3a6b 100644
--- a/src/auto4_lua.h
+++ b/src/auto4_lua.h
@@ -78,6 +78,8 @@ namespace Automation4 {
 		std::deque<PendingCommit> pending_commits;
 		/// Lines to delete once processing complete successfully
 		std::vector<std::unique_ptr<AssEntry>> lines_to_delete;
+		/// Lines that were allocated here and need to be deleted if the script is cancelled.
+		std::vector<AssEntry *> allocated_lines;
 
 		/// Create copies of all of the lines in the script info section if it
 		/// hasn't already happened. This is done lazily, since it only needs
@@ -118,6 +120,8 @@ namespace Automation4 {
 		/// assumes a Lua representation of AssEntry on the top of the stack, and creates an AssEntry object of it
 		static std::unique_ptr<AssEntry> LuaToAssEntry(lua_State *L, AssFile *ass=nullptr);
 
+		std::unique_ptr<AssEntry> LuaToTrackedAssEntry(lua_State *L);
+
 		/// @brief Signal that the script using this file is now done running
 		/// @param set_undo If there's any uncommitted changes to the file,
 		///                 they will be automatically committed with this
diff --git a/src/auto4_lua_assfile.cpp b/src/auto4_lua_assfile.cpp
index 3aef3b769..e73a48920 100644
--- a/src/auto4_lua_assfile.cpp
+++ b/src/auto4_lua_assfile.cpp
@@ -328,6 +328,12 @@ namespace Automation4 {
 		return result;
 	}
 
+	std::unique_ptr<AssEntry> LuaAssFile::LuaToTrackedAssEntry(lua_State *L) {
+		std::unique_ptr<AssEntry> e = LuaToAssEntry(L, ass);
+		allocated_lines.push_back(e.get());
+		return e;
+	}
+
 	int LuaAssFile::ObjectIndexRead(lua_State *L)
 	{
 		switch (lua_type(L, 2)) {
@@ -453,7 +459,7 @@ namespace Automation4 {
 				// insert
 				CheckBounds(n);
 
-				auto e = LuaToAssEntry(L, ass);
+				auto e = LuaToTrackedAssEntry(L);
 				modification_type |= modification_mask(e.get());
 				QueueLineForDeletion(n - 1);
 				AssignLine(n - 1, std::move(e));
@@ -542,7 +548,7 @@ namespace Automation4 {
 
 		for (int i = 1; i <= n; i++) {
 			lua_pushvalue(L, i);
-			auto e = LuaToAssEntry(L, ass);
+			auto e = LuaToTrackedAssEntry(L);
 			modification_type |= modification_mask(e.get());
 
 			if (lines.empty()) {
@@ -586,7 +592,7 @@ namespace Automation4 {
 		new_entries.reserve(n - 1);
 		for (int i = 2; i <= n; i++) {
 			lua_pushvalue(L, i);
-			auto e = LuaToAssEntry(L, ass);
+			auto e = LuaToTrackedAssEntry(L);
 			modification_type |= modification_mask(e.get());
 			InsertLine(new_entries, i - 2, std::move(e));
 			lua_pop(L, 1);
@@ -625,7 +631,7 @@ namespace Automation4 {
 
 	int LuaAssFile::LuaParseKaraokeData(lua_State *L)
 	{
-		auto e = LuaToAssEntry(L, ass);
+		auto e = LuaToTrackedAssEntry(L);
 		auto dia = check_cast_constptr<AssDialogue>(e.get());
 		argcheck(L, !!dia, 1, "Subtitle line must be a dialogue line");
 
@@ -734,6 +740,7 @@ namespace Automation4 {
 	void LuaAssFile::Cancel()
 	{
 		for (auto& line : lines_to_delete) line.release();
+		for (AssEntry *line : allocated_lines) delete line;
 		references--;
 		if (!references) delete this;
 	}