octree.h

Go to the documentation of this file.

/*******************************************************************************
** Copyright (c) 2020 MAK Technologies, Inc.
** All rights reserved.
*******************************************************************************/



#pragma once

#include <vrfutil/kinematicTools.h>

#include <cmdLine/cmdStdOutput.h>

#include <geometry/extent.h>

#include <atomic>
#include <tbb/spin_rw_mutex.h>

#include <cstdarg>
#include <map>
#include <unordered_map>
#include <vector>


template<class T> class DtOctreeNode;


#define OCTREE_DEBUG 0

#if OCTREE_DEBUG

static int octree_debug_level = 0;


static void DEBUG_OUTPUT(const char* fmt, ...)
{
   char buffer[8192];
   va_list ap;
   va_start(ap, fmt);
   vsprintf(buffer, fmt, ap);
   va_end(ap);

   DtInfo("OCT %*s%s\n", octree_debug_level*2, "", buffer);
}

#  define DEBUG_PUSH do { octree_debug_level++; } while(0)
#  define DEBUG_POP do { octree_debug_level--; } while(0)
#else

static inline void DEBUG_OUTPUT(const char* fmt, ...)
{ }

#  define DEBUG_PUSH do { } while(0)
#  define DEBUG_POP do { } while(0)
#endif

static const int NUM_BRANCHES = 8;

static const int MAX_BRANCH_DEPTH = 12;

static const int MAX_CHILD_MEMBERS = 16;

static int theBitmaskList[NUM_BRANCHES];
static bool theBitmaskListCreated = false;

#define EXT_FMT "[%lf, %lf, %lf] - [%lf, %lf, %lf]"
#define EXT_ARGS(x) (x).minX(), (x).minY(), (x).minZ(), (x).maxX(), (x).maxY(), (x).maxZ()
#define OBJ_ARGS(x) ((const char*) (x)->objectName())


template<class T> class DtOctreeMember
{
public:
   DtOctreeMember(T* object, const DtExtent& extent)
      : myObject(object)
      , myExtent(extent)
   {
   }

   void bind(DtOctreeNode<T>* node);

   void unbind(DtOctreeNode<T>* node);

   void unbind();

   bool testKey(const unsigned int key)
   {
      unsigned int threadId = DtVRFGetThreadId();

      tbb::spin_rw_mutex::scoped_lock lock(mySearchKeysMutex, false);
      std::unordered_map<unsigned int, unsigned int>::const_iterator iter = mySearchKeys.find(threadId);

      if ((iter != mySearchKeys.end()) && (key == iter->second))
      {
         return false; 
      }

      lock.upgrade_to_writer();
      mySearchKeys[threadId] = key;

      return true;
   }

protected:
   template<class U> friend class DtOctreeNode;
   template<class U> friend class DtOctree;

   T* myObject;

   DtExtent myExtent;

   typedef std::pair<DtOctreeNode<T>*, typename std::list<DtOctreeMember*>::iterator> NodeBinding;

   std::list<NodeBinding> myBindings;

   std::unordered_map<unsigned int, unsigned int> mySearchKeys;
   tbb::spin_rw_mutex mySearchKeysMutex;
};


template<class T>
class DtSpatialSelectionFunctorTemplate
{
public:
   virtual ~DtSpatialSelectionFunctorTemplate() {};

   virtual bool operator()(const T *object) = 0;
};


