From d24c47fdcef709048a9e91a8fd54eb067b46302f Mon Sep 17 00:00:00 2001
From: greg7mdp <greg7mdp@gmail.com>
Date: Mon, 25 Jul 2022 21:37:43 -0400
Subject: [PATCH] Add `with_submap` api for parallel maps allowing access to
 submaps under lock protection

---
 parallel_hashmap/phmap.h        | 45 +++++++++++++++++++++++----------
 tests/parallel_hash_map_test.cc | 11 ++++++++
 tests/raw_hash_set_test.cc      | 26 +++++++++----------
 3 files changed, 56 insertions(+), 26 deletions(-)

diff --git a/parallel_hashmap/phmap.h b/parallel_hashmap/phmap.h
index 86032cb..3bc8706 100644
--- a/parallel_hashmap/phmap.h
+++ b/parallel_hashmap/phmap.h
@@ -3318,19 +3318,6 @@ public:
         }
     }
 
-#if __cplusplus >= 201703L
-    template <class ExecutionPolicy, class F>
-    void for_each(ExecutionPolicy&& policy, F&& fCallback) const {
-        std::for_each(
-            std::forward<ExecutionPolicy>(policy), sets_.begin(), sets_.end(),
-            [&](auto const& inner) {
-                typename Lockable::SharedLock m(const_cast<Inner&>(inner));
-                std::for_each(inner.set_.begin(), inner.set_.end(), fCallback);
-            }
-        );
-    }
-#endif
-
     // this version allows to modify the values
     template <class F>
     void for_each_m(F&& fCallback) {
@@ -3341,6 +3328,17 @@ public:
     }
 
 #if __cplusplus >= 201703L
+    template <class ExecutionPolicy, class F>
+    void for_each(ExecutionPolicy&& policy, F&& fCallback) const {
+        std::for_each(
+            std::forward<ExecutionPolicy>(policy), sets_.begin(), sets_.end(),
+            [&](auto const& inner) {
+                typename Lockable::SharedLock m(const_cast<Inner&>(inner));
+                std::for_each(inner.set_.begin(), inner.set_.end(), fCallback);
+            }
+        );
+    }
+
     template <class ExecutionPolicy, class F>
     void for_each_m(ExecutionPolicy&& policy, F&& fCallback) {
         std::for_each(
@@ -3353,6 +3351,27 @@ public:
     }
 #endif
 
+    // Extension API: access internal submaps by index
+    // under lock protection
+    // ex: m.with_submap(i, [&](const Map::EmbeddedSet& set) {
+    //        for (auto& p : set) { ...; }});
+    // -------------------------------------------------
+    template <class F>
+    void with_submap(size_t idx, F&& fCallback) const {
+        const Inner& inner     = sets_[idx];
+        const auto&  set = inner.set_;
+        typename Lockable::SharedLock m(const_cast<Inner&>(inner));
+        fCallback(set);
+    }
+
+    template <class F>
+    void with_submap_m(size_t idx, F&& fCallback) const {
+        Inner& inner   = sets_[idx];
+        auto&  set     = inner.set_;
+        typename Lockable::UniqueLock m(const_cast<Inner&>(inner));
+        fCallback(set);
+    }
+
     // Extension API: support for heterogeneous keys.
     //
     //   std::unordered_set<std::string> s;
diff --git a/tests/parallel_hash_map_test.cc b/tests/parallel_hash_map_test.cc
index a7ef1ee..4812e06 100644
--- a/tests/parallel_hash_map_test.cc
+++ b/tests/parallel_hash_map_test.cc
@@ -116,6 +116,17 @@ TEST(THIS_TEST_NAME, ForEach) {
             EXPECT_EQ(pair.first + 7, pair.second);
         });
     EXPECT_EQ(counter, 3);
