[indexer] Add distance for closest object selection

2 files changed, 65 insertions, 21 deletions

diff --git a/indexer/interval_index.hpp b/indexer/interval_index.hpp
index 49e8888da2..c63fe9e954 100644
--- a/indexer/interval_index.hpp
+++ b/indexer/interval_index.hpp
@@ -66,15 +66,14 @@ public:
         end = KeyEnd();
       --end;  // end is inclusive in ForEachImpl().
       ForEachNode(f, beg, end, m_Header.m_Levels, 0,
-                  m_LevelOffsets[m_Header.m_Levels + 1] - m_LevelOffsets[m_Header.m_Levels]);
+                  m_LevelOffsets[m_Header.m_Levels + 1] - m_LevelOffsets[m_Header.m_Levels], 0);
     }
   }
 
 private:
-
   template <typename F>
   void ForEachLeaf(F const & f, uint64_t const beg, uint64_t const end,
-                   uint32_t const offset, uint32_t const size) const
+                   uint32_t const offset, uint32_t const size, uint64_t keyBase) const
   {
     buffer_vector<uint8_t, 1024> data;
     data.resize_no_init(size);
@@ -93,19 +92,19 @@ private:
         break;
       value += ReadVarInt<int64_t>(src);
       if (key >= beg)
-        f(value);
+        Invoke(f, keyBase + key, value, 0);
     }
   }
 
   template <typename F>
   void ForEachNode(F const & f, uint64_t beg, uint64_t end, int level,
-                   uint32_t offset, uint32_t size) const
+                   uint32_t offset, uint32_t size, uint64_t nodeKey) const
   {
     offset += m_LevelOffsets[level];
 
     if (level == 0)
     {
-      ForEachLeaf(f, beg, end, offset, size);
+      ForEachLeaf(f, beg, end, offset, size, nodeKey);
       return;
     }
 
@@ -139,7 +138,7 @@ private:
           {
             uint64_t const beg1 = (i == beg0) ? (beg & levelBytesFF) : 0;
             uint64_t const end1 = (i == end0) ? (end & levelBytesFF) : levelBytesFF;
-            ForEachNode(f, beg1, end1, level - 1, childOffset, childSize);
+            ForEachNode(f, beg1, end1, level - 1, childOffset, childSize, nodeKey + (uint64_t{i} << skipBits));
           }
           childOffset += childSize;
         }
@@ -161,13 +160,24 @@ private:
         {
           uint64_t const beg1 = (i == beg0) ? (beg & levelBytesFF) : 0;
           uint64_t const end1 = (i == end0) ? (end & levelBytesFF) : levelBytesFF;
-          ForEachNode(f, beg1, end1, level - 1, childOffset, childSize);
+          ForEachNode(f, beg1, end1, level - 1, childOffset, childSize, nodeKey + (uint64_t{i} << skipBits));
         }
         childOffset += childSize;
       }
     }
   }
 
+  template <typename F, typename = decltype(std::declval<F>()(uint64_t{0}, uint64_t{0}))>
+  static void Invoke(F const & f, uint64_t key, uint64_t storedId, int)
+  {
+    f(key, storedId);
+  }
+  template <typename F>
+  static void Invoke(F const & f, uint64_t key, uint64_t storedId, ...)
+  {
+    f(storedId);
+  }
+
   ReaderT m_Reader;
   Header m_Header;
   buffer_vector<uint32_t, 7> m_LevelOffsets;
diff --git a/indexer/locality_index.hpp b/indexer/locality_index.hpp
index a8665d291f..bc72a544c9 100644
--- a/indexer/locality_index.hpp
+++ b/indexer/locality_index.hpp
@@ -12,11 +12,17 @@
 
 #include "base/geo_object_id.hpp"
 
+#include <boost/multi_index_container.hpp>
+#include <boost/multi_index/hashed_index.hpp>
+#include <boost/multi_index/member.hpp>
+#include <boost/multi_index/ordered_index.hpp>
+
 #include <cstdint>
 #include <functional>
 #include <memory>
 #include <set>
 #include <string>