template<class T> class DtOctreeNode
{
protected: 

   typedef          std::list<DtOctreeMember<T>*>           MemberContainer;
   typedef typename std::list<DtOctreeMember<T>*>::iterator MemberIterator;

   static int sNumNodes;
   static const int MAX_NUM_NODES = 1000;

   DtOctreeNode(const DtExtent &region, const unsigned int depth)
      : myRegion(region)
      , myDepth(depth)
      , myCenter(region.center())
      , myChildren(0)
   { 
      DEBUG_OUTPUT("Create node " EXT_FMT, EXT_ARGS(region));

      ++sNumNodes;
   }

   ~DtOctreeNode()
   {
      DEBUG_OUTPUT("Destroy node " EXT_FMT, EXT_ARGS(myRegion));
      DEBUG_PUSH;

      if (myChildren)
      {
         for(int i=0; i<NUM_BRANCHES; i++)
            delete myChildren[i];
      }
      delete[] myChildren;

      DEBUG_POP;

      --sNumNodes;
   }

   
   virtual void visitObjects(const unsigned int key, const DtExtent& extent, DtSpatialSelectionFunctorTemplate<T> &callback)
   {
      DEBUG_OUTPUT(EXT_FMT "::visit(" EXT_FMT ", cb)", EXT_ARGS(myRegion), EXT_ARGS(extent));
      DEBUG_PUSH;

      MemberIterator iter, end;

      for(iter = myMembers.begin(), end = myMembers.end(); iter != end; ++iter)
      {
         DEBUG_OUTPUT("Member %s " EXT_FMT, OBJ_ARGS((*iter)->myObject), EXT_ARGS((*iter)->myExtent));
         if ((*iter)->testKey(key))
         {
            if (!(*iter)->myObject->isDeleted() && callback((*iter)->myObject))
            {
               break;
            }
         }
         else
         {
            return;
         }
      }

      if (myChildren)
      {
         int mask = intersectionMask(extent);
         for(int i = 0; i < NUM_BRANCHES; ++i)
         {
            if (myChildren[i] && (mask & theBitmaskList[i]))
            {
               myChildren[i]->visitObjects(key, extent, callback);
            }
         }
      }

      DEBUG_POP;
   }
   
   virtual void addMember(DtOctreeMember<T>* member, int& depth)
   {
      depth++;

      DEBUG_OUTPUT(EXT_FMT "::addMember(%s)", EXT_ARGS(myRegion), OBJ_ARGS(member->myObject));
      DEBUG_PUSH;
      
#if OCTREE_DEBUG
      if (!myRegion.contains(member->myExtent))
      {
         DEBUG_OUTPUT("FAIL!  " EXT_FMT " does not contain " EXT_FMT, EXT_ARGS(myRegion), EXT_ARGS(member->myExtent));
      }
#endif

      if (myChildren)
      {
         int mask = intersectionMask(member->myExtent);

         if (mask == 0xff)
         {
            member->bind(this);
         }
         else
         {
            for (int childIndex = 0; childIndex < NUM_BRANCHES; ++childIndex)
            {
               if (mask & theBitmaskList[childIndex])
               {
                  addMemberToChild(member, childIndex, depth);
               }
            }
         }
      }
      else 
      {
         DEBUG_OUTPUT("add as child member (not enough child members yet)");

         member->bind(this);

         //if (numChildCandidates() >= MAX_CHILD_MEMBERS && myDepth < MAX_BRANCH_DEPTH)
         if (numChildCandidates() >= MAX_CHILD_MEMBERS && sNumNodes < MAX_NUM_NODES - 8)
         {
            split(depth);
         }
      }

      DEBUG_POP;
   }

   bool coalesceEmptyChildren()
   {
      DEBUG_OUTPUT(EXT_FMT "::coalesceEmptyChildren()", EXT_ARGS(myRegion));
      DEBUG_PUSH;

      bool retVal = (myMembers.size() == 0);
      
      if (myChildren)
      {
         for(int i=0; i<NUM_BRANCHES; i++)
         {
            if (myChildren[i])
            {
               if (myChildren[i]->coalesceEmptyChildren())
               {
                  DEBUG_OUTPUT("Deleting child %d\n", i);
                  delete myChildren[i];
                  myChildren[i] = 0;
               } else {
                  retVal = false;
               }
            }
         }
      }

      DEBUG_POP;

      return retVal;
   }

   int numChildCandidates()
   {
      double retVal = 0;
      DtExtent matchExtent;

      MemberIterator iter, end;
      for(iter = myMembers.begin(), end = myMembers.end(); iter != end; ++iter)
      {
         double numBuckets = 0;
         int mask = intersectionMask((*iter)->myExtent);

         if ((*iter)->myExtent == matchExtent)
         {
            continue;
         }

         matchExtent = (*iter)->myExtent;

         if (mask != 0xff)
         {
            for(int childIndex = 0; childIndex < NUM_BRANCHES; ++childIndex)
            {
               if (mask & theBitmaskList[childIndex])
               {
                  numBuckets++;
               }
            }

            retVal += (1.0 / numBuckets);

            if (retVal >= MAX_CHILD_MEMBERS)
            {
               break;
            }
         }
      }

      return retVal;
   }

   int intersectionMask(const DtExtent &extent) const
   {
      int result = 0xff;                  

      if (extent.minX() >= myCenter.x())
         result &= 0xf0;                  
      if (extent.maxX() <= myCenter.x())
         result &= 0x0f;                  

      if (extent.minY() >= myCenter.y())
         result &= 0xcc;                  
      if (extent.maxY() <= myCenter.y())
         result &= 0x33;                  

      if (extent.minZ() >= myCenter.z())
         result &= 0xaa;                  
      if (extent.maxZ() <= myCenter.z())
         result &= 0x55;                  

      if (result == 0)
      {
         return 0xff;                     
      }

      return result;
   }

   void addMemberToChild(DtOctreeMember<T>* member, int childIndex, int& depth)
   {
      DEBUG_OUTPUT(EXT_FMT "::addMemberToChild(%s, %d)", EXT_ARGS(myRegion), OBJ_ARGS(member->myObject), childIndex);
      DEBUG_PUSH;

      if (myChildren[childIndex] == 0)
      {
         DtPoint childCorner;

         if (childIndex & 4)
            childCorner.setX(myRegion.maxX());
         else
            childCorner.setX(myRegion.minX());

         if (childIndex & 2)
            childCorner.setY(myRegion.maxY());
         else
            childCorner.setY(myRegion.minY());

         if (childIndex & 1)
            childCorner.setZ(myRegion.maxZ());
         else
            childCorner.setZ(myRegion.minZ());

         myChildren[childIndex] = new DtOctreeNode(DtExtent(childCorner, myCenter), myDepth + 1);
      }

      myChildren[childIndex]->addMember(member, depth);

      DEBUG_POP;
   }

   void split(int& depth)
   {
      DEBUG_OUTPUT(EXT_FMT "::split()", EXT_ARGS(myRegion));
      DEBUG_PUSH;

      myChildren = new DtOctreeNode<T>* [NUM_BRANCHES];

      for(int i=0; i<NUM_BRANCHES; i++)
      {
         myChildren[i] = 0;
      }

      MemberIterator iter, next, end;
      for(iter = myMembers.begin(), end = myMembers.end(); iter != end; iter = next)
      {
         next = iter;
         ++next;

         DtOctreeMember<T>* member = *iter;

         int mask = intersectionMask(member->myExtent);

         if (mask != 0xff)
         {
            member->unbind(this);

            for(int childIndex = 0; childIndex < NUM_BRANCHES; ++childIndex)
            {
               if (mask & theBitmaskList[childIndex])
               {
                  addMemberToChild(member, childIndex, depth);
               }
            }
         }
      }

      DEBUG_POP;
   }

   const DtExtent myRegion;
   
   const DtPoint myCenter;


   unsigned int myDepth;

   DtOctreeNode<T> **myChildren;

   MemberContainer myMembers;

   template<class U> friend class DtOctree;
   template<class U> friend class DtOctreeMember;
   template<class U> friend class DtOctreeDebugger;
};

template<class T> int DtOctreeNode<T>::sNumNodes = 0;


template<class T> class DtOctree
{
public:

   typedef          std::map<T*, DtOctreeMember<T>*>           MemberContainer;
   typedef typename std::map<T*, DtOctreeMember<T>*>::iterator MemberIterator;

   typedef          std::list<T*>                            OutsideAreaContainer;
   typedef typename std::list<T*>::iterator                  OutsideAreaIterator;

   typedef DtSpatialSelectionFunctorTemplate<T> SpatialSelectionFunctor;

   DtOctree(const DtExtent &region) : mySearchKey(0)
   {
      myRootNode = new DtOctreeNode<T>(region, 0);

      if (!theBitmaskListCreated)
      {
         theBitmaskListCreated = true;
         for (int childIndex = 0; childIndex < NUM_BRANCHES; ++childIndex)
         {
            theBitmaskList[childIndex] = (1 << childIndex);
         }
      }
   }

   virtual ~DtOctree()
   {
      delete myRootNode;

      MemberIterator iter, end;
      for(iter = myMemberMap.begin(), end = myMemberMap.end(); iter != end; ++iter)
      {
         delete (*iter).second;
      }
   }

   virtual void visitObjects(const DtExtent& extent, DtSpatialSelectionFunctorTemplate<T> &callback)
   {
      tbb::spin_rw_mutex::scoped_lock lock(myMutex, false);
      ++mySearchKey; 

      if (myRootNode->myRegion.intersects(extent))
      {
         myRootNode->visitObjects(mySearchKey, extent, callback);
      }

      OutsideAreaIterator iter, end;
      for(iter = myObjectsOutsideArea.begin(), end = myObjectsOutsideArea.end(); iter != end; ++iter)
      {
         callback(*iter);
      }
   }


   virtual void addObject(T* object, const DtExtent& extent)
   {
      DEBUG_OUTPUT("DtOctree::addObject(%s, " EXT_FMT ")", OBJ_ARGS(object), EXT_ARGS(extent));
      DEBUG_PUSH;

      tbb::spin_rw_mutex::scoped_lock lock(myMutex);
      if (myRootNode->myRegion.contains(extent))
      {
         DtOctreeMember<T>* newMember = new DtOctreeMember<T>(object, extent);
         int depth = 0;

         myRootNode->addMember(newMember, depth);
         myMemberMap[object] = newMember;

      }
      else
      {
         DtWarn("Member %s (" EXT_FMT ") added outside octree (" EXT_FMT ")\n",
                OBJ_ARGS(object), EXT_ARGS(extent), EXT_ARGS(myRootNode->myRegion));
         myObjectsOutsideArea.push_back(object);
      }

      DEBUG_POP;
   }

   virtual void removeObject(T* object, bool needsLock = true)
   {
      DEBUG_OUTPUT("DtOctree::removeObject(%s)", OBJ_ARGS(object));
      DEBUG_PUSH;

      tbb::spin_rw_mutex::scoped_lock* lock = 0;
      
      if (needsLock)
      {
         lock = new tbb::spin_rw_mutex::scoped_lock(myMutex, true);
      }

      MemberIterator iter;
      iter = myMemberMap.find(object);
      if (iter != myMemberMap.end())
      {
         DtOctreeMember<T>* member = (*iter).second;
         member->unbind();
         myMemberMap.erase(iter);
         delete member;
      }
      else if (!removeObjectOutsideArea(object))
      {

         //DtWarn("Tried to remove %s from octree, but couldn't find!\n", OBJ_ARGS(object));
      }

      DEBUG_POP;

      if (lock)
      {
         delete lock;
      }
   }

   virtual void updatePositions()
   {
      tbb::spin_rw_mutex::scoped_lock lock(myMutex, false);
      MemberIterator cur = myMemberMap.begin();
      while (cur != myMemberMap.end())
      {
         MemberIterator iter = cur++;

         T* object = (*iter).first;

         if (object->isValid())
         {
            DtOctreeMember<T>* member = (*iter).second;
            DtExtent roughExtent = object->extent();

            if (member->myExtent != roughExtent)
            {
               member->myExtent = roughExtent;

               if (!myRootNode->myRegion.contains(member->myExtent))
               {

                  DtWarn("Object %s (" EXT_FMT ") moving outside octree (" EXT_FMT ")\n",
                     OBJ_ARGS(object), EXT_ARGS(member->myExtent), EXT_ARGS(myRootNode->myRegion));
                  lock.upgrade_to_writer();
                  removeObject(object, false);
                  myObjectsOutsideArea.push_back(object);
                  lock.downgrade_to_reader();
               }
               else
               {

                  member->unbind();
                  int depth = 0;

                  myRootNode->addMember(member, depth);
               }
            }
         }
      }

      myRootNode->coalesceEmptyChildren();
   }

   bool containsObject(T* object)
   {
      tbb::spin_rw_mutex::scoped_lock lock(myMutex, false);
      if (myMemberMap.find(object) != myMemberMap.end())
      {
         return true;
      }
      else
      {
         OutsideAreaIterator iter, end;
         for(iter = myObjectsOutsideArea.begin(), end = myObjectsOutsideArea.end(); iter != end; ++iter)
         {
            if (*iter == object)
            {
               return true;
            }
         }
      }
      return false;
   }

   int numObjects()
   {
      return myMemberMap.size() + myObjectsOutsideArea.size();
   }

   int numObjectsOutside()
   {
      return myObjectsOutsideArea.size();
   }

protected:
   bool removeObjectOutsideArea(T* object)
   {
      OutsideAreaIterator iter, end;
      for(iter = myObjectsOutsideArea.begin(), end = myObjectsOutsideArea.end(); iter != end; ++iter)
      {
         if (*iter == object)
         {
            myObjectsOutsideArea.erase(iter);
            return true;
         }
      }

      return false;
   }

   tbb::spin_rw_mutex myMutex;
   DtOctreeNode<T>* myRootNode;
   
   OutsideAreaContainer myObjectsOutsideArea;
   
   MemberContainer myMemberMap;

   std::atomic<unsigned int> mySearchKey;
};


