From 3965fe357bbdbf0c85e019aeac32cae3c85ac9e4 Mon Sep 17 00:00:00 2001
From: Robert Osfield <robert@openscenegraph.com>
Date: Sun, 6 Jul 2008 12:14:19 +0000
Subject: [PATCH] Moved KdTree build code into osg::KdTree

---
 examples/osgkdtree/variabledivision.h |   1 -
 include/osg/KdTree                    | 129 ++++++++++--
 src/osg/KdTree.cpp                    | 281 +++++++++++++++++++++++++-
 3 files changed, 387 insertions(+), 24 deletions(-)

diff --git a/examples/osgkdtree/variabledivision.h b/examples/osgkdtree/variabledivision.h
index 72eccf48b..2470a7b67 100644
--- a/examples/osgkdtree/variabledivision.h
+++ b/examples/osgkdtree/variabledivision.h
@@ -118,7 +118,6 @@ class KDTree : public osg::Shape
     
         typedef std::vector< unsigned int > AxisStack;
         typedef std::vector< KDNode >       KDNodeList;
-
         typedef std::vector< KDLeaf > KDLeafList;
 
         /// note, leafNum is negative to distinguish from nodeNum
diff --git a/include/osg/KdTree b/include/osg/KdTree
index c0e26b820..eaff4c808 100644
--- a/include/osg/KdTree
+++ b/include/osg/KdTree
@@ -26,17 +26,28 @@ class OSG_EXPORT KdTree : public osg::Shape
     public:
     
         
-        KdTree() {}
+        KdTree();
         
-        KdTree(const KdTree& rhs, const osg::CopyOp& copyop=osg::CopyOp::SHALLOW_COPY):
-            Shape(rhs,copyop) {}
+        KdTree(const KdTree& rhs, const osg::CopyOp& copyop=osg::CopyOp::SHALLOW_COPY);
 
         META_Shape(osg, KdTree)
         
+        struct BuildOptions
+        {
+            BuildOptions():
+                _numVerticesProcessed(0),
+                _targetNumTrianglesPerLeaf(4),
+                _maxNumLevels(24) {}
+                        
+            unsigned int _numVerticesProcessed;
+            unsigned int _targetNumTrianglesPerLeaf;
+            unsigned int _maxNumLevels;
+        };
+        
+        
         /** Build the kdtree from the specified source geometry object.
           * retun true on success. */ 
