wip adjacient nodes

src/orthtree.hpp

@@ -2,12 +2,20 @@
 #include <algorithm>
 #include <array>
 #include <bitset>
+#include <deque>
 #include <execution>
+#include <forward_list>
 #include <functional>
 #include <iostream>
+#include <list>
 #include <map>
 #include <memory>
 #include <mutex>
+#include <queue>
+#include <random>
+#include <shared_mutex>
+#include <tuple>
+#include <unordered_set>
 
 template <unsigned int dim, typename number, typename T> class Orthtree {
 public:
@@ -43,18 +51,26 @@ public:
       assert((min.array() < max.array()).all());
     }
   };
+  typedef std::bitset<dim> coord;
+  class Node;
+  struct NodeInfo {
+    const Node *node;
+    int depth;
+    number width;
+    coord crd;
+  };
+  typedef std::forward_list<NodeInfo> branch;
   class Node {
   public:
-    typedef std::bitset<dim> coord;
     Point center;
-    std::unique_ptr<std::array<Node, (2 << dim)>> children;
+    std::unique_ptr<std::array<Node, (1 << dim)>> children;
     std::unique_ptr<T> data;
 
-    mutable std::mutex mtx;
+    mutable std::shared_mutex mtx;
 
     void create_children(const Point &parent_center, number parent_width) {
-      children = std::make_unique<std::array<Node, (2 << dim)>>();
+      children = std::make_unique<std::array<Node, (1 << dim)>>();
-      for (unsigned int i = 0; i < (2 << dim); ++i) {
+      for (unsigned int i = 0; i < (1 << dim); ++i) {
         Point offset = Point::Ones() * (parent_width / 4);
         for (unsigned int j = 0; j < dim; ++j)
           if (!coord(i)[j])
@@ -63,39 +79,56 @@ public:
       }
     };
 
-    unsigned int get_child_coord(const Point &p) const {
+    coord get_child_coord(const Point &p) const {
       coord child_coord;
       for (unsigned int i = 0; i < dim; ++i)
         child_coord[i] = center[i] < p[i];
-      return child_coord.to_ulong();
+      return child_coord;
     };
 
     void depth(const Point &point, int depth, data_access_f f, number width) {
       if (depth == 0) {
-        std::lock_guard<std::mutex> guard(mtx);
+        std::unique_lock guard(mtx);
         f(data, center);
       } else {
+        mtx.lock_shared();
+        if (!children) {
+          mtx.unlock_shared();
           {
-          std::lock_guard<std::mutex> guard(mtx);
+            std::unique_lock uguard(mtx);
             if (!children)
               create_children(center, width);
           }
-        (*children)[get_child_coord(point)].depth(point, depth - 1, f,
+          mtx.lock_shared();
-                                                  width / 2);
+        }
+        (*children)[get_child_coord(point).to_ulong()].depth(point, depth - 1,
+                                                             f, width / 2);
+
+        mtx.unlock_shared();
       }
     };
+
     BBox bb(number width) const {
       Point offset = Point::Ones() * (width / 2);
       return BBox(center - offset, center + offset);
     };
+
     void access_data(cdata_access_f f) const {
-      std::lock_guard<std::mutex> guard(mtx);
+      std::unique_lock guard(mtx);
       if (data) {
         f(data, center);
       }
     }
+    bool access_data_ret(cdata_access_f f, bool &ret) const {
+      std::unique_lock guard(mtx);
+      if (data) {
+        ret = f(data, center);
+        return true;
+      }
+      return false;
+    }
     void all_data(cdata_access_f f) const {
-      std::lock_guard<std::mutex> guard(mtx);
+      std::unique_lock guard(mtx);
       if (data) {
         f(data, center);
       }
@@ -106,9 +139,19 @@ public:
       }
     }
     bool is_leaf() const {
-      std::lock_guard<std::mutex> guard(mtx);
+      std::shared_lock guard(mtx);
       return !bool(children);
     }
+    void nodes_to_leaf(const Point &p, const number width, const int depth,
+                       branch &l) const {
+      std::shared_lock guard(mtx);
+      if (!children)
+        return;
+      auto c = get_child_coord(p);
+      auto child = &(*children)[c.to_ullong()];
+      l.push_front({child, depth - 1, width / 2, c});
+      child->nodes_to_leaf(p, width / 2, depth - 1, l);
+    }
   };
   std::unique_ptr<Node> root;
   number root_width;
@@ -125,6 +168,122 @@ public:
     root->depth(point, depth, f, root_width);
   }
 
