Erase zero-count entries in AliveVariableHash to stop unbounded memory growth

nautilus/src/nautilus/tracing/ExceptionBasedTraceContext.hpp

@@ -9,14 +9,15 @@
 #include "nautilus/tracing/TracingInterface.hpp"
 #include "tag/Tag.hpp"
 #include "tag/TagRecorder.hpp"
-#include <array>
+#include <cstddef>
 #include <cstdint>
 #include <cstring>
 #include <functional>
 #include <list>
 #include <memory>
 #include <unordered_map>
 #include <unordered_set>
+#include <vector>
 
 namespace nautilus {
 class NautilusFunctionDefinition;
@@ -55,6 +56,8 @@ inline size_t getStaticVarValue(const StaticVarHolder& holder) {
  * - Each variable ID is mixed with a constant multiplier for better hash distribution
  * - The hash incorporates both variable identity (ID) and reference count
  * - Uses unordered_map for sparse storage (only allocates for variables that exist)
+ * - Entries are erased when their count reaches zero, so the map stays bounded by the
+ *   peak number of alive variables instead of growing with every value ref ever seen
  *
  * Performance characteristics:
  * - increment(): O(1) average - hash map lookup + two XOR operations, two multiplications
@@ -96,13 +99,23 @@ class AliveVariableHash {
 	 * The hash is updated by XOR-ing out the old contribution ((id * HASH_MULTIPLIER) * old_count)
 	 * and XOR-ing in the new contribution ((id * HASH_MULTIPLIER) * new_count).
 	 *
+	 * Entries that reach a count of zero are erased. This is hash-neutral (a zero count
+	 * contributes 0 to the XOR hash, identical to an absent entry) and keeps the map size
+	 * proportional to the number of currently-alive variables. Without the erase, the map
+	 * grows with every value ref ever seen and — because the trace context is thread_local —
+	 * persists across trace iterations and across traced functions.
+	 *
 	 * @param id Variable identifier (32-bit value)
 	 */
 	inline void decrement(uint32_t id) noexcept {
-		uint32_t& c = counts[id];
+		const auto it = counts.try_emplace(id, 0).first;
+		uint32_t& c = it->second;
 		alive_hash ^= (id * HASH_MULTIPLIER) * c;
 		--c;
 		alive_hash ^= (id * HASH_MULTIPLIER) * c;
+		if (c == 0) {
+			counts.erase(it);
+		}
 	}
 
 	/**
@@ -118,17 +131,31 @@ class AliveVariableHash {
 		return alive_hash;
 	}
 
+	/**
+	 * @brief Returns the number of tracked entries.
+	 *
+	 * Since zero-count entries are erased by decrement(), this is the number of
+	 * currently-alive variables (variables with a non-zero reference count).
+	 *
+	 * @return Number of entries in the underlying map
+	 */
+	inline size_t size() const noexcept {
+		return counts.size();
+	}
+
 	/**
 	 * @brief Resets all reference counts and hash to initial state.
 	 *
 	 * This efficiently clears all counts without creating a temporary object.
-	 * Optimized: if hash is already 0, we assume counts are already empty and skip the clear.
+	 * Since decrement() erases entries when they reach zero, a balanced trace iteration
+	 * leaves the map empty and this is a no-op. The emptiness check (rather than a hash
+	 * check) also guards against the rare XOR collision where live entries hash to 0.
 	 */
 	inline void reset() noexcept {
-		if (alive_hash != 0) {
+		if (!counts.empty()) {
 			counts.clear();
-			alive_hash = 0;
 		}
+		alive_hash = 0;
 	}
 };
 

nautilus/test/CMakeLists.txt

@@ -14,6 +14,7 @@ add_subdirectory(benchmark)
 add_subdirectory(execution-tests)
 if (ENABLE_TRACING)
     add_subdirectory(ir-pass-tests)
+    add_subdirectory(tracing-tests)
 endif ()
 add_subdirectory(val-tests)
 

nautilus/test/tracing-tests/AliveVariableHashTest.cpp

@@ -0,0 +1,98 @@
+#include "nautilus/tracing/ExceptionBasedTraceContext.hpp"
+#include <catch2/catch_test_macros.hpp>
+
+namespace nautilus::tracing {
+
+TEST_CASE("AliveVariableHash tracks the number of alive variables", "[AliveVariableHash]") {
+	AliveVariableHash avh;
+	REQUIRE(avh.size() == 0);
+	REQUIRE(avh.hash() == 0);
+
+	avh.increment(1);
+	avh.increment(2);
+	avh.increment(3);
+	REQUIRE(avh.size() == 3);
+
+	avh.decrement(2);
+	REQUIRE(avh.size() == 2);
+
+	avh.decrement(1);
+	avh.decrement(3);
+	REQUIRE(avh.size() == 0);
+	REQUIRE(avh.hash() == 0);
+}
+
+TEST_CASE("AliveVariableHash erasure is hash-neutral", "[AliveVariableHash]") {
+	// A variable that was seen and fully released must leave the hash identical
+	// to one that was never seen. Snapshots key on this hash, so any difference
+	// would change trace deduplication and the generated IR.
+	AliveVariableHash seenAndReleased;
+	seenAndReleased.increment(7);
+	seenAndReleased.increment(42);
+	seenAndReleased.decrement(42);
+
+	AliveVariableHash neverSeen;
+	neverSeen.increment(7);
+
+	REQUIRE(seenAndReleased.hash() == neverSeen.hash());
+	REQUIRE(seenAndReleased.size() == neverSeen.size());
+}
+
+TEST_CASE("AliveVariableHash re-incrementing a released id reproduces the original hash", "[AliveVariableHash]") {
+	AliveVariableHash avh;
+	avh.increment(5);
+	const uint64_t hashWithFive = avh.hash();
+
+	avh.decrement(5);
+	avh.increment(5);
+	REQUIRE(avh.hash() == hashWithFive);
+	REQUIRE(avh.size() == 1);
+}
+
+TEST_CASE("AliveVariableHash hash reflects reference counts", "[AliveVariableHash]") {
+	AliveVariableHash once;
+	once.increment(9);
+
+	AliveVariableHash twice;
+	twice.increment(9);
+	twice.increment(9);
+
+	REQUIRE(once.hash() != twice.hash());
+	REQUIRE(twice.size() == 1);
+
+	twice.decrement(9);
+	REQUIRE(twice.hash() == once.hash());
+}
+
+TEST_CASE("AliveVariableHash id churn neither perturbs the hash nor grows the map", "[AliveVariableHash]") {
+	// Monotonically increasing ids that immediately die mimic the tracing access
+	// pattern: value refs are assigned monotonically and never reused. The churn must
+	// not affect the hash or the size while a long-lived entry survives it.
+	AliveVariableHash avh;
+	avh.increment(1000000);
+	const uint64_t hashWithSurvivor = avh.hash();
+
+	for (uint32_t id = 0; id < 10000; ++id) {
+		avh.increment(id);
+		avh.decrement(id);
+		REQUIRE(avh.hash() == hashWithSurvivor);
+	}
+	REQUIRE(avh.size() == 1);
+
+	avh.decrement(1000000);
+	REQUIRE(avh.size() == 0);
+	REQUIRE(avh.hash() == 0);
+}
+
+TEST_CASE("AliveVariableHash reset clears live entries", "[AliveVariableHash]") {
+	AliveVariableHash avh;
+	avh.increment(1);
+	avh.increment(1);
+	avh.increment(2);
+
+	avh.reset();
+	REQUIRE(avh.size() == 0);
+	REQUIRE(avh.hash() == 0);
+}
+
+} // namespace nautilus::tracing

nautilus/test/tracing-tests/CMakeLists.txt

@@ -0,0 +1,15 @@
+include(CTest)
+include(Catch)
+
+add_executable(nautilus-tracing-unit-tests
+        AliveVariableHashTest.cpp
+)
+
+target_link_libraries(nautilus-tracing-unit-tests PRIVATE nautilus Catch2::Catch2WithMain nautilus-sanitizer-suppressions)
+target_include_directories(nautilus-tracing-unit-tests PRIVATE
+        $<BUILD_INTERFACE:${CMAKE_CURRENT_SOURCE_DIR}/../common>
+        $<BUILD_INTERFACE:${CMAKE_CURRENT_SOURCE_DIR}/../../src>
+)
+
+list(APPEND CMAKE_MODULE_PATH ${catch2_SOURCE_DIR}/extras)
+catch_discover_tests(nautilus-tracing-unit-tests EXTRA_ARGS --allow-running-no-tests)