[esan] Add iterator to esan's generic hashtable

Summary: Adds simple iterator support to the esan hashtable.

Reviewers: aizatsky

Subscribers: vitalybuka, zhaoqin, kcc, eugenis, llvm-commits, kubabrecka

Differential Revision: https://reviews.llvm.org/D22682

llvm-svn: 278027

compiler-rt/lib/esan/esan_hashtable.h

@@ -66,19 +66,58 @@ public:
   // If the table is internally-synchronized, this lock must not be held
   // while a hashtable function is called as it will deadlock: the lock
   // is not recursive.  This is meant for use with externally-synchronized
-  // tables.
+  // tables or with an iterator.
   void lock();
   void unlock();
 
 private:
-  void resize();
-
   struct HashEntry {
     KeyTy Key;
     DataTy Payload;
     HashEntry *Next;
   };
 
+public:
+  struct HashPair {
+    HashPair(KeyTy Key, DataTy Data) : Key(Key), Data(Data) {}
+    KeyTy Key;
+    DataTy Data;
+  };
+
+  // This iterator does not perform any synchronization.
+  // It expects the caller to lock the table across the whole iteration.
+  // Calling HashTable functions while using the iterator is not supported.
+  // The iterator returns copies of the keys and data.
+  class iterator {
+  public:
+    iterator(
+        HashTable<KeyTy, DataTy, ExternalLock, HashFuncTy, EqualFuncTy> *Table);
+    iterator(const iterator &Src) = default;
+    iterator &operator=(const iterator &Src) = default;
+    HashPair operator*();
+    iterator &operator++();
+    iterator &operator++(int);
+    bool operator==(const iterator &Cmp) const;
+    bool operator!=(const iterator &Cmp) const;
+
+  private:
+    iterator(
+        HashTable<KeyTy, DataTy, ExternalLock, HashFuncTy, EqualFuncTy> *Table,
+        int Idx);
+    friend HashTable<KeyTy, DataTy, ExternalLock, HashFuncTy, EqualFuncTy>;
+    HashTable<KeyTy, DataTy, ExternalLock, HashFuncTy, EqualFuncTy> *Table;
+    int Idx;
+    HashTable<KeyTy, DataTy, ExternalLock, HashFuncTy, EqualFuncTy>::HashEntry
+        *Entry;
+  };
+
+  // No erase or insert iterator supported
+  iterator begin();
+  iterator end();
+
+private:
+  void resize();
+
   HashEntry **Table;
   u32 Capacity;
   u32 Entries;
@@ -247,4 +286,96 @@ void HashTable<KeyTy, DataTy, ExternalLock, HashFuncTy, EqualFuncTy>::unlock() {
   Mutex.Unlock();
 }
 