+#include <utility>
 
 #include "defines.hpp"
 
@@ -30,6 +36,7 @@ class LocalityIndex
 {
 public:
   using ProcessObject = std::function<void(base::GeoObjectId const &)>;
+  using ProcessClosestObject = std::function<void(base::GeoObjectId const &, double distance)>;
 
   LocalityIndex() = default;
   explicit LocalityIndex(Reader const & reader)
@@ -59,7 +66,7 @@ public:
   // and stop after |sizeHint| objects have been processed. For stability, if an object from
   // an index cell has been processed, all other objects from this cell will be processed too,
   // thus probably overflowing the |sizeHint| limit.
-  void ForClosestToPoint(ProcessObject const & processObject, m2::PointD const & center,
+  void ForClosestToPoint(ProcessClosestObject const & processObject, m2::PointD const & center,
                          double radiusM, uint32_t sizeHint) const
   {
     auto const rect =
@@ -67,17 +74,6 @@ public:
     covering::CoveringGetter cov(rect, covering::CoveringMode::Spiral);
     covering::Intervals const & intervals = cov.Get<DEPTH_LEVELS>(scales::GetUpperScale());
 
-    std::set<uint64_t> objects;
-    auto processAll = [&objects, &processObject](uint64_t storedId) {
-      if (objects.insert(storedId).second)
-        processObject(LocalityObject::FromStoredId(storedId));
-    };
-
-    auto process = [&](uint64_t storedId) {
-      if (objects.size() < sizeHint)
-        processAll(storedId);
-    };
-
     CHECK_EQUAL(intervals.begin()->first, intervals.begin()->second - 1, ());
     auto cellDepth = covering::GetCodingDepth<DEPTH_LEVELS>(scales::GetUpperScale());
     auto bestCell = m2::CellId<DEPTH_LEVELS>::FromInt64(intervals.begin()->first, cellDepth);
@@ -88,6 +84,41 @@ public:
       bestCell = bestCell.Parent();
     }
 
+    struct ResultItem
+    {
+      uint64_t m_ObjectId;
+      double m_Distance;
+    };
+    using namespace boost::multi_index;
+    auto result = multi_index_container<ResultItem,
+                      indexed_by<hashed_unique<member<ResultItem, uint64_t, &ResultItem::m_ObjectId>>,
+                                 ordered_non_unique<member<ResultItem, double, &ResultItem::m_Distance>>>>{};
+
+    auto processAll = [&result, &bestCells, cellDepth, &center](int64_t cellNumber, uint64_t storedId) {
+      auto distance = 0.0;
+      if (bestCells.find(cellNumber) == bestCells.end())
+      {
+        using Converter = CellIdConverter<MercatorBounds, m2::CellId<DEPTH_LEVELS>>;
+
+        auto cell = m2::CellId<DEPTH_LEVELS>::FromInt64(cellNumber, cellDepth);
+        auto cellCenter = Converter::FromCellId(cell);
+        distance = MercatorBounds::DistanceOnEarth(center, cellCenter);
+      }
+
+      auto objectId = LocalityObject::FromStoredId(storedId).GetEncodedId();
+      auto resultInsert = result.insert({objectId, distance});
+      if (resultInsert.second)
+        return;
+      auto const & object = *resultInsert.first;
+      if (distance < object.m_Distance)
+        result.replace(resultInsert.first, {objectId, distance});
+    };
+
+    auto process = [&](int64_t cellNumber, uint64_t storedId) {
+      if (result.size() < sizeHint)
+        processAll(cellNumber, storedId);
+    };
+
     for (auto const & i : intervals)
     {
       if (bestCells.find(i.first) != bestCells.end())
@@ -97,10 +128,13 @@ public:
       else
       {
         m_intervalIndex->ForEach(process, i.first, i.second);
-        if (objects.size() >= sizeHint)
+        if (result.size() >= sizeHint)
           return;
       }
     }
+
+    for (auto const & object : result.template get<1>())
+      processObject(base::GeoObjectId(object.m_ObjectId), object.m_Distance);
   }
 
 private: