From 5bc7e24c0915e1bdfdaa6a861c58fe9ac7301ebc Mon Sep 17 00:00:00 2001
From: Maxim Pimenov <m@maps.me>
Date: Tue, 25 Dec 2018 19:07:32 +0300
Subject: [PATCH] [base] A benchmark for Beam.

---
 base/base_tests/beam_tests.cpp | 47 ++++++++++++++++++++++++++++++----
 base/beam.hpp                  | 35 +++++++++++++++----------
 2 files changed, 64 insertions(+), 18 deletions(-)

diff --git a/base/base_tests/beam_tests.cpp b/base/base_tests/beam_tests.cpp
index 99ac331781..a575643d38 100644
--- a/base/base_tests/beam_tests.cpp
+++ b/base/base_tests/beam_tests.cpp
@@ -8,6 +8,7 @@
 
 #include <cstddef>
 #include <cstdint>
+#include <random>
 #include <string>
 #include <vector>
 
@@ -17,10 +18,10 @@ using namespace std;
 namespace
 {
 template <template <typename, typename> class Beam>
-void Smoke(string const & beamType)
+void Smoke()
 {
-  size_t const kCapacity = 100;
-  size_t const kTotal = 1000;
+  size_t const kCapacity = 10;
+  size_t const kTotal = 100;
 
   CHECK_LESS_OR_EQUAL(kCapacity, kTotal, ());
 
@@ -41,10 +42,46 @@ void Smoke(string const & beamType)
   sort(actual.rbegin(), actual.rend());
   CHECK_EQUAL(expected, actual, ());
 }
+
+template <template <typename, typename> class Beam>
+void Benchmark(string const & beamType, uint64_t const numResets, size_t const capacity,
+               uint64_t const numEvents)
+{
+  base::Timer timer;
+  SCOPE_GUARD(timerGuard,
+              [&] { LOG(LINFO, ("type:", beamType, "\ttime passed:", timer.ElapsedSeconds())); });
+
+  CHECK_LESS_OR_EQUAL(capacity, numEvents, ());
+
+  mt19937 rng(0);
+  uniform_real_distribution<double> dis(0.0, 1.0);
+  for (uint64_t wave = 0; wave <= numResets; ++wave)
+  {
+    Beam<uint64_t, double> beam(capacity);
+
+    uint64_t const begin = wave * numEvents / (numResets + 1);
+    uint64_t const end = (wave + 1) * numEvents / (numResets + 1);
+    for (uint64_t i = begin; i < end; ++i)
+      beam.Add(i, dis(rng));
+  }
+}
 }  // namespace
 
 UNIT_TEST(Beam_Smoke)
 {
-  Smoke<base::Beam>("Vector-based");
-  Smoke<base::HeapBeam>("Heap-based");
+  Smoke<base::Beam>();
+  Smoke<base::HeapBeam>();
+}
+
+UNIT_TEST(Beam_Benchmark)
+{
+  size_t const kCapacity = 100;
+  uint64_t const kNumEvents = 1000000;
+
+  for (uint64_t numResets = 0; numResets < 1000; numResets += 200)
+  {
+    LOG(LINFO, ("Resets =", numResets, "Capacity =", kCapacity, "Total events =", kNumEvents));
+    Benchmark<base::Beam>("Vector-based", numResets, kCapacity, kNumEvents);
+    Benchmark<base::HeapBeam>("Heap-based", numResets, kCapacity, kNumEvents);
+  }
 }
diff --git a/base/beam.hpp b/base/beam.hpp
index 9f41983120..2d5d7001d9 100644
--- a/base/beam.hpp
+++ b/base/beam.hpp
@@ -58,14 +58,17 @@ public:
     m_entries.insert(it, e);
   }
 
-  // template <typename Fn>
-  // void ForEachEntry(Fn && fn) const
-  // {
-  //   for (Entry const & e : m_entries)
-  //     fn(e.m_key, e.m_value);
-  // }
+  // Calls |fn| for all entries currently held in the beam.
+  // The order of calls is not specified.
+  template <typename Fn>
+  void ForEachEntry(Fn && fn) const
+  {
+    for (Entry const & e : m_entries)
+      fn(e.m_key, e.m_value);
+  }
 
-  // todo(@m) Specify whether the entries are supposed to be sorted.
+  // Returns all entries currently held in the beam.
+  // The order of returned entries is not specified.
   std::vector<Entry> const & GetEntries() const { return m_entries; }
 
 private:
@@ -73,6 +76,8 @@ private:
   std::vector<Entry> m_entries;
 };
 
+// A data structure to perform the beam search with.
+// Maintains a binary heap of (Key, Value) pairs.
 template <typename TKey, typename TValue>
 class HeapBeam
 {
@@ -122,13 +127,17 @@ public:
     std::push_heap(m_entries.begin(), m_entries.begin() + m_size);
   }
 
-  // template <typename Fn>
-  // void ForEachEntry(Fn && fn) const
-  // {
-  //   for (Entry const & e : m_entries)
-  //     fn(e.m_key, e.m_value);
-  // }
+  // Calls |fn| for all entries currently held in the beam.
+  // The order of calls is not specified.
+  template <typename Fn>
+  void ForEachEntry(Fn && fn) const
+  {
+    for (Entry const & e : m_entries)
+      fn(e.m_key, e.m_value);
+  }
 
+  // Returns all entries currently held in the beam.
+  // The order of returned entries is not specified.
   std::vector<Entry> const & GetEntries() const { return m_entries; }
 
 private: