From 03979ec4b5a873bd2d41907d1315a5e55c1cf8bb Mon Sep 17 00:00:00 2001
From: Tessil <tessil@gmx.com>
Date: Wed, 31 Jul 2019 22:51:01 +0200
Subject: [PATCH] Start fix for issue #20.

---
 include/tsl/robin_hash.h  | 24 +++++++++++++----
 tests/robin_map_tests.cpp | 56 +++++++++++++++++++++++++++++++++++++++
 2 files changed, 75 insertions(+), 5 deletions(-)

diff --git a/include/tsl/robin_hash.h b/include/tsl/robin_hash.h
index e944363..113c308 100644
--- 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
index b1ae6fc..20d8263 100644
--- 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
  */