Start fix for issue #20.

Tessil · Jul 31, 2019 · 03979ec · 03979ec
1 parent 7d079e8
commit 03979ec
Show file tree

Hide file tree

Showing 2 changed files with 75 additions and 5 deletions.
diff --git a/include/tsl/robin_hash.h b/include/tsl/robin_hash.h
@@ -805,7 +805,12 @@ class robin_hash: private Hash, private KeyEqual, private GrowthPolicy {
      * we use an `iterator` instead of a `const_iterator`.
      */
     iterator erase(iterator pos) {
-        erase_from_bucket(pos);
+        const bool end_reached = erase_from_bucket(pos);
+        m_try_skrink_on_next_insert = true;
+
+        if(end_reached) {
+            return end();
+        }
 
         /**
          * Erase bucket used a backward shift after clearing the bucket.
@@ -815,8 +820,6 @@ class robin_hash: private Hash, private KeyEqual, private GrowthPolicy {
             ++pos;
         }
 
-        m_try_skrink_on_next_insert = true;
-
         return pos;
     }
 
@@ -1146,7 +1149,10 @@ class robin_hash: private Hash, private KeyEqual, private GrowthPolicy {
         return cend();
     }
 
-    void erase_from_bucket(iterator pos) {
+    /**
+     * Return true if we have have reached the end of the bucket.
+     */
+    bool erase_from_bucket(iterator pos) {
         pos.m_bucket->clear();
         m_nb_elements--;
 
@@ -1156,7 +1162,8 @@ class robin_hash: private Hash, private KeyEqual, private GrowthPolicy {
          * 
          * We try to move the values closer to their ideal bucket.
          */
-        std::size_t previous_ibucket = static_cast<std::size_t>(pos.m_bucket - m_buckets);
+        const std::size_t initial_ibucket = static_cast<std::size_t>(pos.m_bucket - m_buckets);
+        std::size_t previous_ibucket = initial_ibucket;
         std::size_t ibucket = next_bucket(previous_ibucket);
 
         while(m_buckets[ibucket].dist_from_ideal_bucket() > 0) {
@@ -1170,6 +1177,13 @@ class robin_hash: private Hash, private KeyEqual, private GrowthPolicy {
             previous_ibucket = ibucket;
             ibucket = next_bucket(ibucket);
         }
+
+        /**
+         * If we have wrapped around during the backward shift, we then have reached the end.
+         * TODO Improve documentation.
+         */
+        const bool end_reached = initial_ibucket > previous_ibucket;
+        return end_reached;
     }
 
     template<class K, class... Args>

diff --git a/tests/robin_map_tests.cpp b/tests/robin_map_tests.cpp
@@ -487,6 +487,62 @@ BOOST_AUTO_TEST_CASE(test_range_erase_same_iterators) {
     BOOST_CHECK_EQUAL(it_const.value(), -100);
 }
 
+BOOST_AUTO_TEST_CASE(test_erase_backward_shiftwrap_around_1) {
+    struct hash {
+        std::size_t operator()(std::uint32_t v) const {
+            return static_cast<std::size_t>(v);
+        }
+    };
+
+    tsl::robin_map<std::uint32_t, int> map(32);
+    BOOST_REQUIRE_EQUAL(map.bucket_count(), 32);
+
+    map.insert({31, 1});
+    map.insert({32+31, 2});
+
+    auto it = map.begin();
+    BOOST_CHECK_EQUAL(it->first, 32 + 31);
+
+    ++it;
+    BOOST_CHECK_EQUAL(it->first, 31);
+
+    it = map.erase(it);
+    BOOST_CHECK(it == map.end());
+    BOOST_CHECK_EQUAL(map.size(), 1);
+}
+
+BOOST_AUTO_TEST_CASE(test_erase_backward_shiftwrap_around_2) {
+    struct hash {
+        std::size_t operator()(std::uint32_t v) const {
+            return static_cast<std::size_t>(v);
+        }
+    };
+
+    tsl::robin_map<std::uint32_t, int> map(32);
+    BOOST_REQUIRE_EQUAL(map.bucket_count(), 32);
+
+    map.insert({31, 1});
+    map.insert({32+31, 2});
+    map.insert({64+31, 3});
+    map.insert({96+31, 3});
+
+    auto it = map.begin();
+    BOOST_CHECK_EQUAL(it->first, 32 + 31);
+
+    ++it;
+    BOOST_CHECK_EQUAL(it->first, 64 + 31);
+
+    ++it;
+    BOOST_CHECK_EQUAL(it->first, 96 + 31);
+
+    ++it;
+    BOOST_CHECK_EQUAL(it->first, 31);
+
+    it = map.erase(it);
+    BOOST_CHECK(it == map.end());
+    BOOST_CHECK_EQUAL(map.size(), 3);
+}
+
 /**
  * max_load_factor
  */