From 5e01cf40ddc6dad07965a82f06ed09c5e87b3b9c Mon Sep 17 00:00:00 2001
From: Amael Marquez <amael.marquez@irit.fr>
Date: Tue, 31 Oct 2023 17:05:52 +0100
Subject: [PATCH] [SpatialPartitioning] Add kd-tree sample index iterator and
 query

---
 .../Iterator/kdTreeSampleIndexIterator.h      | 42 ++++++++++++++++++
 .../KdTree/Query/kdTreeSampleIndexQuery.h     | 44 +++++++++++++++++++
 Ponca/src/SpatialPartitioning/KdTree/kdTree.h |  9 +++-
 3 files changed, 93 insertions(+), 2 deletions(-)
 create mode 100644 Ponca/src/SpatialPartitioning/KdTree/Iterator/kdTreeSampleIndexIterator.h
 create mode 100644 Ponca/src/SpatialPartitioning/KdTree/Query/kdTreeSampleIndexQuery.h

diff --git a/Ponca/src/SpatialPartitioning/KdTree/Iterator/kdTreeSampleIndexIterator.h b/Ponca/src/SpatialPartitioning/KdTree/Iterator/kdTreeSampleIndexIterator.h
new file mode 100644
index 000000000..55df0aa7a
--- /dev/null
+++ b/Ponca/src/SpatialPartitioning/KdTree/Iterator/kdTreeSampleIndexIterator.h
@@ -0,0 +1,42 @@
+/*
+ This Source Code Form is subject to the terms of the Mozilla Public
+ License, v. 2.0. If a copy of the MPL was not distributed with this
+ file, You can obtain one at http://mozilla.org/MPL/2.0/.
+*/
+
+#pragma once
+
+namespace Ponca {
+
+template <typename Traits>
+class KdTreeSampleIndexQuery;
+
+template <typename Traits>
+class KdTreeSampleIndexIterator
+{
+public:
+    using IndexType = typename Traits::IndexType;
+
+    void operator++() { ++m_index; }
+
+    bool operator!=(const KdTreeSampleIndexIterator& other) const {
+        return m_points != other.m_points || m_samples != other.m_samples || m_index != other.m_index;
+    }
+
+    IndexType operator*() const { return (*m_samples)[m_index]; }
+
+private:
+    friend KdTreeSampleIndexQuery<Traits>;
+
+    using PointContainer = typename Traits::PointContainer;
+    using IndexContainer = typename Traits::IndexContainer;
+
+    IndexType m_index;
+    const PointContainer* m_points;
+    const IndexContainer* m_samples;
+
+    KdTreeSampleIndexIterator(IndexType index, const PointContainer& points, const IndexContainer& samples)
+        : m_index(index), m_points(&points), m_samples(&samples) {}
+};
+
+}  // namespace Ponca
diff --git a/Ponca/src/SpatialPartitioning/KdTree/Query/kdTreeSampleIndexQuery.h b/Ponca/src/SpatialPartitioning/KdTree/Query/kdTreeSampleIndexQuery.h
new file mode 100644
index 000000000..d0cc77a44
--- /dev/null
+++ b/Ponca/src/SpatialPartitioning/KdTree/Query/kdTreeSampleIndexQuery.h
@@ -0,0 +1,44 @@
+/*
+ This Source Code Form is subject to the terms of the Mozilla Public
+ License, v. 2.0. If a copy of the MPL was not distributed with this
+ file, You can obtain one at http://mozilla.org/MPL/2.0/.
+*/
+
+#pragma once
+
+#include "../Iterator/kdTreeSampleIndexIterator.h"
+
+namespace Ponca {
+
+template <typename Traits>
+class KdTreeBase;
+
+template <typename Traits>
+class KdTreeSampleIndexQuery
+{
+public:
+    using IndexType = typename Traits::IndexType;
+    using Iterator  = KdTreeSampleIndexIterator<Traits>;
+
+    Iterator begin() {
+        return Iterator(0, *m_points, *m_samples);
+    }
+
+    Iterator end() {
+        return Iterator((IndexType)m_samples->size(), *m_points, *m_samples);
+    }
+
+private:
+    friend KdTreeBase<Traits>;
+
+    using PointContainer = typename Traits::PointContainer;
+    using IndexContainer = typename Traits::IndexContainer;
+
+    const PointContainer* m_points;
+    const IndexContainer* m_samples;
+
+    KdTreeSampleIndexQuery(const PointContainer& points, const IndexContainer& samples)
+        : m_points(&points), m_samples(&samples) {}
+};
+
+} // namespace Ponca
diff --git a/Ponca/src/SpatialPartitioning/KdTree/kdTree.h b/Ponca/src/SpatialPartitioning/KdTree/kdTree.h
index b25f8759a..c84da5b6a 100644
--- a/Ponca/src/SpatialPartitioning/KdTree/kdTree.h
+++ b/Ponca/src/SpatialPartitioning/KdTree/kdTree.h
@@ -21,6 +21,7 @@
 #include "Query/kdTreeNearestQueries.h"
 #include "Query/kdTreeKNearestQueries.h"
 #include "Query/kdTreeRangeQueries.h"
+#include "Query/kdTreeSampleIndexQuery.h"
 
 namespace Ponca {
 template <typename Traits> class KdTreeBase;
@@ -77,7 +78,7 @@ class KdTreeBase
 
     static_assert(std::is_same<typename PointContainer::value_type, DataPoint>::value,
         "PointContainer must contain DataPoints");
-    
+
     // Queries use a value of -1 for invalid indices
     static_assert(std::is_signed<IndexType>::value, "Index type must be signed");
     static_assert(std::is_same<typename IndexContainer::value_type, IndexType>::value, "Index type mismatch");
@@ -289,7 +290,11 @@ public :
     {
         return KdTreeRangeIndexQuery<Traits>(this, r, index);
     }
-    
+
+    KdTreeSampleIndexQuery<Traits> sample_point_indices() const
+    {
+        return KdTreeSampleIndexQuery<Traits>(m_points, m_indices);
+    }
 
     // Data --------------------------------------------------------------------
 protected: