From 190dc246b36c129b633368a91169c9ae157badb7 Mon Sep 17 00:00:00 2001
From: konimarti <30975830+konimarti@users.noreply.github.com>
Date: Fri, 10 Jan 2025 13:33:33 +0100
Subject: [PATCH] mem: add macro to assert on memory leak in scope (#1792)

* mem: add macro to assert on memory leak in scope

Implement `mem::@assert_leak` to assert on a memory leak in the scope of
the macro body. Memory report for the leak can be disabled by setting
the boolean argument to false.

* fix: add conditional compilation flags

* Moved the code into `mem.c3` and made it a builtin.

---------

Co-authored-by: Christoffer Lerno <christoffer@aegik.com>
---
 lib/std/core/mem.c3 | 35 ++++++++++++++++++++++++++++++++++-
 releasenotes.md     |  1 +
 2 files changed, 35 insertions(+), 1 deletion(-)

diff --git a/lib/std/core/mem.c3 b/lib/std/core/mem.c3
index c7ed62a00..f508d46b0 100644
--- a/lib/std/core/mem.c3
+++ b/lib/std/core/mem.c3
@@ -490,6 +490,40 @@ macro void @report_heap_allocs_in_scope($enabled = true; @body())
 	@body();
 }
 
+<*
+ Assert on memory leak in the scope of the macro body.
+
+ @param $report "Set to false to disable memory report"
+*>
+macro void @assert_leak($report = true; @body()) @builtin
+{
+	$if env::DEBUG_SYMBOLS || $feature(MEMORY_ASSERTS):
+		TrackingAllocator tracker;
+		tracker.init(allocator::thread_allocator);
+		Allocator old_allocator = allocator::thread_allocator;
+		allocator::thread_allocator = &tracker;
+		defer
+		{
+			allocator::thread_allocator = old_allocator;
+			defer tracker.free();
+			usz allocated = tracker.allocated();
+			if (allocated)
+			{
+				DString report = dstring::new();
+				defer report.free();
+				$if $report:
+					report.append_char('\n');
+					(void)tracker.fprint_report(&report);
+				$endif
+				assert(allocated == 0, "Memory leak detected"
+					" (%d bytes allocated).%s",
+					allocated, report.str_view());
+			}
+		}
+	$endif
+	@body();
+}
+
 <*
  Allocate [size] bytes on the stack to use for allocation,
  with the heap allocator as the backing allocator.
@@ -832,4 +866,3 @@ fn void* __memcpy(void* dst, void* src, usz n) @weak @export("memcpy")
 	}
 	return dst;
 }
-
diff --git a/releasenotes.md b/releasenotes.md
index 91079338a..967cc15bb 100644
--- a/releasenotes.md
+++ b/releasenotes.md
@@ -83,6 +83,7 @@
 - Added `String.trim_left()` and `.trim_right()`.
 - Deprecation of several `&` macros.
 - Format functions for timedates.
+- Add `@assert_leak()` to assert on memory leaks in the scope.
 
 ## 0.6.5 Change list