+    
+    counter = 0;
+    for (size_t i=0; i<m.subcnt(); ++i) {
+        m.with_submap(i, [&](const Map::EmbeddedSet& set) {
+            for (auto& p : set) {
+                ++counter;
+                EXPECT_EQ(p.first + 7, p.second);
+            }
+        });
+    }
+    EXPECT_EQ(counter, 3);
 }
 
 TEST(THIS_TEST_NAME, EmplaceSingle) {
diff --git a/tests/raw_hash_set_test.cc b/tests/raw_hash_set_test.cc
index 0e46959..feb9709 100644
--- a/tests/raw_hash_set_test.cc
+++ b/tests/raw_hash_set_test.cc
@@ -138,27 +138,27 @@ TEST(BitMask, WithShift) {
   EXPECT_EQ(0x0000000080800000, mask);
 
   BitMask<uint64_t, 8, 3> b(mask);
-  EXPECT_EQ(*b, 2);
+  EXPECT_EQ(*b, 2u);
 }
 
 TEST(BitMask, LeadingTrailing) {
-  EXPECT_EQ((BitMask<uint32_t, 16>(0x00001a40).LeadingZeros()), 3);
-  EXPECT_EQ((BitMask<uint32_t, 16>(0x00001a40).TrailingZeros()), 6);
+  EXPECT_EQ((BitMask<uint32_t, 16>(0x00001a40).LeadingZeros()), 3u);
+  EXPECT_EQ((BitMask<uint32_t, 16>(0x00001a40).TrailingZeros()), 6u);
 
-  EXPECT_EQ((BitMask<uint32_t, 16>(0x00000001).LeadingZeros()), 15);
-  EXPECT_EQ((BitMask<uint32_t, 16>(0x00000001).TrailingZeros()), 0);
+  EXPECT_EQ((BitMask<uint32_t, 16>(0x00000001).LeadingZeros()), 15u);
+  EXPECT_EQ((BitMask<uint32_t, 16>(0x00000001).TrailingZeros()), 0u);
 
-  EXPECT_EQ((BitMask<uint32_t, 16>(0x00008000).LeadingZeros()), 0);
-  EXPECT_EQ((BitMask<uint32_t, 16>(0x00008000).TrailingZeros()), 15);
+  EXPECT_EQ((BitMask<uint32_t, 16>(0x00008000).LeadingZeros()), 0u);
+  EXPECT_EQ((BitMask<uint32_t, 16>(0x00008000).TrailingZeros()), 15u);
 
-  EXPECT_EQ((BitMask<uint64_t, 8, 3>(0x0000008080808000).LeadingZeros()), 3);
-  EXPECT_EQ((BitMask<uint64_t, 8, 3>(0x0000008080808000).TrailingZeros()), 1);
+  EXPECT_EQ((BitMask<uint64_t, 8, 3>(0x0000008080808000).LeadingZeros()), 3u);
+  EXPECT_EQ((BitMask<uint64_t, 8, 3>(0x0000008080808000).TrailingZeros()), 1u);
 
-  EXPECT_EQ((BitMask<uint64_t, 8, 3>(0x0000000000000080).LeadingZeros()), 7);
-  EXPECT_EQ((BitMask<uint64_t, 8, 3>(0x0000000000000080).TrailingZeros()), 0);
+  EXPECT_EQ((BitMask<uint64_t, 8, 3>(0x0000000000000080).LeadingZeros()), 7u);
+  EXPECT_EQ((BitMask<uint64_t, 8, 3>(0x0000000000000080).TrailingZeros()), 0u);
 
-  EXPECT_EQ((BitMask<uint64_t, 8, 3>(0x8000000000000000).LeadingZeros()), 0);
-  EXPECT_EQ((BitMask<uint64_t, 8, 3>(0x8000000000000000).TrailingZeros()), 7);
+  EXPECT_EQ((BitMask<uint64_t, 8, 3>(0x8000000000000000).LeadingZeros()), 0u);
+  EXPECT_EQ((BitMask<uint64_t, 8, 3>(0x8000000000000000).TrailingZeros()), 7u);
 }
 
 TEST(Group, EmptyGroup) {