From aaaf85993ac7cdd6cd29778db21aefc71a16e26f Mon Sep 17 00:00:00 2001
From: Runxi Yu <me@runxiyu.org>
Date: Tue, 23 Sep 2025 11:50:32 +0800
Subject: [PATCH] Fix memory leaks

---
 ds/map/btree/del.ha      |  5 +++--
 ds/map/btree/internal.ha | 10 ++++++++++
 ds/map/btree/iter.ha     | 12 ++++++++++++
 ds/map/hashmap/iter.ha   | 11 ++++++++++-
 ds/map/test.ha           |  7 +++++++

diff --git a/ds/map/btree/del.ha b/ds/map/btree/del.ha
index edc6a8e7d22da33512f58f68ab146768fc802f64..370dd3fda63f71741c3ffe60306849d4ea493cec 100644
--- a/ds/map/btree/del.ha
+++ b/ds/map/btree/del.ha
@@ -7,8 +7,9 @@ // Deletes an item from a [[map]]. Returns the removed value or void.
 export fn del(m: *map, key: []u8) (*opaque | void) = {
 	const r = delete_rec(m, m.root, key);
 	if (len(m.root.keys) == 0 && !m.root.leaf) {
-		m.root = m.root.children[0];
+		let old = m.root;
+		m.root = old.children[0];
+		node_shallow_finish(old);
 	};
 	return r;
 };
-
diff --git a/ds/map/btree/internal.ha b/ds/map/btree/internal.ha
index 022903ed1e5ec7326d004c31861432399375ae85..6b9d3a8e5fc6cb4e75e3a0ab7a63b8025332f5b9 100644
--- a/ds/map/btree/internal.ha
+++ b/ds/map/btree/internal.ha
@@ -46,6 +46,15 @@ 	})?;
 	return nd;
 };
 
+fn node_shallow_finish(n: *node) void = {
+	if (!n.leaf) {
+		free(n.children);
+	};
+	free(n.keys);
+	free(n.vals);
+	free(n);
+};
+
 fn split_child(m: *map, x: *node, i: size) (void | nomem) = {
 	const t = m.t;
 	let y = x.children[i];
@@ -128,6 +137,7 @@
 	delete(x.keys[i]);
 	delete(x.vals[i]);
 	delete(x.children[i + 1]);
+	node_shallow_finish(right);
 };
 
 fn ensure_child_has_space(m: *map, x: *node, i: size) void = {
diff --git a/ds/map/btree/iter.ha b/ds/map/btree/iter.ha
index c29058f2b9b47f4c3d2efc3db32864755864e9ec..37f39ead4134dfdd30094706e491ff57c938b240 100644
--- a/ds/map/btree/iter.ha
+++ b/ds/map/btree/iter.ha
@@ -7,6 +7,7 @@ 	vt: map::iterator,
 	nodes: []*node,
 	idxs: []size,
 	visit_key: []bool,
+	finished: bool,
 };
 
 const _itvt: map::vtable_iterator = map::vtable_iterator {
@@ -21,6 +22,7 @@ 		vt = &_itvt,
 		nodes = [],
 		idxs = [],
 		visit_key = [],
+		finished = false,
 	})?;
 
 	match (append(it.nodes, m.root)) {
@@ -40,6 +42,9 @@ 	return (it: *map::iterator);
 };
 
 export fn next(it: *iterator) (([]u8, *opaque) | done) = {
+	if (it.finished) {
+		return done;
+	};
 	for (len(it.nodes) != 0) {
 		let top = len(it.nodes) - 1;
 		let x = it.nodes[top];
@@ -106,5 +111,12 @@ 		delete(it.idxs[top]);
 		delete(it.visit_key[top]);
 	};
 
+	free(it.nodes);
+	free(it.idxs);
+	free(it.visit_key);
+	it.nodes = [];
+	it.idxs = [];
+	it.visit_key = [];
+	it.finished = true;
 	return done;
 };
