From 785878f062aac4e52b791ad726504afddcfc26c2 Mon Sep 17 00:00:00 2001
From: Pier Angelo Vendrame <PieroV@users.noreply.github.com>
Date: Thu, 18 Apr 2024 01:03:47 +0200
Subject: [PATCH] Simplify data_ and data_interface_ in KDTreeFlann. (#6734)

KDTreeFlann::SetRawData had some confusion around data, data_, and
data_interface_. As a result, the data from the parameter was copied to
the member std::vector, but the interface (which takes an Eigen::Map)
used the data passed as an argument to the method.
As a result, the searches risked to be run on a dangling pointer,
instead of the intended internal storage.

However, rather than adjusting the pointer for Eigen::Map, I preferred
simplifying the code, and copy the argument to an Eigen::MatrixXd.
---
 CHANGELOG.md                        |  2 ++
 cpp/open3d/geometry/KDTreeFlann.cpp | 21 ++++++---------------
 cpp/open3d/geometry/KDTreeFlann.h   | 14 +++++---------
 3 files changed, 13 insertions(+), 24 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index af429044ee2..b68b0bdbd02 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -1,4 +1,5 @@
 ## Main
+
 -   Fix TriangleMesh::SamplePointsUniformly not sampling triangle meshes uniformly (PR #6653)
 -   Fix tensor based TSDF integration example.
 -   Use GLIBCXX_USE_CXX11_ABI=ON by default
@@ -33,6 +34,7 @@
 -   Add Python pathlib support for file IO (PR #6619)
 -   Fix log error message for `probability` argument validation in `PointCloud::SegmentPlane` (PR #6622)
 -   Fix macOS arm64 builds, add CI runner for macOS arm64 (PR #6695)
+-   Fix KDTreeFlann possibly using a dangling pointer instead of internal storage and simplified its members (PR #6734)
 
 ## 0.13
 
diff --git a/cpp/open3d/geometry/KDTreeFlann.cpp b/cpp/open3d/geometry/KDTreeFlann.cpp
index d93175a8d43..c81814a724d 100644
--- a/cpp/open3d/geometry/KDTreeFlann.cpp
+++ b/cpp/open3d/geometry/KDTreeFlann.cpp
@@ -97,8 +97,7 @@ int KDTreeFlann::SearchKNN(const T &query,
     // This is optimized code for heavily repeated search.
     // Other flann::Index::knnSearch() implementations lose performance due to
     // memory allocation/deallocation.
-    if (data_.empty() || dataset_size_ <= 0 ||
-        size_t(query.rows()) != dimension_ || knn < 0) {
+    if (data_.size() == 0 || query.rows() != data_.rows() || knn < 0) {
         return -1;
     }
     indices.resize(knn);
@@ -121,8 +120,7 @@ int KDTreeFlann::SearchRadius(const T &query,
     // Since max_nn is not given, we let flann to do its own memory management.
     // Other flann::Index::radiusSearch() implementations lose performance due
     // to memory management and CPU caching.
-    if (data_.empty() || dataset_size_ <= 0 ||
-        size_t(query.rows()) != dimension_) {
+    if (data_.size() == 0 || query.rows() != data_.rows()) {
         return -1;
     }
     std::vector<nanoflann::ResultItem<Eigen::Index, double>> indices_dists;
@@ -148,8 +146,7 @@ int KDTreeFlann::SearchHybrid(const T &query,
     // It is also the recommended setting for search.
     // Other flann::Index::radiusSearch() implementations lose performance due
     // to memory allocation/deallocation.
-    if (data_.empty() || dataset_size_ <= 0 ||
-        size_t(query.rows()) != dimension_ || max_nn < 0) {
+    if (data_.size() == 0 || query.rows() != data_.rows() || max_nn < 0) {
         return -1;
     }
     distance2.resize(max_nn);
@@ -166,18 +163,12 @@ int KDTreeFlann::SearchHybrid(const T &query,
 }
 
 bool KDTreeFlann::SetRawData(const Eigen::Map<const Eigen::MatrixXd> &data) {
-    dimension_ = data.rows();
-    dataset_size_ = data.cols();
-    if (dimension_ == 0 || dataset_size_ == 0) {
+    if (data.size() == 0) {
         utility::LogWarning("[KDTreeFlann::SetRawData] Failed due to no data.");
         return false;
     }
-    data_.resize(dataset_size_ * dimension_);
-    memcpy(data_.data(), data.data(),
-           dataset_size_ * dimension_ * sizeof(double));
-    data_interface_.reset(new Eigen::Map<const Eigen::MatrixXd>(data));
-    nanoflann_index_.reset(
-            new KDTree_t(dimension_, std::cref(*data_interface_), 15));
+    data_ = data;
+    nanoflann_index_ = std::make_unique<KDTree_t>(data_.rows(), data_, 15);
     nanoflann_index_->index_->buildIndex();
     return true;
 }
diff --git a/cpp/open3d/geometry/KDTreeFlann.h b/cpp/open3d/geometry/KDTreeFlann.h
index 3d3ab2c5bea..cb41fccafb7 100644
--- a/cpp/open3d/geometry/KDTreeFlann.h
+++ b/cpp/open3d/geometry/KDTreeFlann.h
@@ -97,17 +97,13 @@ class KDTreeFlann {
     bool SetRawData(const Eigen::Map<const Eigen::MatrixXd> &data);
 
 protected:
-    using KDTree_t = nanoflann::KDTreeEigenMatrixAdaptor<
-            Eigen::Map<const Eigen::MatrixXd>,
-            -1,
-            nanoflann::metric_L2,
-            false>;
+    using KDTree_t = nanoflann::KDTreeEigenMatrixAdaptor<const Eigen::MatrixXd,
+                                                         -1,
+                                                         nanoflann::metric_L2,
+                                                         false>;
 
-    std::vector<double> data_;
-    std::unique_ptr<Eigen::Map<const Eigen::MatrixXd>> data_interface_;
+    Eigen::MatrixXd data_;
     std::unique_ptr<KDTree_t> nanoflann_index_;
-    size_t dimension_ = 0;
-    size_t dataset_size_ = 0;
 };
 
 }  // namespace geometry