-        virtual bool build(osg::Geometry* geometry);
-        
+        virtual bool build(BuildOptions& buildOptions, osg::Geometry* geometry);
         
         struct LineSegmentIntersection
         {
@@ -68,13 +79,13 @@ class OSG_EXPORT KdTree : public osg::Shape
         typedef int value_type;
         typedef std::vector< value_type >   Indices;
 
-        struct KDNode
+        struct KdNode
         {
-            KDNode():
+            KdNode():
                 first(0),
                 second(0) {}
 
-            KDNode(value_type f, value_type s):
+            KdNode(value_type f, value_type s):
                 first(f),
                 second(s) {}
 
@@ -85,13 +96,13 @@ class OSG_EXPORT KdTree : public osg::Shape
         };
 
 
-        struct KDLeaf
+        struct KdLeaf
         {
-            KDLeaf():
+            KdLeaf():
                 first(0),
                 second(0) {}
 
-            KDLeaf(value_type f, value_type s):
+            KdLeaf(value_type f, value_type s):
                 first(f),
                 second(s) {}
 
@@ -120,6 +131,93 @@ class OSG_EXPORT KdTree : public osg::Shape
             unsigned int _p3;    
         };
 
+        typedef std::vector< unsigned int > AxisStack;
+        typedef std::vector< KdNode >       KDNodeList;
+        typedef std::vector< KdLeaf > KDLeafList;
+
+        /// note, leafNum is negative to distinguish from nodeNum
+        int addLeaf(const KdLeaf& leaf) { int num = _kdLeaves.size(); _kdLeaves.push_back(leaf); return -(num+1); }
+
+        int replaceLeaf(int leafNum, const KdLeaf& leaf)
+        {
+            int num = -leafNum-1; 
+            
+            if (num>_kdLeaves.size()-1)
+            {
+                osg::notify(osg::NOTICE)<<"Warning: replaceChild("<<leafNum<<", leaf), num = "<<num<<" _kdLeaves.size()="<<_kdLeaves.size()<<std::endl;
+                return leafNum;
+            }
+            
+            _kdLeaves[num] = leaf; return leafNum;
+        }
+
+        /// note, leafNum is negative to distinguish from nodeNum
+        KdLeaf& getLeaf(int leafNum)
+        {
+            int num = -leafNum-1;
+            if (num<0 || num>_kdLeaves.size()-1)
+            {
+                osg::notify(osg::NOTICE)<<"Warning: getLeaf("<<leafNum<<", num = "<<num<<") _kdLeaves.size()="<<_kdLeaves.size()<<std::endl;
+            }
+
+            return _kdLeaves[num];
+        }
+
+        int addNode(const KdNode& node)
+        {
+            int num = _kdNodes.size(); 
+            _kdNodes.push_back(node); 
+            return num;
+        }
+
+        /// note, nodeNum is positive to distinguish from leftNum
+        KdNode& getNode(int nodeNum)
+        {
+            if (nodeNum<0 || nodeNum>_kdNodes.size()-1)
+            {
+                osg::notify(osg::NOTICE)<<"Warning: getNode("<<nodeNum<<") _kdNodes.size()="<<_kdNodes.size()<<std::endl;
+            }
+            return _kdNodes[nodeNum];
+        }
+        
+        osg::BoundingBox& getBounindingBox(int nodeNum)
+        {
+            if (nodeNum<0)
+            {
+                int num = -nodeNum-1;
+                return _kdLeaves[num].bb;            
+            }
+            else
+            {
+                return _kdNodes[nodeNum].bb;
+            }
+        }
+
+
+        void computeDivisions(BuildOptions& options);
+
+        int divide(BuildOptions& options, osg::BoundingBox& bb, int nodeIndex, unsigned int level);
+
+
+        typedef std::vector< osg::BoundingBox >                 BoundingBoxList;
+        typedef std::vector< Triangle >                         TriangleList;
+        typedef std::vector< osg::Vec3 >                        CenterList;
+
+        osg::observer_ptr<osg::Geometry>    _geometry;
+
+        osg::BoundingBox                    _bb;
+
+        AxisStack                           _axisStack;
+        KDNodeList                          _kdNodes;
+        KDLeafList                          _kdLeaves;
+
+        osg::ref_ptr<osg::Vec3Array>        _vertices;
+        
+        Indices                             _primitiveIndices;
+        
+        BoundingBoxList                     _boundingBoxes;
+        TriangleList                        _triangles;
+        CenterList                          _centers;
 
 };
 
@@ -131,13 +229,10 @@ class OSG_EXPORT KdTreeBuilder : public osg::NodeVisitor
     
         void apply(osg::Geode& geode);
     
+        KdTree::BuildOptions _buildOptions;
+
         osg::ref_ptr<osg::KdTree> _kdTreePrototype;
-        
-        unsigned int _maxNumLevels;
-        unsigned int _targetNumTrianglesPerLeaf;
-        
-        unsigned int _numVerticesProcessed;
-        
+                
     protected:
     
         virtual ~KdTreeBuilder() {}
diff --git a/src/osg/KdTree.cpp b/src/osg/KdTree.cpp
index 40663bb3e..e0f2e1b14 100644
--- a/src/osg/KdTree.cpp
+++ b/src/osg/KdTree.cpp
@@ -14,20 +14,292 @@
 #include <osg/KdTree>
 #include <osg/Geode>
 #include <osg/io_utils>
+#include <osg/TriangleIndexFunctor>
 
 using namespace osg;
 
