Refactor: Convert mjCOctree::MakeOctree from recursive to iterative.

The octree construction is now performed using a breadth-first search with a `std::deque` to manage tasks. A new `OctreeTask` struct encapsulates the state for each node to be processed, and a helper function `Subdivide` is introduced for subdividing the Octree nodes.

Also reduce the number of (center, size) to (min, max) conversion for the bounding box representation in order to avoid duplicated vertices in the Octree due to floating point errors.

PiperOrigin-RevId: 806214441
Change-Id: Iff9015ebe4705522db743c478cef5acf1a7312b7

src/user/user_objects.cc

@@ -21,6 +21,7 @@
 #include <cstdint>
 #include <cstdio>
 #include <cstring>
+#include <deque>
 #include <functional>
 #include <limits>
 #include <memory>
@@ -573,9 +574,12 @@ void mjCOctree::CopyChild(int* child) const {
 
 void mjCOctree::CopyAabb(mjtNum* aabb) const {
   for (int i = 0; i < node_.size(); ++i) {
-    for (int j = 0; j < 6; ++j) {
+    aabb[i * 6 + 0] = (node_[i].aamm[0] + node_[i].aamm[3]) / 2;
-      aabb[i * 6 + j] = node_[i].aabb[j];
+    aabb[i * 6 + 1] = (node_[i].aamm[1] + node_[i].aamm[4]) / 2;
-    }
+    aabb[i * 6 + 2] = (node_[i].aamm[2] + node_[i].aamm[5]) / 2;
+    aabb[i * 6 + 3] = (node_[i].aamm[3] - node_[i].aamm[0]) / 2;
+    aabb[i * 6 + 4] = (node_[i].aamm[4] - node_[i].aamm[1]) / 2;
+    aabb[i * 6 + 5] = (node_[i].aamm[5] - node_[i].aamm[2]) / 2;
   }
 }
 
@@ -624,7 +628,7 @@ void mjCOctree::CreateOctree(const double aamm[6]) {
   std::transform(elements.begin(), elements.end(), elements_ptrs.begin(),
                  [](Triangle& triangle) { return &triangle; });
   std::unordered_map<Point, int> vert_map;
-  MakeOctree(elements_ptrs, box, 0, vert_map);
+  MakeOctree(elements_ptrs, box, vert_map);
 }
 
 
@@ -685,23 +689,25 @@ static bool boxTriangle(const Triangle& v, const double aamm[6]) {
 }
 
 
-int mjCOctree::MakeOctree(const std::vector<Triangle*>& elements, const double aamm[6], int lev,
+void mjCOctree::TaskToNode(const OctreeTask& task, OctNode& node,
-                          std::unordered_map<Point, int>& vert_map) {
+                           std::unordered_map<Point, int>& vert_map) {
-  node_.push_back(OctNode());
+  node.level = task.lev;
-  OctNode& node = node_.back();
+
-  node.level = lev;
+  if (task.parent_index != -1) {
+    node_[task.parent_index].child[task.child_slot] = task.node_index;
+    const auto parent_aamm = node_[task.parent_index].aamm;
+    node.aamm[0] = task.child_slot & 1 ? parent_aamm[0] : (parent_aamm[3] + parent_aamm[0]) / 2;
+    node.aamm[1] = task.child_slot & 2 ? parent_aamm[1] : (parent_aamm[4] + parent_aamm[1]) / 2;
+    node.aamm[2] = task.child_slot & 4 ? parent_aamm[2] : (parent_aamm[5] + parent_aamm[2]) / 2;
+    node.aamm[3] = task.child_slot & 1 ? (parent_aamm[0] + parent_aamm[3]) / 2 : parent_aamm[3];
+    node.aamm[4] = task.child_slot & 2 ? (parent_aamm[1] + parent_aamm[4]) / 2 : parent_aamm[4];
+    node.aamm[5] = task.child_slot & 4 ? (parent_aamm[2] + parent_aamm[5]) / 2 : parent_aamm[5];
+  }
 
-  // create a new node
-  int index = nnode_++;
-  std::array<double, 6> aabb = {
-      (aamm[0] + aamm[3]) / 2, (aamm[1] + aamm[4]) / 2,
-      (aamm[2] + aamm[5]) / 2, (aamm[3] - aamm[0]) / 2,
-      (aamm[4] - aamm[1]) / 2, (aamm[5] - aamm[2]) / 2};
-  node.aabb = aabb;
   for (int i = 0; i < 8; i++) {
-    Point v = {{(i & 1) ? aamm[3] : aamm[0],
+    Point v = {{(i & 1) ? node.aamm[3] : node.aamm[0],
-                (i & 2) ? aamm[4] : aamm[1],
+                (i & 2) ? node.aamm[4] : node.aamm[1],
-                (i & 4) ? aamm[5] : aamm[2]}};
+                (i & 4) ? node.aamm[5] : node.aamm[2]}};
     auto it = vert_map.find(v);
     if (it != vert_map.end()) {
       node.vertid[i] = it->second;
@@ -712,37 +718,63 @@ int mjCOctree::MakeOctree(const std::vector<Triangle*>& elements, const double a
     }
     node.child[i] = -1;
   }
+}
 
-  // find all triangles that intersect the current box
+
-  std::vector<Triangle*> colliding;
+void mjCOctree::Subdivide(std::deque<OctreeTask>& queue, const std::vector<Triangle*>& colliding,
-  for (auto* element : elements) {
+                          const OctreeTask& task, std::unordered_map<Point, int>& vert_map) {
-    if (boxTriangle(*element, aamm)) {
+  for (int i = 0; i < 8; i++) {
-      colliding.push_back(element);
+    OctreeTask new_task;
+    new_task.elements = colliding;
+    new_task.lev = task.lev + 1;
+    new_task.parent_index = task.node_index;
+    new_task.node_index = nnode_++;
+    new_task.child_slot = i;
+
+    node_.push_back(OctNode());
+    TaskToNode(new_task, node_.back(), vert_map);
+    queue.push_back(std::move(new_task));
+  }
+}
+
+
+void mjCOctree::MakeOctree(const std::vector<Triangle*>& elements, const double aamm[6],
+                           std::unordered_map<Point, int>& vert_map) {
+  std::deque<OctreeTask> queue;
+  OctreeTask initial_task;
+  initial_task.elements = elements;
+  initial_task.lev = 0;
+  initial_task.parent_index = -1;
+  initial_task.child_slot = -1;
+  initial_task.node_index = nnode_++;
+  queue.push_back(std::move(initial_task));
+
+  // create root node
+  node_.push_back(OctNode());
+  OctNode& root = node_.back();
+  std::copy(aamm, aamm + 6, root.aamm.data());
+  TaskToNode(queue.front(), node_.back(), vert_map);
+
+  while (!queue.empty()) {
+    OctreeTask task = std::move(queue.front());
+    queue.pop_front();
+
+    // find all triangles that intersect the current box
+    std::vector<Triangle*> colliding;
+    for (auto* element : task.elements) {
+      if (boxTriangle(*element, node_[task.node_index].aamm.data())) {
+        colliding.push_back(element);
+      }
     }
-  }
 
-  // return if the box is empty
+    // skip if the box is empty
-  if (colliding.empty() || lev >= 6) {
+    if (colliding.empty() || task.lev >= 6) {
-    return index;
+      continue;
-  }
+    }
 
-  // split the box into 8 sub-boxes
+    // subdivide the node
-  double new_aamm[8][6];
+    Subdivide(queue, colliding, task, vert_map);
-  for (int i = 0; i < 8; i++) {
-    new_aamm[i][0] = aabb[0] + aabb[3] * (i & 1 ? -1 : 0);
-    new_aamm[i][1] = aabb[1] + aabb[4] * (i & 2 ? -1 : 0);
-    new_aamm[i][2] = aabb[2] + aabb[5] * (i & 4 ? -1 : 0);
-    new_aamm[i][3] = aabb[0] + aabb[3] * (i & 1 ? 0 : 1);
-    new_aamm[i][4] = aabb[1] + aabb[4] * (i & 2 ? 0 : 1);
-    new_aamm[i][5] = aabb[2] + aabb[5] * (i & 4 ? 0 : 1);
   }
-
-  // recursive calls to create sub-boxes
-  for (int i = 0; i < 8; i++) {
-    node_[index].child[i] = MakeOctree(colliding, new_aamm[i], lev + 1, vert_map);
-  }
-
-  return index;
 }
 
 //------------------------- class mjCDef implementation --------------------------------------------

src/user/user_objects.h

@@ -20,6 +20,7 @@
 #include <cmath>
 #include <cstddef>
 #include <array>
+#include <deque>
 #include <functional>
 #include <string>
 #include <string_view>
@@ -251,10 +252,18 @@ struct OctNode {
   int level = 0;                       // level of the node
   std::array<int, 8> child = {-1};     // children nodes
   std::array<int, 8> vertid = {-1};    // vertex id's
-  std::array<double, 6> aabb = {0};    // bounding box
+  std::array<double, 6> aamm = {0};    // bounding box
   std::array<double, 8> coeff = {0};   // interpolation coefficients
 };
 
+struct OctreeTask {
+  std::vector<Triangle*> elements;
+  int lev;
+  int parent_index;
+  int child_slot;
+  int node_index;
+};
+
 struct mjCOctree_ {
   int nnode_ = 0;
   int nvert_ = 0;
@@ -290,8 +299,11 @@ class mjCOctree : public mjCOctree_ {
 
  private:
   void Make(std::vector<Triangle>& elements);
-  int MakeOctree(const std::vector<Triangle*>& elements, const double aamm[6], int lev,
+  void MakeOctree(const std::vector<Triangle*>& elements, const double aamm[6],
-                 std::unordered_map<Point, int>& vert_map);
+                  std::unordered_map<Point, int>& vert_map);
+  void TaskToNode(const OctreeTask& task, OctNode& node, std::unordered_map<Point, int>& vert_map);
+  void Subdivide(std::deque<OctreeTask>& queue, const std::vector<Triangle*>& colliding,
+                 const OctreeTask& task, std::unordered_map<Point, int>& vert_map);
 };