Re-implement RBSet as a subclass of RBMap with empty values.

This deduplicates RBTree logic to make the code more maintainable.
Ivorforce · Jan 3, 2025 · 5d00dc4 · 5d00dc4
1 parent 2582793
commit 5d00dc4
Show file tree

Hide file tree

Showing 6 changed files with 65 additions and 681 deletions.
diff --git a/core/templates/rb_map.h b/core/templates/rb_map.h
@@ -40,7 +40,8 @@
 // based on the very nice implementation of rb-trees by:
 // https://web.archive.org/web/20120507164830/https://web.mit.edu/~emin/www/source_code/red_black_tree/index.html
 
-template <typename K, typename V, typename C = Comparator<K>, typename A = DefaultAllocator>
+// value_is_key is a hack to make the RBSet easier to implement.
+template <typename K, typename V, typename C = Comparator<K>, typename A = DefaultAllocator, bool value_is_key = false>
 class RBMap {
 	enum Color {
 		RED,
@@ -50,8 +51,7 @@ class RBMap {
 
 public:
 	class Element {
-	private:
-		friend class RBMap<K, V, C, A>;
+		friend class RBMap<K, V, C, A, value_is_key>;
 		int color = RED;
 		Element *right = nullptr;
 		Element *left = nullptr;
@@ -85,11 +85,19 @@ class RBMap {
 		const V &value() const {
 			return _data.value;
 		}
-		V &get() {
-			return _data.value;
+		auto &get() {
+			if constexpr (value_is_key) {
+				return (const K &)_data.key;
+			} else {
+				return _data.value;
+			}
 		}
-		const V &get() const {
-			return _data.value;
+		const auto &get() const {
+			if constexpr (value_is_key) {
+				return _data.key;
+			} else {
+				return _data.value;
+			}
 		}
 		Element(const KeyValue<K, V> &p_data) :
 				_data(p_data) {}
@@ -132,10 +140,20 @@ class RBMap {
 	};
 
 	struct ConstIterator {
-		_FORCE_INLINE_ const KeyValue<K, V> &operator*() const {
-			return E->key_value();
+		_FORCE_INLINE_ const auto &operator*() const {
+			if constexpr (value_is_key) {
+				return E->key();
+			} else {
+				return E->key_value();
+			}
+		}
+		_FORCE_INLINE_ const auto *operator->() const {
+			if constexpr (value_is_key) {
+				return &E->key();
+			} else {
+				return &E->key_value();
+			}
 		}
-		_FORCE_INLINE_ const KeyValue<K, V> *operator->() const { return &E->key_value(); }
 		_FORCE_INLINE_ ConstIterator &operator++() {
 			E = E->next();
 			return *this;