+//===----------------------------------------------------------------------===//
+// Iterator implementation
+//===----------------------------------------------------------------------===//
+
+template <typename KeyTy, typename DataTy, bool ExternalLock,
+          typename HashFuncTy, typename EqualFuncTy>
+HashTable<KeyTy, DataTy, ExternalLock, HashFuncTy, EqualFuncTy>::iterator::
+    iterator(
+        HashTable<KeyTy, DataTy, ExternalLock, HashFuncTy, EqualFuncTy> *Table)
+    : Table(Table), Idx(-1), Entry(nullptr) {
+  operator++();
+}
+
+template <typename KeyTy, typename DataTy, bool ExternalLock,
+          typename HashFuncTy, typename EqualFuncTy>
+HashTable<KeyTy, DataTy, ExternalLock, HashFuncTy, EqualFuncTy>::iterator::
+    iterator(
+        HashTable<KeyTy, DataTy, ExternalLock, HashFuncTy, EqualFuncTy> *Table,
+        int Idx)
+    : Table(Table), Idx(Idx), Entry(nullptr) {
+  CHECK(Idx >= (int)Table->Capacity); // Only used to create end().
+}
+
+template <typename KeyTy, typename DataTy, bool ExternalLock,
+          typename HashFuncTy, typename EqualFuncTy>
+typename HashTable<KeyTy, DataTy, ExternalLock, HashFuncTy,
+                   EqualFuncTy>::HashPair
+    HashTable<KeyTy, DataTy, ExternalLock, HashFuncTy, EqualFuncTy>::iterator::
+    operator*() {
+  CHECK(Idx >= 0 && Idx < (int)Table->Capacity);
+  CHECK(Entry != nullptr);
+  return HashPair(Entry->Key, Entry->Payload);
+}
+
+template <typename KeyTy, typename DataTy, bool ExternalLock,
+          typename HashFuncTy, typename EqualFuncTy>
+typename HashTable<KeyTy, DataTy, ExternalLock, HashFuncTy,
+                   EqualFuncTy>::iterator &
+    HashTable<KeyTy, DataTy, ExternalLock, HashFuncTy, EqualFuncTy>::iterator::
+    operator++() {
+  if (Entry != nullptr)
+    Entry = Entry->Next;
+  while (Entry == nullptr) {
+    ++Idx;
+    if (Idx >= (int)Table->Capacity)
+      break; // At end().
+    Entry = Table->Table[Idx];
+  }
+  return *this;
+}
+
+template <typename KeyTy, typename DataTy, bool ExternalLock,
+          typename HashFuncTy, typename EqualFuncTy>
+typename HashTable<KeyTy, DataTy, ExternalLock, HashFuncTy,
+                   EqualFuncTy>::iterator &
+    HashTable<KeyTy, DataTy, ExternalLock, HashFuncTy, EqualFuncTy>::iterator::
+    operator++(int) {
+  iterator Temp(*this);
+  operator++();
+  return Temp;
+}
+
+template <typename KeyTy, typename DataTy, bool ExternalLock,
+          typename HashFuncTy, typename EqualFuncTy>
+bool HashTable<KeyTy, DataTy, ExternalLock, HashFuncTy, EqualFuncTy>::iterator::
+operator==(const iterator &Cmp) const {
+  return Cmp.Table == Table && Cmp.Idx == Idx && Cmp.Entry == Entry;
+}
+
+template <typename KeyTy, typename DataTy, bool ExternalLock,
+          typename HashFuncTy, typename EqualFuncTy>
+bool HashTable<KeyTy, DataTy, ExternalLock, HashFuncTy, EqualFuncTy>::iterator::
+operator!=(const iterator &Cmp) const {
+  return Cmp.Table != Table || Cmp.Idx != Idx || Cmp.Entry != Entry;
+}
+
+template <typename KeyTy, typename DataTy, bool ExternalLock,
+          typename HashFuncTy, typename EqualFuncTy>
+typename HashTable<KeyTy, DataTy, ExternalLock, HashFuncTy,
+                   EqualFuncTy>::iterator
+HashTable<KeyTy, DataTy, ExternalLock, HashFuncTy, EqualFuncTy>::begin() {
+  return iterator(this);
+}
+
+template <typename KeyTy, typename DataTy, bool ExternalLock,
+          typename HashFuncTy, typename EqualFuncTy>
+typename HashTable<KeyTy, DataTy, ExternalLock, HashFuncTy,
+                   EqualFuncTy>::iterator
+HashTable<KeyTy, DataTy, ExternalLock, HashFuncTy, EqualFuncTy>::end() {
+  return iterator(this, Capacity);
+}
+
 } // namespace __esan

compiler-rt/test/esan/Unit/hashtable.cpp

@@ -56,6 +56,15 @@ int main()
 {
   __esan::HashTable<int, int> IntTable;
   assert(IntTable.size() == 0);
+
+  // Test iteration on an empty table.
+  int Count = 0;
+  for (auto Iter = IntTable.begin(); Iter != IntTable.end();
+       ++Iter, ++Count) {
+    // Empty.
+  }
+  assert(Count == 0);
+
   bool Added = IntTable.add(4, 42);
   assert(Added);
   assert(!IntTable.add(4, 42));
@@ -63,9 +72,21 @@ int main()
   int Value;
   bool Found = IntTable.lookup(4, Value);
   assert(Found && Value == 42);
+
+  // Test iterator.
+  IntTable.lock();
+  for (auto Iter = IntTable.begin(); Iter != IntTable.end();
+       ++Iter, ++Count) {
+    assert((*Iter).Key == 4);
+    assert((*Iter).Data == 42);
+  }
+  IntTable.unlock();
+  assert(Count == 1);
+  assert(Count == IntTable.size());
   assert(!IntTable.remove(5));
   assert(IntTable.remove(4));
 
+  // Test a more complex payload.
   __esan::HashTable<int, MyDataPayload> DataTable(4);
   MyDataPayload NewData(new MyData("mystring"));
   Added = DataTable.add(4, NewData);
@@ -86,6 +107,30 @@ int main()
     Found = DataTable.lookup(i+1, FoundData);
     assert(Found && strcmp(FoundData.Data->Buf, "delete-at-end") == 0);
   }
+  DataTable.lock();
+  Count = 0;
+  for (auto Iter = DataTable.begin(); Iter != DataTable.end();
+       ++Iter, ++Count) {
+    int Key = (*Iter).Key;
+    FoundData = (*Iter).Data;
+    assert(Key >= 1 && Key <= 4);
+    assert(strcmp(FoundData.Data->Buf, "delete-at-end") == 0);
+  }
+  DataTable.unlock();
+  assert(Count == 4);
+  assert(Count == DataTable.size());
+
+  // Ensure the iterator supports a range-based for loop.
+  DataTable.lock();
+  Count = 0;
+  for (auto Pair : DataTable) {
+    assert(Pair.Key >= 1 && Pair.Key <= 4);
+    assert(strcmp(Pair.Data.Data->Buf, "delete-at-end") == 0);
+    ++Count;
+  }
+  DataTable.unlock();
+  assert(Count == 4);
+  assert(Count == DataTable.size());
 
   // Test payload freeing via smart pointer wrapper.
   __esan::HashTable<MyDataPayload, MyDataPayload, true> DataKeyTable;