[SpatialPartitioning] KdTreeQueries: Move implementation to headers

poncateam · Jul 26, 2023 · 557aa8c · 557aa8c
1 parent 360e4ea
commit 557aa8c
Show file tree

Hide file tree

Showing 13 changed files with 176 additions and 305 deletions.
diff --git a/Ponca/src/SpatialPartitioning/KdTree/Query/kdTreeKNearestIndexQuery.h b/Ponca/src/SpatialPartitioning/KdTree/Query/kdTreeKNearestIndexQuery.h
@@ -30,12 +30,24 @@ class KdTreeKNearestIndexQuery : public KdTreeQuery<Traits>,
     }
 
 public:
-    KdTreeKNearestIterator<IndexType, DataPoint> begin();
-    KdTreeKNearestIterator<IndexType, DataPoint> end();
+    KdTreeKNearestIterator<IndexType, DataPoint> begin(){
+        QueryAccelType::reset();
+        QueryType::reset();
+        this->search();
+        return KdTreeKNearestIterator<IndexType, DataPoint>(QueryType::m_queue.begin());
+    }
+    KdTreeKNearestIterator<IndexType, DataPoint> end(){
+        return KdTreeKNearestIterator<IndexType, DataPoint>(QueryType::m_queue.end());
+    }
 
 protected:
-    void search();
+    void search(){
+        KdTreeQuery<Traits>::search_internal(QueryAccelType::m_kdtree->point_data()[QueryType::input()].pos(),
+                                             [](IndexType, IndexType){},
+                                             [this](){return QueryType::descentDistanceThreshold();},
+                                             [this](IndexType idx){return QueryType::input() == idx;},
+                                             [this](IndexType idx, IndexType, Scalar d){QueryType::m_queue.push({idx, d}); return false;}
+        );
+    }
 };
-
-#include "./kdTreeKNearestIndexQuery.hpp"
 } // namespace ponca
diff --git a/Ponca/src/SpatialPartitioning/KdTree/Query/kdTreeKNearestIndexQuery.hpp b/Ponca/src/SpatialPartitioning/KdTree/Query/kdTreeKNearestIndexQuery.hpp
diff --git a/Ponca/src/SpatialPartitioning/KdTree/Query/kdTreeKNearestPointQuery.h b/Ponca/src/SpatialPartitioning/KdTree/Query/kdTreeKNearestPointQuery.h
@@ -30,13 +30,25 @@ class KdTreeKNearestPointQuery : public KdTreeQuery<Traits>,
     }
 
 public:
-    KdTreeKNearestIterator<IndexType, DataPoint> begin();
-    KdTreeKNearestIterator<IndexType, DataPoint> end();
+    KdTreeKNearestIterator<IndexType, DataPoint> begin(){
+        QueryAccelType::reset();
+        QueryType::reset();
+        this->search();
+        return KdTreeKNearestIterator<IndexType, DataPoint>(QueryType::m_queue.begin());
+    }
+    KdTreeKNearestIterator<IndexType, DataPoint> end(){
+        return KdTreeKNearestIterator<IndexType, DataPoint>(QueryType::m_queue.end());
+    }
 
 protected:
-   void search();
+   void search(){
+       KdTreeQuery<Traits>::search_internal(QueryType::input(),
+                                            [](IndexType, IndexType){},
+                                            [this](){return QueryType::descentDistanceThreshold();},
+                                            [](IndexType){return false;},
+                                            [this](IndexType idx, IndexType, Scalar d){QueryType::m_queue.push({idx, d}); return false;}
+       );
+   }
 };
-
-#include "./kdTreeKNearestPointQuery.hpp"
 } // namespace Ponca
 
diff --git a/Ponca/src/SpatialPartitioning/KdTree/Query/kdTreeKNearestPointQuery.hpp b/Ponca/src/SpatialPartitioning/KdTree/Query/kdTreeKNearestPointQuery.hpp
diff --git a/Ponca/src/SpatialPartitioning/KdTree/Query/kdTreeNearestIndexQuery.h b/Ponca/src/SpatialPartitioning/KdTree/Query/kdTreeNearestIndexQuery.h
@@ -30,12 +30,29 @@ class KdTreeNearestIndexQuery : public KdTreeQuery<Traits>,
     }
 
 public:
-    KdTreeNearestIterator<IndexType> begin();
-    KdTreeNearestIterator<IndexType> end();
+    KdTreeNearestIterator<IndexType> begin(){
+        QueryAccelType::reset();
+        QueryType::reset();
+        this->search();
+        return KdTreeNearestIterator<IndexType>(QueryType::m_nearest);
+    }
+    KdTreeNearestIterator<IndexType> end(){
+        return KdTreeNearestIterator<IndexType>(QueryType::m_nearest + 1);
+    }
 
 protected:
