From 0f811c601ed09746bdbf44eea16d4310196de388 Mon Sep 17 00:00:00 2001
From: Maxim Pimenov <m@maps.me>
Date: Wed, 19 Dec 2018 17:50:38 +0300
Subject: [PATCH] [base] An alternative implementation of Beam.

---
 base/base_tests/beam_tests.cpp | 25 +++++++++---
 base/beam.hpp                  | 75 +++++++++++++++++++++++++++++++++-
 2 files changed, 93 insertions(+), 7 deletions(-)

diff --git a/base/base_tests/beam_tests.cpp b/base/base_tests/beam_tests.cpp
index 02d66be135..99ac331781 100644
--- a/base/base_tests/beam_tests.cpp
+++ b/base/base_tests/beam_tests.cpp
@@ -3,21 +3,28 @@
 #include "base/assert.hpp"
 #include "base/beam.hpp"
 
+#include "base/scope_guard.hpp"
+#include "base/timer.hpp"
+
 #include <cstddef>
 #include <cstdint>
+#include <string>
 #include <vector>
 
 using namespace base;
 using namespace std;
 
-UNIT_TEST(Beam_Smoke)
+namespace
 {
-  size_t const kCapacity = 10;
-  size_t const kTotal = 100;
+template <template <typename, typename> class Beam>
+void Smoke(string const & beamType)
+{
+  size_t const kCapacity = 100;
+  size_t const kTotal = 1000;
 
   CHECK_LESS_OR_EQUAL(kCapacity, kTotal, ());
-  
-  base::Beam<uint32_t, double> beam(kCapacity);
+
+  Beam<uint32_t, double> beam(kCapacity);
 
   for (uint32_t i = 0; i < kTotal; ++i)
     beam.Add(i, static_cast<double>(i));
@@ -31,5 +38,13 @@ UNIT_TEST(Beam_Smoke)
   for (auto const & e : beam.GetEntries())
     actual.emplace_back(e.m_value);
 
+  sort(actual.rbegin(), actual.rend());
   CHECK_EQUAL(expected, actual, ());
 }
+}  // namespace
+
+UNIT_TEST(Beam_Smoke)
+{
+  Smoke<base::Beam>("Vector-based");
+  Smoke<base::HeapBeam>("Heap-based");
+}
diff --git a/base/beam.hpp b/base/beam.hpp
index bb9d21f94a..9f41983120 100644
--- a/base/beam.hpp
+++ b/base/beam.hpp
@@ -3,6 +3,8 @@
 #include "base/geo_object_id.hpp"
 #include "base/macros.hpp"
 
+#include <algorithm>
+#include <utility>
 #include <vector>
 
 namespace base
@@ -29,8 +31,12 @@ public:
 
   explicit Beam(size_t capacity) : m_capacity(capacity) { m_entries.reserve(m_capacity); }
 
-  // O(log(n) + n) for |n| entries.
-  // O(|m_capacity|) in the worst case.
+  // Tries to add a key-value pair to the beam. Evicts the entry with the lowest
+  // value if the insertion would result in an overspill.
+  //
+  // Complexity: O(log(m_capacity)) in the best case (a typical application has a considerable
+  //                                fraction of calls that do not result in an eviction).
+  //             O(m_capacity) in the worst case.
   void Add(Key const & key, Value const & value)
   {
     if (PREDICT_FALSE(m_capacity == 0))
@@ -59,10 +65,75 @@ public:
   //     fn(e.m_key, e.m_value);
   // }
 
+  // todo(@m) Specify whether the entries are supposed to be sorted.
   std::vector<Entry> const & GetEntries() const { return m_entries; }
 
 private:
   size_t m_capacity;
   std::vector<Entry> m_entries;
 };
+
+template <typename TKey, typename TValue>
+class HeapBeam
+{
+public:
+  using Key = TKey;
+  using Value = TValue;
+
+  struct Entry
+  {
+    Key m_key;
+    Value m_value;
+
+    Entry(Key const & key, Value const & value) : m_key(key), m_value(value) {}
+
+    bool operator<(Entry const & rhs) const { return m_value > rhs.m_value; }
+  };
+
+  explicit HeapBeam(size_t capacity) : m_capacity(capacity), m_size(0)
+  {
+    m_entries.reserve(m_capacity);
+  }
+
+  // Tries to add a key-value pair to the beam. Evicts the entry with the lowest
+  // value if the insertion would result in an overspill.
+  //
+  // Complexity: O(log(m_capacity)).
+  void Add(Key const & key, Value const & value)
+  {
+    if (PREDICT_FALSE(m_capacity == 0))
+      return;
+
+    Entry const e(key, value);
+    if (PREDICT_FALSE(m_size < m_capacity))
+    {
+      ++m_size;
+      m_entries.emplace_back(e);
+      std::push_heap(m_entries.begin(), m_entries.begin() + m_size);
+      return;
+    }
+
+    ASSERT_GREATER(m_size, 0, ());
+    if (m_entries.front() < e)
+      return;
+
+    std::pop_heap(m_entries.begin(), m_entries.begin() + m_size);
+    m_entries[m_size - 1] = std::move(e);
+    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);
+  // }
+
+  std::vector<Entry> const & GetEntries() const { return m_entries; }
+
+private:
+  size_t m_capacity;
+  size_t m_size;
+  std::vector<Entry> m_entries;
+};
 }  // namespace base