From 1f9282ff0fe14e50b26d5945f39f8f2c51b43f04 Mon Sep 17 00:00:00 2001
From: Runxi Yu <me@runxiyu.org>
Date: Sun, 21 Sep 2025 19:37:02 +0800
Subject: [PATCH] Maps should own their keys

---
 ds/map/btree/finish.ha        |  3 +++
 ds/map/btree/internal.ha      | 18 +++++++++++++++++-
 ds/map/rbtree/del.ha          |  1 +
 ds/map/rbtree/finish.ha       |  1 +
 ds/map/rbtree/set.ha          | 16 +++++++++++++++-
 ds/map/slice_basic/del.ha     |  1 +
 ds/map/slice_basic/finish.ha  |  3 +++
 ds/map/slice_basic/set.ha     | 13 ++++++++++++-
 ds/map/slice_sorted/del.ha    |  1 +
 ds/map/slice_sorted/finish.ha |  5 +++++
 ds/map/slice_sorted/set.ha    | 13 ++++++++++++-
 ds/map/swiss/del.ha           |  2 +-
 ds/map/swiss/finish.ha        |  8 ++++++++
 ds/map/swiss/set.ha           | 15 +++++++++++++--

diff --git a/ds/map/btree/finish.ha b/ds/map/btree/finish.ha
index e8a6f231a1e25d99ad642bbd01848ca438a454f5..bb826e4cd1ffae6121f198b2e3083bbb0223615a 100644
--- a/ds/map/btree/finish.ha
+++ b/ds/map/btree/finish.ha
@@ -13,6 +13,9 @@ 			node_finish(n.children[i]);
 		};
 		free(n.children);
 	};
+	for (let i = 0z; i < len(n.keys); i += 1) {
+		free(n.keys[i]);
+	};
 	free(n.keys);
 	free(n.vals);
 	free(n);
diff --git a/ds/map/btree/internal.ha b/ds/map/btree/internal.ha
index d73038ca42922c5f9957db12c1c1d895742126ef..022903ed1e5ec7326d004c31861432399375ae85 100644
--- a/ds/map/btree/internal.ha
+++ b/ds/map/btree/internal.ha
@@ -71,6 +71,13 @@ 	insert(x.vals[i], medv)?;
 	insert(x.children[i + 1], z)?;
 };
 
+fn dup_u8(src: []u8) ([]u8 | nomem) = {
+	return match (alloc(src, len(src))) {
+	case let b: []u8 => yield b;
+	case nomem => return nomem;
+	};
+};
+
 fn insert_nonfull(m: *map, x: *node, key: []u8, val: *opaque) (void | nomem) = {
 	let i = sort::lbisect((x.keys: []const opaque), size([]u8),
 		(&key: const *opaque), &cmp_u8slice);
@@ -81,7 +88,11 @@ 		return;
 	};
 
 	if (x.leaf) {
-		insert(x.keys[i], key)?;
+		let kcopy = match (dup_u8(key)) {
+		case let b: []u8 => yield b;
+		case nomem => return nomem;
+		};
+		insert(x.keys[i], kcopy)?;
 		insert(x.vals[i], val)?;
 		return;
 	};
@@ -205,6 +216,7 @@
 	if (i < len(x.keys) && bytes::equal(x.keys[i], key)) {
 		if (x.leaf) {
 			let ret = x.vals[i];
+			free(x.keys[i]);
 			delete(x.keys[i]);
 			delete(x.vals[i]);
 			return ret;
@@ -217,14 +229,18 @@
 		if (len(y.keys) >= t) {
 			let (pk, pv) = pop_max(m, y);
 			let ret = x.vals[i];
+			let oldk = x.keys[i];
 			x.keys[i] = pk;
 			x.vals[i] = pv;
+			free(oldk);
 			return ret;
 		} else if (len(z.keys) >= t) {
 			let (sk, sv) = pop_min(m, z);
 			let ret = x.vals[i];
+			let oldk = x.keys[i];
 			x.keys[i] = sk;
 			x.vals[i] = sv;
+			free(oldk);
 			return ret;
 		} else {
 			merge_children(m, x, i);
diff --git a/ds/map/rbtree/del.ha b/ds/map/rbtree/del.ha
index be529d63c1b8fd69eff8fcf519471e879b3fc8a9..e18e248a768d79d25cced6c01c0164d76e7bb51a 100644
--- a/ds/map/rbtree/del.ha
+++ b/ds/map/rbtree/del.ha
@@ -47,6 +47,7 @@ 			set_parent(y.left, y);
 			y.color = nodez.color;
 		};
 
+		free(nodez.key);
 		free(nodez);
 
 		if (y_orig == color::BLACK) {
diff --git a/ds/map/rbtree/finish.ha b/ds/map/rbtree/finish.ha
index 6bc923e6de2ec8211bad266e4675c621aa565f6c..78981e774a2e718487226803db6232c3b42154e9 100644
--- a/ds/map/rbtree/finish.ha
+++ b/ds/map/rbtree/finish.ha
@@ -4,6 +4,7 @@ 	case null => return;
 	case let p: *node =>
 		free_subtree(p.left);
 		free_subtree(p.right);
+		free(p.key);
 		free(p);
 	};
 };
diff --git a/ds/map/rbtree/set.ha b/ds/map/rbtree/set.ha
index 247514e06bc2211728ea2b05581dedb70183c97d..07e8ab6226bbbf18460a3dc033e2e9a0680f9705 100644
--- a/ds/map/rbtree/set.ha
+++ b/ds/map/rbtree/set.ha
@@ -1,5 +1,14 @@
+// SPDX-License-Identifier: MPL-2.0
+
 use bytes;
 
+fn dup_u8(src: []u8) ([]u8 | nomem) = {
+	return match (alloc(src, len(src))) {
+	case let b: []u8 => yield b;
+	case nomem => return nomem;
+	};
+};
+
 export fn set(m: *map, key: []u8, value: *opaque) (void | nomem) = {
 	match (find_node(m, key)) {
 	case let ex: *node =>
@@ -8,9 +17,14 @@ 		return;
 	case null => void;
 	};
 
+	let kcopy = match (dup_u8(key)) {
+	case let b: []u8 => yield b;
+	case nomem => return nomem;
+	};
+
 	let z = alloc(node {
 		color = color::RED,
-		key = key,
+		key = kcopy,
 		val = value,
 		left = null,
 		right = null,
diff --git a/ds/map/slice_basic/del.ha b/ds/map/slice_basic/del.ha
index b10e179bc055eb5c9acc6492c96ac4f536920ebc..a56d49566f495cc51ba6f3a9dc0485ccfa60e521 100644
--- a/ds/map/slice_basic/del.ha
+++ b/ds/map/slice_basic/del.ha
@@ -7,6 +7,7 @@ export fn del(m: *map, key: []u8) (*opaque | void) = {
 	for (let i = 0z; i < len(m.items); i += 1) {
 		if (bytes::equal(m.items[i].0, key)) {
 			let v = m.items[i].1;
+			free(m.items[i].0);
 			delete(m.items[i]);
 			return v;
 		};
diff --git a/ds/map/slice_basic/finish.ha b/ds/map/slice_basic/finish.ha
index d07898c0e7499d8080b5b1d5c93d08da417c0e42..464ebe6c5ad519a61615773f2c1a865508129336 100644
--- a/ds/map/slice_basic/finish.ha
+++ b/ds/map/slice_basic/finish.ha
@@ -2,6 +2,9 @@ // SPDX-License-Identifier: MPL-2.0
 
 // Frees resources associated with a [[map]].
 export fn finish(m: *map) void = {
+	for (let i = 0z; i < len(m.items); i += 1) {
+		free(m.items[i].0);
+	};
 	free(m.items);
 	free(m);
 };
diff --git a/ds/map/slice_basic/set.ha b/ds/map/slice_basic/set.ha
index 6915cb01feb73fe4e27747a4a287a143abd7bf75..1f8d493485897c5daf44a26de12e8f8133a8ce03 100644
--- a/ds/map/slice_basic/set.ha
+++ b/ds/map/slice_basic/set.ha
@@ -2,6 +2,13 @@ // SPDX-License-Identifier: MPL-2.0
 
 use bytes;
 
+fn dup_u8(src: []u8) ([]u8 | nomem) = {
+	return match (alloc(src, len(src))) {
+	case let b: []u8 => yield b;
+	case nomem => return nomem;
+	};
+};
+
 // Sets an item in a [[map]], replacing any existing item with the same key.
 export fn set(m: *map, key: []u8, value: *opaque) (void | nomem) = {
 	for (let i = 0z; i < len(m.items); i += 1) {
@@ -10,5 +17,9 @@ 			m.items[i].1 = value;
 			return;
 		};
 	};
-	append(m.items, (key, value))?;
+	let kb = match (dup_u8(key)) {
+	case let b: []u8 => yield b;
+	case nomem => return nomem;
+	};
+	append(m.items, (kb, value))?;
 };
diff --git a/ds/map/slice_sorted/del.ha b/ds/map/slice_sorted/del.ha
index 3dd5078a79a331b1dc73ed14582ae0e83525de73..0044358dc25f18cfc2ad2d23dc2ffc7c2ae0d6c3 100644
--- a/ds/map/slice_sorted/del.ha
+++ b/ds/map/slice_sorted/del.ha
@@ -9,6 +9,7 @@ 		size(([]u8, *opaque)),
 		(&probe: const *opaque), &cmp_kv)) {
 	case let idx: size =>
 		let v = m.items[idx].1;
+		free(m.items[idx].0);
 		delete(m.items[idx]);
 		return v;
 	case void =>
diff --git a/ds/map/slice_sorted/finish.ha b/ds/map/slice_sorted/finish.ha
index 88e8eb85f2404ed13519653348e4bf92ba86e217..464ebe6c5ad519a61615773f2c1a865508129336 100644
--- a/ds/map/slice_sorted/finish.ha
+++ b/ds/map/slice_sorted/finish.ha
@@ -1,5 +1,10 @@
+// SPDX-License-Identifier: MPL-2.0
+
 // Frees resources associated with a [[map]].
 export fn finish(m: *map) void = {
+	for (let i = 0z; i < len(m.items); i += 1) {
+		free(m.items[i].0);
+	};
 	free(m.items);
 	free(m);
 };
diff --git a/ds/map/slice_sorted/set.ha b/ds/map/slice_sorted/set.ha
index 34a74059105c4d77d4c7809ac84ab4d265090566..7eda109fad7c5728fa763e6a1193ddad33bc205c 100644
--- a/ds/map/slice_sorted/set.ha
+++ b/ds/map/slice_sorted/set.ha
@@ -1,5 +1,12 @@
 use sort;
 
+fn dup_u8(src: []u8) ([]u8 | nomem) = {
+	return match (alloc(src, len(src))) {
+	case let b: []u8 => yield b;
+	case nomem => return nomem;
+	};
+};
+
 // Sets an item in a [[map]], replacing any existing item with the same key.
 export fn set(m: *map, key: []u8, value: *opaque) (void | nomem) = {
 	let dummy = 0;
@@ -14,6 +21,10 @@ 	case void =>
 		let ins = sort::lbisect((m.items: []const opaque),
 			size(([]u8, *opaque)),
 			(&probe: const *opaque), &cmp_kv);
-		insert(m.items[ins], (key, value))?;
+		let kb = match (dup_u8(key)) {
+		case let b: []u8 => yield b;
+		case nomem => return nomem;
+		};
+		insert(m.items[ins], (kb, value))?;
 	};
 };
diff --git a/ds/map/swiss/del.ha b/ds/map/swiss/del.ha
index 653066fa12edcd0f9b256a2063744e4bfb4a5ca7..b9b73ef6b813e0cee2136f7059e433bdecf0a15c 100644
--- a/ds/map/swiss/del.ha
+++ b/ds/map/swiss/del.ha
@@ -21,11 +21,11 @@ 			if (is_full_ctrl(c) && c == t) {
 				if (bytes::equal(g.keys[i], key)) {
 					let v = g.vals[i];
 					g.ctrl[i] = CTRL_DELETED;
+					free(g.keys[i]);
 					g.keys[i] = [];
 					g.vals[i] = null;
 					m.used -= 1;
 					m.tombs += 1;
-					// elide the tombstones if exceed 1/3 of the capacity
 					if (m.tombs * 3 >= capacity_slots(m)) {
 						rehash_in_place(m);
 					};
diff --git a/ds/map/swiss/finish.ha b/ds/map/swiss/finish.ha
index 640dd673472c4a9817496e481eb97f2eded8d80f..38b223df90e47a24e2df300b655715c06749d025 100644
--- a/ds/map/swiss/finish.ha
+++ b/ds/map/swiss/finish.ha
@@ -5,6 +5,14 @@
 // Frees resources associated with a [[map]].
 export fn finish(m: *map) void = {
 	if (len(m.groups) != 0) {
+		for (let gi = 0z; gi <= m.group_mask; gi += 1) {
+			let g = &m.groups[gi];
+			for (let si = 0z; si < GROUP_SIZE; si += 1) {
+				let c = g.ctrl[si];
+				if (!is_full_ctrl(c)) continue;
+				free(g.keys[si]);
+			};
+		};
 		free(m.groups);
 	};
 	free(m);
diff --git a/ds/map/swiss/set.ha b/ds/map/swiss/set.ha
index c8cac47251f5b01049a1a0584574f71fdd1f542b..26aa0e57197878ff50704d3483643fdfd5d9426b 100644
--- a/ds/map/swiss/set.ha
+++ b/ds/map/swiss/set.ha
@@ -1,8 +1,15 @@
-// SPDX-License-Identifier: Apache-2.0 AND MPL-2.0
 // SPDX-FileCopyrightText: 2024 The Cockroach Authors
 // SPDX-FileCopyrightText: 2025 Runxi Yu
 
+use ds::map;
 use bytes;
+
+fn dup_u8(src: []u8) ([]u8 | nomem) = {
+	return match (alloc(src, len(src))) {
+	case let b: []u8 => yield b;
+	case nomem => return nomem;
+	};
+};
 
 // Sets an item in a [[map]], replacing any existing item with the same key.
 export fn set(m: *map, key: []u8, value: *opaque) (void | nomem) = {
@@ -78,7 +85,11 @@ 				let slot = match (first_dead) {
 				case void => yield i;
 				case let di: size => yield di;
 				};
-				g.keys[slot] = key;
+				let kb = match (dup_u8(key)) {
+				case let b: []u8 => yield b;
+				case nomem => return nomem;
+				};
+				g.keys[slot] = kb;
 				g.vals[slot] = value;
 				g.ctrl[slot] = t;
 				m.used += 1;

-- 
2.48.1