-    void search();
+    void search(){
+        KdTreeQuery<Traits>::search_internal(QueryAccelType::m_kdtree->point_data()[QueryType::input()].pos(),
+                                             [](IndexType, IndexType){},
+                                             [this](){return QueryType::descentDistanceThreshold();},
+                                             [this](IndexType idx){return QueryType::input() == idx;},
+                                             [this](IndexType idx, IndexType, Scalar d)
+                                             {
+                                                 QueryType::m_nearest = idx;
+                                                 QueryType::m_squared_distance = d;
+                                                 return false;
+                                             }
+        );
+    }
 };
-
-#include "./kdTreeNearestIndexQuery.hpp"
 } // namespace ponca
diff --git a/Ponca/src/SpatialPartitioning/KdTree/Query/kdTreeNearestIndexQuery.hpp b/Ponca/src/SpatialPartitioning/KdTree/Query/kdTreeNearestIndexQuery.hpp
diff --git a/Ponca/src/SpatialPartitioning/KdTree/Query/kdTreeNearestPointQuery.h b/Ponca/src/SpatialPartitioning/KdTree/Query/kdTreeNearestPointQuery.h
@@ -30,12 +30,29 @@ class KdTreeNearestPointQuery : public KdTreeQuery<Traits>,
     }
 
 public:
-    KdTreeNearestIterator<IndexType> begin();
-    KdTreeNearestIterator<IndexType> end();
+    KdTreeNearestIterator<IndexType> begin(){
+        QueryAccelType::reset();
+        QueryType::reset();
+        this->search();
+        return KdTreeNearestIterator<IndexType>(QueryType::m_nearest);
+    }
+    KdTreeNearestIterator<IndexType> end(){
+        return KdTreeNearestIterator<IndexType>(QueryType::m_nearest + 1);
+    }
 
 protected:
-    void search();
+    void search(){
+        KdTreeQuery<Traits>::search_internal(QueryType::input(),
+                                             [](IndexType, IndexType){},
+                                             [this](){return QueryType::descentDistanceThreshold();},
+                                             [](IndexType){return false;},
+                                             [this](IndexType idx, IndexType, Scalar d)
+                                             {
+                                                 QueryType::m_nearest = idx;
+                                                 QueryType::m_squared_distance = d;
+                                                 return false;
+                                             }
+        );
+    }
 };
-
-#include "./kdTreeNearestPointQuery.hpp"
 } // namespace ponca
diff --git a/Ponca/src/SpatialPartitioning/KdTree/Query/kdTreeNearestPointQuery.hpp b/Ponca/src/SpatialPartitioning/KdTree/Query/kdTreeNearestPointQuery.hpp
diff --git a/Ponca/src/SpatialPartitioning/KdTree/Query/kdTreeRangeIndexQuery.h b/Ponca/src/SpatialPartitioning/KdTree/Query/kdTreeRangeIndexQuery.h
@@ -36,13 +36,56 @@ class KdTreeRangeIndexQuery : public KdTreeQuery<Traits>,
     }
 
 public:
-    inline Iterator begin();
-    inline Iterator end();
+    inline Iterator begin(){
+        QueryAccelType::reset();
+        QueryType::reset();
+        Iterator it(this);
+        this->advance(it);
+        return it;
+    }
+    inline Iterator end(){
+        return Iterator(this, QueryAccelType::m_kdtree->point_count());
+    }
 
 protected:
-    inline void initialize();
-    inline void advance(Iterator& iterator);
-};
+    inline void advance(Iterator& it){
+        const auto& points  = QueryAccelType::m_kdtree->point_data();
+        const auto& indices = QueryAccelType::m_kdtree->index_data();
+        const auto& point   = points[QueryType::input()].pos();
+
+        if (points.empty() || indices.empty())
+            throw std::invalid_argument("Empty KdTree");
 
-#include "./kdTreeRangeIndexQuery.hpp"
+        for(IndexType i=it.m_start; i<it.m_end; ++i)
+        {
+            IndexType idx = indices[i];
+            if(idx == QueryType::input()) continue;
+
+            Scalar d = (point - points[idx].pos()).squaredNorm();
+            if(d < QueryType::descentDistanceThreshold())
+            {
+                it.m_index = idx;
+                it.m_start = i+1;
+                return;
+            }
+        }
+
+        if (KdTreeQuery<Traits>::search_internal(point,
+                                                 [&it](IndexType start, IndexType end)
+                                                 {
+                                                     it.m_start = start;
+                                                     it.m_end   = end;
+                                                 },
+                                                 [this](){return QueryType::descentDistanceThreshold();},
+                                                 [this](IndexType idx){return QueryType::input() == idx;},
+                                                 [&it](IndexType idx, IndexType i, Scalar)
+                                                 {
+                                                     it.m_index = idx;
+                                                     it.m_start = i+1;
+                                                     return true;
+                                                 }
+        ))
+            it.m_index = static_cast<IndexType>(points.size());
+    }
+};
 } // namespace ponca