From bd415ae9353779aab3f78e62d831af23adfc7370 Mon Sep 17 00:00:00 2001
From: Rui Ueyama <ruiu@cs.stanford.edu>
Date: Wed, 23 Dec 2020 17:15:33 +0900
Subject: [PATCH] Use concurrent_unordered_set

---
 input_sections.cc |  3 ++-
 main.cc           |  2 +-
 mold.h            | 54 +++++++++++++++++++++++++----------------------
 object_file.cc    |  6 ++++--
 output_chunks.cc  |  5 +++--
 5 files changed, 39 insertions(+), 31 deletions(-)

diff --git a/input_sections.cc b/input_sections.cc
index 14f5ad7b..109239a4 100644
--- a/input_sections.cc
+++ b/input_sections.cc
@@ -333,7 +333,8 @@ MergeableSection::MergeableSection(InputSection *isec, std::string_view data)
     std::string_view substr = data.substr(0, end + 1);
     data = data.substr(end + 1);
 
-    StringPiece *piece = parent.map.insert(substr, StringPiece(substr));
+    auto [it, inserted] = parent.set.insert(StringPiece(substr));
+    StringPiece *piece = &*it;
     pieces.push_back({piece, offset});
     offset += substr.size();
   }
diff --git a/main.cc b/main.cc
index 93ee73dc..ee375d52 100644
--- a/main.cc
+++ b/main.cc
@@ -1057,7 +1057,7 @@ static void show_stats() {
 
   static Counter merged_strings("merged_strings");
   for (MergedSection *osec : MergedSection::instances)
-    merged_strings.inc(osec->map.size());
+    merged_strings.inc(osec->set.size());
 
   Counter num_output_chunks("output_out::chunks", out::chunks.size());
   Counter num_objs("num_objs", out::objs.size());
diff --git a/mold.h b/mold.h
index d5017283..a738e3fb 100644
--- a/mold.h
+++ b/mold.h
@@ -7,6 +7,7 @@
 #include "elf.h"
 
 #include "tbb/concurrent_hash_map.h"
+#include "tbb/concurrent_unordered_set.h"
 #include "tbb/parallel_for_each.h"
 #include "tbb/parallel_invoke.h"
 #include "tbb/parallel_reduce.h"
@@ -110,28 +111,6 @@ struct tbb_hash_compare<std::string_view> {
 };
 }
 
-template<typename ValueT>
-class ConcurrentMap {
-public:
-  typedef tbb::concurrent_hash_map<std::string_view, ValueT> MapT;
-
-  ValueT *insert(std::string_view key, const ValueT &val) {
-    typename MapT::const_accessor acc;
-    map.insert(acc, std::make_pair(key, val));
-    return const_cast<ValueT *>(&acc->second);
-  }
-
-  void for_each_value(std::function<void(const ValueT &)> fn) {
-    for (typename MapT::const_iterator it = map.begin(); it != map.end(); ++it)
-      fn(it->second);
-  }
-
-  size_t size() const { return map.size(); }
-
-private:
-  MapT map;
-};
-
 //
 // Symbol
 //
@@ -144,6 +123,18 @@ struct StringPiece {
     : isec(other.isec.load()), data(other.data), size(other.size),
       output_offset(other.output_offset) {}
 
+  struct hash {
+    size_t operator()(const StringPiece &piece) const {
+      return std::hash<std::string_view>()({(char *)piece.data, piece.size});
+    }
+  };
+
+  struct equal {
+    bool operator()(const StringPiece &a, const StringPiece &b) const {
+      return a.size == b.size && memcmp(a.data, b.data, a.size) == 0;
+    }
+  };
+
   inline u64 get_addr() const;
 
   std::atomic<MergeableSection *> isec = ATOMIC_VAR_INIT(nullptr);
@@ -173,9 +164,22 @@ public:
   Symbol(std::string_view name) : name(name) {}
   Symbol(const Symbol &other) : Symbol(other.name) {}
 
+  struct hash {
+    size_t operator()(const Symbol &sym) const {
+      return std::hash<std::string_view>()(sym.name);
+    }
+  };
+
+  struct equal {
+    bool operator()(const Symbol &a, const Symbol &b) const {
+      return a.name == b.name;
+    }
+  };
+
   static Symbol *intern(std::string_view name) {
-    static ConcurrentMap<Symbol> map;
-    return map.insert(name, Symbol(name));
+    static tbb::concurrent_unordered_set<Symbol, Symbol::hash, Symbol::equal> set;
+    auto [it, inserted] = set.insert(Symbol(name));
+    return &*it;
   }
 
   inline u64 get_addr() const;
@@ -565,7 +569,7 @@ public:
   static MergedSection *get_instance(std::string_view name, u32 type, u64 flags);
 
   static inline std::vector<MergedSection *> instances;
-  ConcurrentMap<StringPiece> map;
+  tbb::concurrent_unordered_set<StringPiece, StringPiece::hash, StringPiece::equal> set;
 
   void copy_buf() override;
 
diff --git a/object_file.cc b/object_file.cc
index 40b73177..2e2a350f 100644
--- a/object_file.cc
+++ b/object_file.cc
@@ -130,8 +130,10 @@ void ObjectFile::initialize_sections() {
       if (entries[0] != GRP_COMDAT)
         error(to_string(this) + ": unsupported SHT_GROUP format");
 
-      static ConcurrentMap<ComdatGroup> map;
-      ComdatGroup *group = map.insert(signature, ComdatGroup(nullptr, 0));
+      static tbb::concurrent_hash_map<std::string_view, ComdatGroup> map;
+      decltype(map)::const_accessor acc;
+      map.insert(acc, {signature, ComdatGroup(nullptr, 0)});
+      ComdatGroup *group = const_cast<ComdatGroup *>(&acc->second);
       comdat_groups.push_back({group, entries});
 
       static Counter counter("comdats");
diff --git a/output_chunks.cc b/output_chunks.cc
index bfd5bdbe..b6296881 100644
--- a/output_chunks.cc
+++ b/output_chunks.cc
@@ -663,10 +663,11 @@ MergedSection::get_instance(std::string_view name, u32 type, u64 flags) {
 void MergedSection::copy_buf() {
   u8 *base = out::buf + shdr.sh_offset;
 
-  map.for_each_value([&](const StringPiece &piece) {
+  for (decltype(set)::const_iterator it = set.begin(); it != set.end(); ++it) {
+    StringPiece &piece = *it;
     if (MergeableSection *m = piece.isec)
       memcpy(base + m->offset + piece.output_offset, piece.data, piece.size);
-  });
+  }
 }
 
 void CopyrelSection::add_symbol(Symbol *sym) {