+//#define VERBOSE_OUTPUT
+
+////////////////////////////////////////////////////////////////////////////////
+//
+// Functor for collecting triangle indices from Geometry
+
+struct TriangleIndicesCollector
+{
+    TriangleIndicesCollector():
+        _kdTree(0)
+    {
+    }
+
+    inline void operator () (unsigned int p1, unsigned int p2, unsigned int p3)
+    {
+        unsigned int i = _kdTree->_triangles.size();
+        _kdTree->_triangles.push_back(KdTree::Triangle(p1,p2,p3));
+        
+        osg::BoundingBox bb;
+        bb.expandBy((*(_kdTree->_vertices))[p1]);
+        bb.expandBy((*(_kdTree->_vertices))[p2]);
+        bb.expandBy((*(_kdTree->_vertices))[p3]);
+        _kdTree->_boundingBoxes.push_back(bb);
+        
+        _kdTree->_centers.push_back(bb.center());
+
+        _kdTree->_primitiveIndices.push_back(i);
+        
+    }
+    
+    KdTree* _kdTree;
+
+};
 
 ////////////////////////////////////////////////////////////////////////////////
 //
 // KdTree
 
-bool KdTree::build(osg::Geometry* geometry)
+KdTree::KdTree()
+{
+}
+
+KdTree::KdTree(const KdTree& rhs, const osg::CopyOp& copyop):
+    Shape(rhs)
+{
+}
+
+bool KdTree::build(BuildOptions& options, osg::Geometry* geometry)
 {
     osg::notify(osg::NOTICE)<<"KdTree::build("<<geometry<<")"<<std::endl;
+
+#ifdef VERBOSE_OUTPUT    
+    osg::notify(osg::NOTICE)<<"osg::KDTreeBuilder::createKDTree()"<<std::endl;
+#endif
+
+    osg::Vec3Array* vertices = dynamic_cast<osg::Vec3Array*>(geometry->getVertexArray());
+    if (!vertices) return false;
+    
+    if (vertices->size() <= options._targetNumTrianglesPerLeaf) return false;
+
+    _geometry = geometry;
+    _bb = _geometry->getBound();
+    _vertices = vertices;
+    
+    unsigned int estimatedSize = (unsigned int)(2.0*float(vertices->size())/float(options._targetNumTrianglesPerLeaf));
+
+#ifdef VERBOSE_OUTPUT    
+    osg::notify(osg::NOTICE)<<"kdTree->_kdNodes.reserve()="<<estimatedSize<<std::endl<<std::endl;
+#endif
+
+    _kdNodes.reserve(estimatedSize);
+    _kdLeaves.reserve(estimatedSize);
+    
+    computeDivisions(options);
+
+    options._numVerticesProcessed += vertices->size();
+
+    unsigned int estimatedNumTriangles = vertices->size()*2;
+    _primitiveIndices.reserve(estimatedNumTriangles);
+    _boundingBoxes.reserve(estimatedNumTriangles);
+    _triangles.reserve(estimatedNumTriangles);
+    _centers.reserve(estimatedNumTriangles);
+
+
+
+    osg::TriangleIndexFunctor<TriangleIndicesCollector> collectTriangleIndices;
+    collectTriangleIndices._kdTree = this;
+    geometry->accept(collectTriangleIndices);
+
+    _primitiveIndices.reserve(vertices->size());
+
+
+
+    KdLeaf leaf(0, _primitiveIndices.size());
+
+    int leafNum = addLeaf(leaf);
+
+    osg::BoundingBox bb = _bb;
+    int nodeNum = divide(options, bb, leafNum, 0);
+
+#ifdef VERBOSE_OUTPUT    
+    osg::notify(osg::NOTICE)<<"Root nodeNum="<<nodeNum<<std::endl;
+#endif
+    
+    
+//    osg::notify(osg::NOTICE)<<"_kdNodes.size()="<<k_kdNodes.size()<<"  estimated size = "<<estimatedSize<<std::endl;
+//    osg::notify(osg::NOTICE)<<"_kdLeaves.size()="<<_kdLeaves.size()<<"  estimated size = "<<estimatedSize<<std::endl<<std::endl;
+
+
     return true;
 }
 