+  // std::vector<NodeInfo> const std::forward_list<std::pair<Node *, coord>>&
+  // branch)
+
+  branch get_branch_to_leaf(const Point &p) const {
+    branch b({root.get, 0, root_width, coord{0}});
+    root->node_to_leaf(p, root_width, b);
+    return b;
+  }
+
+  // std::vector<branch> get_leaf_neighours(branch b) {
+  //   auto node = b.front();
+  //   b.pop_front();
+  //   if (b.empty())
+  //     return {};
+  //   std::vector<branch> ret;
+  //   ret.reserve(2<<dim);
+  //   for (int i = 0; i < (1<<dim); i += 1) {
+  //     auto& c = *b.front()->children;
+  //     if (coord{i} != node->coord) {
+  //       ret.push_back(b);
+  //       ret.back().push_front(&c);
+  //     }
+  //   }
+  //   return ret;
+  // }
+  template<typename Container>
+  void get_nodes_in_dir(const branch b, int dir,
+                        std::queue<branch,Container> &nodes) const {
+    nodes.push_back(b);
+    const unsigned int dim_dir = std::abs(dir) - 1;
+    if (!b.front().node->children)
+      return;
+    for (int i = 0; i < 1 << dim; i++) {
+      if (coord{i}[dim_dir] == dir > 0)
+        continue;
+      branch child_branch(b.begin(), b.end());
+      child_branch.push_front({b.front().node->children->at(i),
+                               b.front().depth - 1, b.front().width / 2,
+                               coord{i}});
+      get_nodes_in_dir(child_branch, dir, nodes);
+    }
+  }
+
+  template<typename Container>
+  void get_adjacent_nodes(branch b, int dir,
+                          std::queue<branch,Container> &ad_nodes) const {
+    auto n = b.first();
+    const unsigned int dim_dir = std::abs(dir) - 1;
+    int up = 1;
+    auto b_it = b.begin();
+    if (std::next(b_it) == b.end())
+      return;
+    for (; b_it.crd[dim_dir] == dir > 0; std::advance(b_it, 1)) {
+      if (std::next(b_it) == b.end())
+        return;
+      up++;
+    }
+    branch c(b.begin(), std::next(b_it));
+    branch p(std::next(b_it), b.end()); // TODO remove from b not cpy
+    c.reverse();
+    for (int i = 0; i < up; i += 1) {
+      coord pc = coord(c.front().crd).flip(dim_dir);
+      p.push_front({c.front().node->children->at(pc), c.front().depth - i,
+                    c.front().width / (1 << i), pc});
+      c.pop_front();
+    }
+    get_nodes_in_dir(p, -dir, ad_nodes);
+  }
+
+  void datas_by_distance(const Point &p, cdata_access_f f) const {
+    branch b = get_branch_to_leaf(p);
+    std::deque<branch> to_visit{b};
+    std::unordered_set<const Node *> visited;
+    while (!to_visit.empty()) {
+      const branch& curr_visit = to_visit.front();
+      for (const auto &n : curr_visit) {
+        if (visited.contains(n.node))
+          break;
+        bool ret;
+        if (n.node->data_acces_ret(f, ret) && !ret)
+          return;
+        visited.insert(n.node);
+      }
+      to_visit.pop_front();
+      for (int i = 0; i < dim; i += 1) {
+        get_adjacent_nodes(curr_visit,dim+1,to_visit);//to_visit.push_back({p1, get_branch_to_leaf(p1)});
+        get_adjacent_nodes(curr_visit,-dim-1,to_visit);//to_visit.push_back({p1, get_branch_to_leaf(p1)});
+      }
+    }
+  }
+  // void datas_by_distance(const Point &p, cdata_access_f f) const {
+  //   branch b = get_branch_to_leaf(p);
+  //   std::deque<std::pair<Point, branch>> to_visit{{p, b}};
+  //   std::unordered_set<Node *> visited;
+  //   while (!to_visit.empty()) {
+  //     for (auto &n : to_visit.second.front()) {
+  //       if (visited.contains(n))
+  //         break;
+  //       bool ret;
+  //       if (n.node->data_acces_ret(f, ret) && !ret)
+  //         return;
+  //       visited.insert(n.node);
+  //     }
+  //     auto search_p = to_visit.front().first;
+  //     to_visit.pop_front();
+  //     for (int i = 0; i < dim; i += 1) {
+  //       Point offset = Point::Zero();
+  //       offset(i) = b.width;
+  //       auto p1 = search_p + offset;
+  //       auto p2 = search_p + offset;
+  //       to_visit.push_back({p1, get_branch_to_leaf(p1)});
+  //       to_visit.push_back({p2, get_branch_to_leaf(p2)});
+  //     }
+  //   }
+  // }
+
   void data_in_bb(const BBox &bb, cdata_access_f f) const {
     number curr_width = root_width;
     std::vector<const Node *> to_visit{root.get()};
@@ -135,17 +294,16 @@ public:
       auto bbrel = cbb.relation(bb);
       // std::cout << to_visit.back()->center.transpose() << " - "
       // << cbb.min.transpose() << " - " << bbrel << std::endl;
+      bool ret;
       switch (bbrel) {
       case (Orthtree::BBox::CONTAINED):
         to_visit.back()->all_data(f);
         break;
       case (Orthtree::BBox::OVERLAP):
       case Orthtree::BBox::CONTAINS:
-        to_visit.back()->access_data([&](auto &p, auto c) {
+        to_visit.back()->access_data_ret(f, ret);
-          if (!f(p, c))
+        if (to_visit.back()->access_data_ret(f, ret) && !ret)
           stop = true;
-          return false;
-        });
         if (!to_visit.back()->is_leaf()) {
           // std::transform(std::execution::par_unseq,to_visit.back()->children->cbegin(),
           //                to_visit.back()->children->cend(),

test/orthtree_test.cpp

@@ -80,6 +80,27 @@ TEST(orthtree, insert2) {
   EXPECT_EQ(i, 2);
 }
 
+
+TEST(orthtree, dist_search) {
+  typedef Orthtree<2, float, Eigen::Vector2f> O;
+  O t(O::Point::Zero(), 2<<4);
+
+  for (int i = 0; i < 4; i += 1) {
+    O::Point p = O::Point(-14 + i*5, 6);
+    t.depth(p, 4, [&](auto &tp, [[maybe_unused]] auto c) {
+      tp = std::make_unique<O::Point>(p);
+      return true;
+    });
+  }
+  t.data_in_bb(O::BBox(O::Point(-2, 0), O::Point(0, 32)),
+               [&](auto &tp, auto c) -> bool {
+                 EXPECT_TRUE(tp) << c.transpose();
+                 i += 1;
+                 return true;
+               });
+  EXPECT_EQ(i, 2);
+}
+
 int main(int argc, char **argv) {
   testing::InitGoogleTest(&argc, argv);
   return RUN_ALL_TESTS();