diff --git a/ds/map/hashmap/iter.ha b/ds/map/hashmap/iter.ha
index a06f91c2611f540e4e0caae926964ad01bca1f89..2a3bec5183af8bbb2e8a74a439c8aeda5023c989 100644
--- a/ds/map/hashmap/iter.ha
+++ b/ds/map/hashmap/iter.ha
@@ -9,6 +9,7 @@ 	vt: map::iterator,
 	m: *map,
 	bi: size,
 	cur: nullable *map::iterator,
+	finished: bool,
 };
 
 const _itvt: map::vtable_iterator = map::vtable_iterator {
@@ -23,6 +24,7 @@ 		vt = &_itvt,
 		m = m,
 		bi = 0,
 		cur = null,
+		finished = false,
 	})?;
 	for (it.bi < m.n) {
 		let b = m.buckets[it.bi];
@@ -39,10 +41,16 @@ 	return (it: *map::iterator);
 };
 
 export fn next(it: *iterator) (([]u8, *opaque) | done) = {
+	if (it.finished) {
+		return done;
+	};
 	for (true) {
 		match (it.cur) {
 		case null =>
-			if (it.bi >= it.m.n) return done;
+			if (it.bi >= it.m.n) {
+				it.finished = true;
+				return done;
+			};
 			let b = it.m.buckets[it.bi];
 			it.bi += 1;
 			match (map::iter(b)) {
@@ -56,6 +64,7 @@ 			match (map::next(curi)) {
 			case let kv: ([]u8, *opaque) =>
 				return kv;
 			case done =>
+				free(curi);
 				it.cur = null;
 			};
 		};
diff --git a/ds/map/test.ha b/ds/map/test.ha
index e6b023db8861d9c8a0ad4855b2fb0205029bcb19..3027a03deb081faa6ee03b4fa2fbaa550a20c772 100644
--- a/ds/map/test.ha
+++ b/ds/map/test.ha
@@ -29,11 +29,13 @@ fn must_del(m: *map, key: []u8) (*opaque | void) = del(m, key);
 
 fn verify_iter_matches_oracle(m: *map, exp: oracle) void = {
 	let seen: []bool = alloc([false...], len(exp))!;
+	defer free(seen);
 
 	let it: *iterator = match (iter(m)) {
 	case let p: *iterator => yield p;
 	case nomem => abort("iter: out of memory");
 	};
+	defer free(it);
 
 	for (true) {
 		match (next(it)) {
@@ -72,6 +74,8 @@ export fn stress_test(m: *map, key_space: size) void = {
 	let empty: [KEY_LEN]u8 = [0...];
 	let keybufs: [][KEY_LEN]u8 = alloc([empty...], key_space)!;
 	let keys: [][]u8 = alloc([[0...]...], key_space)!;
+	defer free(keys);
+	defer free(keybufs);
 
 	for (let i = 0z; i < key_space; i += 1) {
 		for (let j = 8z; j < KEY_LEN; j += 1) keybufs[i][j] = 0xABu8;
@@ -80,9 +84,11 @@ 		keys[i] = keybufs[i][..];
 	};
 
 	let vals: []int = alloc([0...], key_space)!;
+	defer free(vals);
 	for (let i = 0z; i < key_space; i += 1) vals[i] = (i: int);
 
 	let exp: oracle = alloc([null...], key_space)!;
+	defer free(exp);
 
 	// Sequential inserts with immediate verification
 	for (let i = 0z; i < key_space; i += 1) {
@@ -239,6 +245,7 @@ 	let it_empty: *iterator = match (iter(m)) {
 	case let p: *iterator => yield p;
 	case nomem => abort("iter(empty): out of memory");
 	};
+	defer free(it_empty);
 	match (next(it_empty)) {
 	case done => void;
 	case ([]u8, *opaque) => abort("final: iterator produced elements after clear");

-- 
2.48.1