Fix FlatMap assign operator bug (#2622)

apache · Apr 29, 2024 · f97cbe5 · f97cbe5
1 parent 0128bb1
commit f97cbe5
Show file tree

Hide file tree

Showing 3 changed files with 42 additions and 17 deletions.
diff --git a/src/butil/containers/flat_map.h b/src/butil/containers/flat_map.h
@@ -151,12 +151,12 @@ class FlatMap {
         key_type key;
     };
 
-    FlatMap(const hasher& hashfn = hasher(), 
-            const key_equal& eql = key_equal(),
-            const allocator_type& alloc = allocator_type());
+    explicit FlatMap(const hasher& hashfn = hasher(),
+                     const key_equal& eql = key_equal(),
+                     const allocator_type& alloc = allocator_type());
     ~FlatMap();
-    FlatMap(const FlatMap& rhs);    
-    void operator=(const FlatMap& rhs);
+    FlatMap(const FlatMap& rhs);
+    FlatMap& operator=(const FlatMap& rhs);
     void swap(FlatMap & rhs);
 
     // Must be called to initialize this map, otherwise insert/operator[]
@@ -307,9 +307,9 @@ class FlatSet {
     typedef typename Map::key_equal key_equal;
     typedef typename Map::allocator_type allocator_type;
 
-    FlatSet(const hasher& hashfn = hasher(), 
-            const key_equal& eql = key_equal(),
-            const allocator_type& alloc = allocator_type())
+    explicit FlatSet(const hasher& hashfn = hasher(),
+                     const key_equal& eql = key_equal(),
+                     const allocator_type& alloc = allocator_type())
         : _map(hashfn, eql, alloc) {}
     void swap(FlatSet & rhs) { _map.swap(rhs._map); }
 

diff --git a/src/butil/containers/flat_map_inl.h b/src/butil/containers/flat_map_inl.h
@@ -261,39 +261,48 @@ FlatMap<_K, _T, _H, _E, _S, _A>::FlatMap(const FlatMap& rhs)
 }
 
 template <typename _K, typename _T, typename _H, typename _E, bool _S, typename _A>
-void
+FlatMap<_K, _T, _H, _E, _S, _A>&
 FlatMap<_K, _T, _H, _E, _S, _A>::operator=(const FlatMap<_K, _T, _H, _E, _S, _A>& rhs) {
     if (this == &rhs) {
-        return;
+        return *this;
     }
     // NOTE: assignment does not change _load_factor/_hashfn/_eql if |this| is
     // initialized
     clear();
-    if (rhs.empty()) {
-        return;
-    }
-    if (!initialized()) {
-        _load_factor = rhs._load_factor;
+    if (!rhs.initialized()) {
+        return *this;
     }
+    bool need_copy = !rhs.empty();
+    _load_factor = rhs._load_factor;
     if (_buckets == NULL || is_too_crowded(rhs._size)) {
         get_allocator().Free(_buckets);
         _nbucket = rhs._nbucket;
         // note: need an extra bucket to let iterator know where buckets end
         _buckets = (Bucket*)get_allocator().Alloc(sizeof(Bucket) * (_nbucket + 1/*note*/));
         if (NULL == _buckets) {
             LOG(ERROR) << "Fail to new _buckets";
-            return;
+            return *this;
+        }
+        // If no need to copy, set buckets invalid.
+        if (!need_copy) {
+            for (size_t i = 0; i < _nbucket; ++i) {
+                _buckets[i].set_invalid();
+            }
+            _buckets[_nbucket].next = NULL;
         }
         if (_S) {
             free(_thumbnail);
             _thumbnail = bit_array_malloc(_nbucket);
             if (NULL == _thumbnail) {
                 LOG(ERROR) << "Fail to new _thumbnail";
-                return;
+                return *this;
             }
             bit_array_clear(_thumbnail, _nbucket);
         }
     }
+    if (!need_copy) {
+        return *this;
+    }
     if (_nbucket == rhs._nbucket) {
         // For equivalent _nbucket, walking through _buckets instead of using
         // iterators is more efficient.
@@ -322,6 +331,7 @@ FlatMap<_K, _T, _H, _E, _S, _A>::operator=(const FlatMap<_K, _T, _H, _E, _S, _A>
                 Element::second_ref_from_value(*it);
         }
     }
+    return *this;
 }
 
 template <typename _K, typename _T, typename _H, typename _E, bool _S, typename _A>

diff --git a/test/flat_map_unittest.cpp b/test/flat_map_unittest.cpp
@@ -1332,6 +1332,21 @@ TEST_F(FlatMapTest, copy) {
     m2 = m1;
     ASSERT_FALSE(m1.is_too_crowded(m1.size()));
     ASSERT_FALSE(m2.is_too_crowded(m1.size()));
+
+    butil::FlatMap<int, int> m3;
+    ASSERT_FALSE(m3.initialized());
+    m1 = m3;
+    ASSERT_TRUE(m1.empty());
+    ASSERT_TRUE(m1.initialized());
+
+    m3 = m2;
+    ASSERT_TRUE(m3.initialized());
+    m3.clear();
+    ASSERT_TRUE(m3.initialized());
+    ASSERT_TRUE(m3.empty());
+    butil::FlatMap<int, int> m4 = m3;
+    ASSERT_TRUE(m4.initialized());
+    ASSERT_TRUE(m4.empty());
 }
 
 }