From 640cb926cad5cff6c4d877085a6978cae02ff8e0 Mon Sep 17 00:00:00 2001
From: lifegpc <root@lifegpc.com>
Date: Tue, 16 Jul 2024 17:12:11 +0800
Subject: [PATCH] Update

---
 binary_search_tree.h      | 38 ++++++++++++++++++++++++++++----------
 binary_tree.h             | 26 ++++++++++++++++++++++++++
 test/binary_tree_test.cpp |  8 ++++++++
 3 files changed, 62 insertions(+), 10 deletions(-)

diff --git a/binary_search_tree.h b/binary_search_tree.h
index 400d48f..97b392f 100644
--- a/binary_search_tree.h
+++ b/binary_search_tree.h
@@ -152,7 +152,7 @@ inline void binary_search_tree_iter(BinarySearchTree<K, V>* root, F callback, Ar
 }
 
 template <typename K, typename V>
-BinarySearchTree<K, V>* binary_search_tree_insert(BinarySearchTree<K, V>*& root, K key, V value, std::function<int(K, K)> comp, std::function<void(K)> free_key = std::function<void(K)>(), std::function<void(V)> free_value = std::function<void(V)>()) {
+BinarySearchTree<K, V>* binary_search_tree_insert(BinarySearchTree<K, V>*& root, K key, V value, std::function<int(K, K)> comp, std::function<void(K)> free_key = std::function<void(K)>(), std::function<void(V)> free_value = std::function<void(V)>(), bool noupdate = false) {
     if (!root) {
         root = binary_tree_new<BinarySearchTreePair<K, V>>({ key, value });
         if (!root) {
@@ -173,9 +173,9 @@ BinarySearchTree<K, V>* binary_search_tree_insert(BinarySearchTree<K, V>*& root,
         re = comp(key, cur->data.key);
     }
     if (re == 0) {
-        if (free_value) free_value(cur->data.value);
+        if (free_value && !noupdate) free_value(cur->data.value);
         if (free_key) free_key(key);
-        cur->data.value = value;
+        if (!noupdate) cur->data.value = value;
         return cur;
     }
     BinarySearchTree<K, V>* node = binary_tree_new<BinarySearchTreePair<K, V>>({key, value});
@@ -221,6 +221,18 @@ public:
     }
     template <class F, class G>
     BinarySearchMap(F free_key, G free_value = std::function<void(V)>(), Compare cmp = Compare()) : BinarySearchMap(std::function<void(K)>(free_key), std::function<void(V)>(free_value), cmp) {}
+    BinarySearchMap(const BinarySearchMap& other) {
+        tree = binary_tree_clone(other.tree, [&other](BinarySearchTreePair<K, V> e) {
+            if (other.free_key_func) other.free_key_func(e.key);
+            if (other.free_value_func) other.free_value_func(e.value);
+        });
+        if (!tree) {
+            throw std::runtime_error("Failed to clone map.");
+        }
+        comp_func = other.comp_func;
+        free_key_func = other.free_key_func;
+        free_value_func = other.free_value_func;
+    }
     ~BinarySearchMap() {
         clear();
     }
@@ -230,6 +242,9 @@ public:
     bool inline del(K key, V* deleted_value = nullptr) {
         return binary_search_tree_delete(tree, key, deleted_value, comp_func, free_key_func, free_value_func);
     }
+    bool inline get(K key, V& value) const {
+        return binary_search_tree_get(tree, key, value, comp_func);
+    }
     bool inline get(K key, V& value) {
         return binary_search_tree_get(tree, key, value, comp_func);
     }
@@ -244,15 +259,18 @@ public:
     void inline iter(F callback, Args... args) {
         return binary_search_tree_iter(tree, std::function<void(K, V, Args...)>(callback), args...);
     }
-    V& operator[](K key) {
+    const V& operator[](K key) const {
         BinarySearchTree<K, V>* node = binary_search_tree_get_node(tree, key, comp_func);
         if (!node) {
-            V tmp;
-            BinarySearchTree<K, V>* node = binary_search_tree_insert(tree, key, tmp, comp_func, free_key_func, free_value_func);
-            if (!node) {
-                throw std::runtime_error("Failed to insert new node");
-            }
-            return node->data.value;
+            throw std::runtime_error("No such node.");
+        }
+        return node->data.value;
+    }
+    V& operator[](K key) {
+        V tmp = V();
+        BinarySearchTree<K, V>* node = binary_search_tree_insert(tree, key, tmp, comp_func, free_key_func, free_value_func, true);
+        if (!node) {
+            throw std::runtime_error("Failed to insert new node.");
         }
         return node->data.value;
     }
diff --git a/binary_tree.h b/binary_tree.h
index 2a3d410..6bf5b84 100644
--- a/binary_tree.h
+++ b/binary_tree.h
@@ -32,6 +32,32 @@ inline void binary_tree_clear(struct BinaryTree<T>*& top, F free_func) {
     binary_tree_clear(top, std::function<void(T)>(free_func));
 }
 
+template <typename T>
+struct BinaryTree<T>* binary_tree_clone(struct BinaryTree<T>* tree, std::function<void(T)> free_func = std::function<void(T)>()) {
+    if (!tree) return nullptr;
+    T newdata(tree->data);
+    struct BinaryTree<T>* now = binary_tree_new(newdata);
+    if (now) {
+        if (tree->left) {
+            now->left = binary_tree_clone(tree->left, free_func);
+            if (!now->left) goto end;
+        }
+        if (tree->right) {
+            now->right = binary_tree_clone(tree->right, free_func);
+            if (!now->right) goto end;
+        }
+    }
+    return now;
+end:
+    binary_tree_clear(now, free_func);
+    return now;
+}
+
+template <typename T, class F>
+inline struct BinaryTree<T>* binary_tree_clone(struct BinaryTree<T>* tree, F free_func) {
+    return binary_tree_clone(tree, std::function<void(T)>(free_func));
+}
+
 template <typename T, typename... Args>
 void binary_tree_dfs(struct BinaryTree<T>* top, std::function<void(struct BinaryTree<T>*, Args...)> callback, Args... args) {
     if (!top) return;
diff --git a/test/binary_tree_test.cpp b/test/binary_tree_test.cpp
index 46f0efc..90b33eb 100644
--- a/test/binary_tree_test.cpp
+++ b/test/binary_tree_test.cpp
@@ -160,4 +160,12 @@ TEST(BinaryTreeTest, BinarySearchTree3) {
     });
     GTEST_ASSERT_EQ(keys, "10,30,33");
     GTEST_ASSERT_EQ(values, "20,40,44");
+    GTEST_ASSERT_EQ(map[20], 0);
+    const BinarySearchMap map2(map);
+    map[20] = 33;
+    GTEST_ASSERT_EQ(map2[20], 0);
+    GTEST_ASSERT_EQ(map[20], 33);
+    int j;
+    GTEST_ASSERT_EQ(map2.get(33, j), true);
+    GTEST_ASSERT_EQ(j, 44);
 }