#if OCTREE_DEBUG

class DtDebugOctreeFunctor : public DtSpatialSelectionFunctor
{
public:
   DtDebugOctreeFunctor(DtSpatialSelectionFunctor& inner)
      : myInner(inner)
   {
   }  

   virtual bool operator()(const DtSimObject* object)
   {
      mySeenObjects.insert(object);
      return myInner(object);
   }

   std::set<const DtSimObject*> mySeenObjects;
   DtSpatialSelectionFunctor& myInner;
};

template<class T> class DtOctreeDebugger : public DtOctree<T>
{
public:
   DtOctreeDebugger(const DtExtent& extent)
      : DtOctree<T>(extent)
   { }

   void addObject(T* object, const DtExtent& extent)
   {
      allObjects[object] = extent;
      DtOctree<T>::addObject(object, extent);
   }

   void removeObject(T* object)
   {
      allObjects.erase(object);
      DtOctree<T>::removeObject(object);
   }

   void updatePositions()
   {
      std::map<T*, DtExtent>::iterator iter = allObjects.begin();
      std::map<T*, DtExtent>::iterator end  = allObjects.end();

      for(; iter != end; ++iter)
      {
         T* object = (*iter).first;
         DtExtent& extent = (*iter).second;

         extent = object->roughExtent();
      }

      DtOctree<T>::updatePositions();
   }

   void visitObjects(const DtExtent& extent, DtSpatialSelectionFunctor &callback)
   {
      DtDebugOctreeFunctor debugFunctor(callback);
      DtOctree<T>::visitObjects(extent, debugFunctor);

      std::map<T*, DtExtent>::iterator iter = allObjects.begin();
      std::map<T*, DtExtent>::iterator end  = allObjects.end();

      for(; iter != end; ++iter)
      {
         T* object = (*iter).first;
         const DtExtent& objectExtent = (*iter).second;

         if (objectExtent.intersects(extent) && debugFunctor.mySeenObjects.count(object) == 0)
         {
            DEBUG_OUTPUT("FAILED to visit object %s", OBJ_ARGS(object));
            DEBUG_PUSH;
            DEBUG_OUTPUT("Old extent is " EXT_FMT, EXT_ARGS(objectExtent));
            DEBUG_OUTPUT("New extent is " EXT_FMT, EXT_ARGS(object->roughExtent()));
            DEBUG_POP;
         }
      }
   }
protected:
   std::map<T*, DtExtent> allObjects;
};


static DtOctree<DtSimObject>* createOctree(const DtExtent& extent)
{
   return new DtOctreeDebugger<DtSimObject>(extent);
}

#else



#endif


template<class T> void DtOctreeMember<T>::bind(DtOctreeNode<T>* node)
{
   for (typename std::list<NodeBinding>::const_iterator it = myBindings.begin(); it != myBindings.end(); ++it)
   {
      if (node == (*it).first)
      {
         return; 
      }
   }

   myBindings.push_back(NodeBinding(node, node->myMembers.insert(node->myMembers.end(), this)));
}

template<class T> void DtOctreeMember<T>::unbind(DtOctreeNode<T>* node)
{
   for (typename std::list<NodeBinding>::iterator it = myBindings.begin(); it != myBindings.end(); ++it)
   {
      if (node == (*it).first)
      {
         node->myMembers.erase((*it).second);
         myBindings.erase(it);
         return;
      }
   }
}

template<class T> void DtOctreeMember<T>::unbind()
{
   for (typename std::list<NodeBinding>::const_iterator it = myBindings.begin(); it != myBindings.end(); ++it)
   {
      (*it).first->myMembers.erase((*it).second);
   }
   myBindings.clear();
}