[search] Prefer prevEmit results in NearbyPointsSweeper

3 files changed, 110 insertions, 29 deletions

diff --git a/geometry/geometry_tests/nearby_points_sweeper_test.cpp b/geometry/geometry_tests/nearby_points_sweeper_test.cpp
index cdcaea5f99..4e84963fb2 100644
--- a/geometry/geometry_tests/nearby_points_sweeper_test.cpp
+++ b/geometry/geometry_tests/nearby_points_sweeper_test.cpp
@@ -5,10 +5,12 @@
 
 #include "base/stl_add.hpp"
 
-#include "std/set.hpp"
-#include "std/vector.hpp"
+#include <set>
+#include <utility>
+#include <vector>
 
 using namespace m2;
+using namespace std;
 
 namespace search
 {
@@ -20,9 +22,10 @@ using TIndexSet = multiset<size_t>;
 UNIT_TEST(NearbyPointsSweeper_Smoke)
 {
   {
+    uint8_t const priority = 0;
     NearbyPointsSweeper sweeper(0.0);
     for (size_t i = 0; i < 10; ++i)
-      sweeper.Add(10.0, 10.0, i);
+      sweeper.Add(10.0, 10.0, i, priority);
 
     TIndexSet expected = {0};
 
@@ -33,13 +36,14 @@ UNIT_TEST(NearbyPointsSweeper_Smoke)
   }
 
   {
+    uint8_t const priority = 0;
     vector<double> const coords = {0.0, 0.5, 1.0, 1.5, 1.4, 1.6};
 
     {
       NearbyPointsSweeper sweeper(0.5);
 
       for (size_t i = 0; i < coords.size(); ++i)
-        sweeper.Add(coords[i], 0.0, i);
+        sweeper.Add(coords[i], 0.0, i, priority);
 
       TIndexSet expected = {0, 2, 5};
 
@@ -53,7 +57,7 @@ UNIT_TEST(NearbyPointsSweeper_Smoke)
       NearbyPointsSweeper sweeper(0.5);
 
       for (size_t i = 0; i < coords.size(); ++i)
-        sweeper.Add(0.0, coords[i], i);
+        sweeper.Add(0.0, coords[i], i, priority);
 
       TIndexSet expected = {0, 2, 5};
 
@@ -65,12 +69,13 @@ UNIT_TEST(NearbyPointsSweeper_Smoke)
   }
 
   {
-    vector<m2::PointD> const points = {m2::PointD(0.0, 0.0), m2::PointD(1.0, 1.0),
-                                       m2::PointD(1.5, 0.0), m2::PointD(1.5 + 1.01, 1.5 + 1.0)};
+    uint8_t const priority = 0;
+    vector<PointD> const points = {PointD(0.0, 0.0), PointD(1.0, 1.0),
+                                   PointD(1.5, 0.0), PointD(1.5 + 1.01, 1.5 + 1.0)};
     NearbyPointsSweeper sweeper(1.0);
 
     for (size_t i = 0; i < points.size(); ++i)
-      sweeper.Add(points[i].x, points[i].y, i);
+      sweeper.Add(points[i].x, points[i].y, i, priority);
 
     TIndexSet expected = {0, 2, 3};
 
@@ -83,9 +88,10 @@ UNIT_TEST(NearbyPointsSweeper_Smoke)
 
 UNIT_TEST(NearbyPointsSweeper_CornerCases)
 {
-  vector<m2::PointD> const points = {m2::PointD(0, 0),     m2::PointD(0, 0), m2::PointD(1, 0),
-                                     m2::PointD(0, 1),     m2::PointD(1, 1), m2::PointD(1, 0),
-                                     m2::PointD(0.5, 0.5), m2::PointD(0, 1)};
+  uint8_t const priority = 0;
+  vector<PointD> const points = {PointD(0, 0),     PointD(0, 0), PointD(1, 0),
+                                 PointD(0, 1),     PointD(1, 1), PointD(1, 0),
+                                 PointD(0.5, 0.5), PointD(0, 1)};
 
   {
     // In accordance with the specification, all points must be left
@@ -95,7 +101,7 @@ UNIT_TEST(NearbyPointsSweeper_CornerCases)
     TIndexSet expected;
     for (size_t i = 0; i < points.size(); ++i)
     {
-      sweeper.Add(points[i].x, points[i].y, i);
+      sweeper.Add(points[i].x, points[i].y, i, priority);
       expected.insert(i);
     }
 
@@ -108,7 +114,7 @@ UNIT_TEST(NearbyPointsSweeper_CornerCases)
   {
     NearbyPointsSweeper sweeper(10.0);
     for (size_t i = 0; i < points.size(); ++i)
-      sweeper.Add(points[i].x, points[i].y, i);
+      sweeper.Add(points[i].x, points[i].y, i, priority);
 
     TIndexSet expected = {0};
     TIndexSet actual;
@@ -117,5 +123,50 @@ UNIT_TEST(NearbyPointsSweeper_CornerCases)
     TEST_EQUAL(expected, actual, ());
   }
 }
+
+UNIT_TEST(NearbyPointsSweeper_Priority)
+{
+  {
+    NearbyPointsSweeper sweeper(0.0);
+    for (size_t i = 0; i < 10; ++i)
+      sweeper.Add(10.0, 10.0, i /* index */, i /* priority */);
+
+    TIndexSet expected = {9};
+
+    TIndexSet actual;
+    sweeper.Sweep(MakeInsertFunctor(actual));
+
+    TEST_EQUAL(expected, actual, ());
+  }
+  {
+    vector<pair<double, uint8_t>> const objects = {{0.0, 0}, {0.5, 1}, {1.0, 1}, {1.5, 1}, {1.4, 0}, {1.6, 0}};
+    NearbyPointsSweeper sweeper(0.5);
+
+    for (size_t i = 0; i < objects.size(); ++i)
+      sweeper.Add(objects[i].first, 0.0, i /* index */, objects[i].second);
+
+    TIndexSet expected = {1, 3};
+
+    TIndexSet actual;
+    sweeper.Sweep(MakeInsertFunctor(actual));
+
+    TEST_EQUAL(expected, actual, ());
+  }
+  {
+    vector<pair<PointD, uint8_t>> const objects = {{PointD(0.0, 0.0), 0}, {PointD(1.0, 1.0), 1},
+                                                   {PointD(1.5, 0.0), 0}, {PointD(1.5 + 1.01, 1.5 + 1.0), 0}};
+    NearbyPointsSweeper sweeper(1.0);
+
+    for (size_t i = 0; i < objects.size(); ++i)
+      sweeper.Add(objects[i].first.x, objects[i].first.y, i /* index */, objects[i].second);
+
+    TIndexSet expected = {1, 3};
+
+    TIndexSet actual;
+    sweeper.Sweep(MakeInsertFunctor(actual));
+
+    TEST_EQUAL(expected, actual, ());
+  }
+}
 }  // namespace
 }  // namespace search
diff --git a/geometry/nearby_points_sweeper.cpp b/geometry/nearby_points_sweeper.cpp
index 65f09e9d34..f50ec52c1b 100644
--- a/geometry/nearby_points_sweeper.cpp
+++ b/geometry/nearby_points_sweeper.cpp
@@ -3,8 +3,8 @@
 namespace m2
 {
 // NearbyPointsSweeper::Event ----------------------------------------------------------------------
-NearbyPointsSweeper::Event::Event(Type type, double y, double x, size_t index)
-  : m_type(type), m_y(y), m_x(x), m_index(index)
+NearbyPointsSweeper::Event::Event(Type type, double y, double x, size_t index, uint8_t priority)
+  : m_type(type), m_y(y), m_x(x), m_index(index), m_priority(priority)
 {
 }
 
@@ -19,7 +19,10 @@ bool NearbyPointsSweeper::Event::operator<(Event const & rhs) const
   if (m_x != rhs.m_x)
     return m_x < rhs.m_x;
 
-  return m_index < rhs.m_index;
+  if (m_index != rhs.m_index)
+    return m_index < rhs.m_index;
+
+  return m_priority < rhs.m_priority;
 }
 
 // NearbyPointsSweeper -----------------------------------------------------------------------------
@@ -27,9 +30,9 @@ NearbyPointsSweeper::NearbyPointsSweeper(double eps) : m_eps(eps), m_heps(std::m
 {
 }
 
-void NearbyPointsSweeper::Add(double x, double y, size_t index)
+void NearbyPointsSweeper::Add(double x, double y, size_t index, uint8_t priority)
 {
-  m_events.emplace_back(Event::TYPE_SEGMENT_START, y - m_heps, x, index);
-  m_events.emplace_back(Event::TYPE_SEGMENT_END, y + m_heps, x, index);
+  m_events.emplace_back(Event::TYPE_SEGMENT_START, y - m_heps, x, index, priority);
+  m_events.emplace_back(Event::TYPE_SEGMENT_END, y + m_heps, x, index, priority);
 }
 }  // namespace m2
diff --git a/geometry/nearby_points_sweeper.hpp b/geometry/nearby_points_sweeper.hpp
index 318ce7da47..9c45da7f2a 100644
--- a/geometry/nearby_points_sweeper.hpp
+++ b/geometry/nearby_points_sweeper.hpp
@@ -1,6 +1,7 @@
 #pragma once
 
 #include "base/assert.hpp"
+#include "base/visitor.hpp"
 
 #include <algorithm>
 #include <cmath>
@@ -28,7 +29,8 @@ public:
   // Adds a new point (|x|, |y|) on the plane. |index| is used to
   // identify individual points, and will be reported for survived
   // points during the Sweep phase.
-  void Add(double x, double y, size_t index);
+  // Points with higher |priority| will be preferred during the Sweep phase.
+  void Add(double x, double y, size_t index, uint8_t priority);
 
   // Emits indexes of all survived points. Complexity: O(n * log(n)),
   // where n is the number of already added points.
@@ -37,7 +39,7 @@ public:
   {
     std::sort(m_events.begin(), m_events.end());
 
-    std::set<std::pair<double, size_t>> line;
+    std::set<LineEvent> line;
 
     for (auto const & event : m_events)
     {
@@ -47,11 +49,12 @@ public:
       {
         if (line.empty())
         {
-          line.emplace(event.m_x, event.m_index);
+          line.insert({event.m_x, event.m_index, event.m_priority});
           break;
         }
 
-        auto it = line.upper_bound(std::make_pair(event.m_x, std::numeric_limits<size_t>::max()));
+        auto it = line.upper_bound(
+            {event.m_x, std::numeric_limits<size_t>::max(), std::numeric_limits<uint8_t>::max()});
 
         bool add = true;
         while (true)
@@ -65,11 +68,16 @@ public:
             continue;
           }
 
-          double const x = it->first;
+          double const x = it->m_x;
           if (fabs(x - event.m_x) <= m_eps)
           {
-            add = false;
-            break;
+            if (it->m_priority >= event.m_priority)
+            {
+              add = false;
+              break;
+            }
+            // New event has higher priority than an existing event nearby.
+            it = line.erase(it);
           }
 
           if (x + m_eps < event.m_x)
@@ -82,14 +90,14 @@ public:
         }
 
         if (add)
-          line.emplace(event.m_x, event.m_index);
+          line.insert({event.m_x, event.m_index, event.m_priority});
 
         break;
       }
 
       case Event::TYPE_SEGMENT_END:
       {
-        auto it = line.find(std::make_pair(event.m_x, event.m_index));
+        auto it = line.find({event.m_x, event.m_index, event.m_priority});
         if (it != line.end())
         {
           emitter(event.m_index);
@@ -112,7 +120,7 @@ private:
       TYPE_SEGMENT_END
     };
 
-    Event(Type type, double y, double x, size_t index);
+    Event(Type type, double y, double x, size_t index, uint8_t priority);
 
     bool operator<(Event const & rhs) const;
 
@@ -121,6 +129,25 @@ private:
     double m_y;
     double m_x;
     size_t m_index;
+    uint8_t m_priority;
+  };
+
+  struct LineEvent
+  {
+    DECLARE_VISITOR_AND_DEBUG_PRINT(LineEvent, visitor(m_x, "x"), visitor(m_index, "index"),
+                                    visitor(m_priority, "priority"));
+    bool operator<(LineEvent const & rhs) const
+    {
+      if (m_x != rhs.m_x)
+        return m_x < rhs.m_x;
+      if (m_index < rhs.m_index)
+        return m_index < rhs.m_index;
+      return m_priority < rhs.m_priority;
+    }
+
+    double m_x;
+    size_t m_index;
+    uint8_t m_priority;
   };
 
   std::vector<Event> m_events;