+void KdTree::computeDivisions(BuildOptions& options)
+{
+    osg::Vec3 dimensions(_bb.xMax()-_bb.xMin(),
+                         _bb.yMax()-_bb.yMin(),
+                         _bb.zMax()-_bb.zMin());
+
+#ifdef VERBOSE_OUTPUT    
+    osg::notify(osg::NOTICE)<<"computeDivisions("<<options._maxNumLevels<<") "<<dimensions<< " { "<<std::endl;
+#endif
+
+    _axisStack.reserve(options._maxNumLevels);
+ 
+    int level = 0;
+    for(unsigned int level=0; level<options._maxNumLevels; ++level)
+    {
+        int axis = 0;
+        if (dimensions[0]>=dimensions[1])
+        {
+            if (dimensions[0]>=dimensions[2]) axis = 0;
+            else axis = 2;
+        }
+        else if (dimensions[1]>=dimensions[2]) axis = 1;
+        else axis = 2;
+
+        _axisStack.push_back(axis);
+        dimensions[axis] /= 2.0f;
+
+#ifdef VERBOSE_OUTPUT    
+        osg::notify(osg::NOTICE)<<"  "<<level<<", "<<dimensions<<", "<<axis<<std::endl;
+#endif
+    }
+
+#ifdef VERBOSE_OUTPUT    
+    osg::notify(osg::NOTICE)<<"}"<<std::endl;
+#endif
+}
+
+int KdTree::divide(BuildOptions& options, osg::BoundingBox& bb, int nodeIndex, unsigned int level)
+{
+    if (_axisStack.size()<=level) return nodeIndex;
+
+    int axis = _axisStack[level];
+
+#ifdef VERBOSE_OUTPUT    
+    //osg::notify(osg::NOTICE)<<"divide("<<nodeIndex<<", "<<level<< "), axis="<<axis<<std::endl;
+#endif
+
+    if (nodeIndex>=0)
+    {
+#ifdef VERBOSE_OUTPUT    
+        osg::notify(osg::NOTICE)<<"  divide node"<<std::endl;
+#endif
+        KdNode& node = getNode(nodeIndex);
+        return nodeIndex;
+    }
+    else
+    {    
+
+        if (getLeaf(nodeIndex).second<=options._targetNumTrianglesPerLeaf) return nodeIndex;
+
+        //osg::notify(osg::NOTICE)<<"  divide leaf"<<std::endl;
+        
+        int nodeNum = addNode(KdNode());
+
+        float original_min = bb._min[axis];
+        float original_max = bb._max[axis];
+
+        float mid = (original_min+original_max)*0.5f;
+
+        {
+            KdLeaf& leaf = getLeaf(nodeIndex);
+
+            //osg::Vec3Array* vertices = kdTree._vertices.get();
+            int end = leaf.first+leaf.second-1;
+            int left = leaf.first;
+            int right = leaf.first+leaf.second-1;
+            
+            while(left<right)
+            {
+                while(left<right && (_centers[_primitiveIndices[left]][axis]<=mid)) { ++left; }
+
+                while(left<right && (_centers[_primitiveIndices[right]][axis]>mid)) { --right; }
+                
+                while(left<right && (_centers[_primitiveIndices[right]][axis]>mid)) { --right; }
+
+                if (left<right)
+                {
+                    std::swap(_primitiveIndices[left], _primitiveIndices[right]);
+                    ++left;
+                    --right;
+                }
+            }
+            
+            if (left==right)
+            {
+                if (_centers[_primitiveIndices[left]][axis]<=mid) ++left;
+                else --right;
+            }
+            
+            KdLeaf leftLeaf(leaf.first, (right-leaf.first)+1);
+            KdLeaf rightLeaf(left, (end-left)+1);
+
+#if 0            
+            osg::notify(osg::NOTICE)<<"In  leaf.first     ="<<leaf.first     <<" leaf.second     ="<<leaf.second<<std::endl;
+            osg::notify(osg::NOTICE)<<"    leftLeaf.first ="<<leftLeaf.first <<" leftLeaf.second ="<<leftLeaf.second<<std::endl;
+            osg::notify(osg::NOTICE)<<"    rightLeaf.first="<<rightLeaf.first<<" rightLeaf.second="<<rightLeaf.second<<std::endl;
+            osg::notify(osg::NOTICE)<<"    left="<<left<<" right="<<right<<std::endl;
+
+            if (leaf.second != (leftLeaf.second +rightLeaf.second))
+            {
+                osg::notify(osg::NOTICE)<<"*** Error in size, leaf.second="<<leaf.second
+                                        <<", leftLeaf.second="<<leftLeaf.second
+                                        <<", rightLeaf.second="<<rightLeaf.second<<std::endl;
+            }
+            else
+            {
+                osg::notify(osg::NOTICE)<<"Size OK, leaf.second="<<leaf.second
+                                        <<", leftLeaf.second="<<leftLeaf.second
+                                        <<", rightLeaf.second="<<rightLeaf.second<<std::endl;
+            }
+#endif
+            if (leftLeaf.second<=0)
+            {
+                //osg::notify(osg::NOTICE)<<"LeftLeaf empty"<<std::endl;
+                getNode(nodeNum).first = 0;
+                getNode(nodeNum).second = replaceLeaf(nodeIndex, rightLeaf);
+            }
+            else if (rightLeaf.second<=0)
+            {
+                //osg::notify(osg::NOTICE)<<"RightLeaf empty"<<std::endl;
+                getNode(nodeNum).first = replaceLeaf(nodeIndex, leftLeaf);
+                getNode(nodeNum).second = 0;
+            }
+            else
+            {
+                getNode(nodeNum).first = replaceLeaf(nodeIndex, leftLeaf);
+                getNode(nodeNum).second = addLeaf(rightLeaf);
+            }
+        }
+
+        
+        int originalLeftChildIndex = getNode(nodeNum).first;
+        int originalRightChildIndex = getNode(nodeNum).second;
+
+
+        float restore = bb._max[axis];
+        bb._max[axis] = mid;
+
+        //osg::notify(osg::NOTICE)<<"  divide leftLeaf "<<kdTree.getNode(nodeNum).first<<std::endl;
+        int leftChildIndex = divide(options, bb, originalLeftChildIndex, level+1);
+
+        bb._max[axis] = restore;
+        
+        restore = bb._min[axis];
+        bb._min[axis] = mid;
+
+        //osg::notify(osg::NOTICE)<<"  divide rightLeaf "<<kdTree.getNode(nodeNum).second<<std::endl;
+        int rightChildIndex = divide(options, bb, originalRightChildIndex, level+1);
+        
+        bb._min[axis] = restore;
+        
+        getNode(nodeNum).first = leftChildIndex;
+        getNode(nodeNum).second = rightChildIndex; 
+        
+        
+        return nodeNum;        
+    }
+}
+
+
 bool KdTree::intersect(const osg::Vec3& start, const osg::Vec3& end, LineSegmentIntersections& intersections)
 {
     osg::notify(osg::NOTICE)<<"KdTree::intersect("<<start<<","<<end<<")"<<std::endl;
@@ -38,10 +310,7 @@ bool KdTree::intersect(const osg::Vec3& start, const osg::Vec3& end, LineSegment
 //
 // KdTreeBuilder
 KdTreeBuilder::KdTreeBuilder():
-    osg::NodeVisitor(osg::NodeVisitor::TRAVERSE_ALL_CHILDREN),
-    _maxNumLevels(16),
-    _targetNumTrianglesPerLeaf(4),
-    _numVerticesProcessed(0)
+    osg::NodeVisitor(osg::NodeVisitor::TRAVERSE_ALL_CHILDREN)
 {            
     _kdTreePrototype = new osg::KdTree;
 }
@@ -60,7 +329,7 @@ void KdTreeBuilder::apply(osg::Geode& geode)
 
             osg::ref_ptr<osg::KdTree> kdTree = dynamic_cast<osg::KdTree*>(_kdTreePrototype->cloneType());
 
-            if (kdTree->build(geom))
+            if (kdTree->build(_buildOptions, geom))
             {
                 geom->setShape(kdTree.get());
             }