quadtree.c

/*
 * This code was initially authored by the Stackoverflow user dragon-energy and
 * posted under following page:
 * https://stackoverflow.com/questions/41946007/efficient-and-well-explained-implementation-of-a-quadtree-for-2d-collision-det
 *
 * As for the license, the author has kindly noted:
 *
 * "Oh and feel free to use this code I post however you want, even for
 * commercial projects. I would really love it if people let me know if they
 * find it useful, but do as you wish."
 *
 * And generally all Stackoverflow-posted code is by default licensed with CC
 * BY-SA 4.0: https://creativecommons.org/licenses/by-sa/4.0/
 */
#include "quadtree.h"
#include <stdlib.h>
#include <string.h>
 
enum {
   // ----------------------------------------------------------------------------------------
   // Element node fields:
   // ----------------------------------------------------------------------------------------
   enode_num = 2,
 
   // Points to the next element in the leaf node. A value of -1
   // indicates the end of the list.
   enode_idx_next = 0,
 
   // Stores the element index.
   enode_idx_elt = 1,
 
   // ----------------------------------------------------------------------------------------
   // Element fields:
   // ----------------------------------------------------------------------------------------
   elt_num = 5,
 
   // Stores the rectangle encompassing the element.
   elt_idx_lft = 0,
   elt_idx_top = 1,
   elt_idx_rgt = 2,
   elt_idx_btm = 3,
 
   // Stores the ID of the element.
   elt_idx_id = 4,
 
   // ----------------------------------------------------------------------------------------
   // Node fields:
   // ----------------------------------------------------------------------------------------
   node_num = 2,
 
   // Points to the first child if this node is a branch or the first element
   // if this node is a leaf.
   node_idx_fc = 0,
 
   // Stores the number of elements in the node or -1 if it is not a leaf.
   node_idx_num = 1,
 
   // ----------------------------------------------------------------------------------------
   // Node data fields:
   // ----------------------------------------------------------------------------------------
   nd_num = 6,
 
   // Stores the extents of the node using a centered rectangle and half-size.
   nd_idx_mx = 0,
   nd_idx_my = 1,
   nd_idx_sx = 2,
   nd_idx_sy = 3,
 
   // Stores the index of the node.
   nd_idx_index = 4,
 
   // Stores the depth of the node.
   nd_idx_depth = 5,
};
 
static void node_insert( Quadtree *qt, int index, int depth, int mx, int my,
                         int sx, int sy, int element );
 
static int intersect( int l1, int t1, int r1, int b1, int l2, int t2, int r2,
                      int b2 )
{
   return l2 <= r1 && r2 >= l1 && t2 <= b1 && b2 >= t1;
}
 
static void leaf_insert( Quadtree *qt, int node, int depth, int mx, int my,
                         int sx, int sy, int element )
{
   // Insert the element node to the leaf.
   const int nd_fc = il_get( &qt->nodes, node, node_idx_fc );
   il_set( &qt->nodes, node, node_idx_fc, il_insert( &qt->enodes ) );
   il_set( &qt->enodes, il_get( &qt->nodes, node, node_idx_fc ), enode_idx_next,
           nd_fc );
   il_set( &qt->enodes, il_get( &qt->nodes, node, node_idx_fc ), enode_idx_elt,
           element );
 
   // If the leaf is full, split it.
   if ( il_get( &qt->nodes, node, node_idx_num ) == qt->max_elements &&
        depth < qt->max_depth ) {
      int     fc   = 0;
      IntList elts = { 0 };
      il_create( &elts, 1 );
 
      // Transfer elements from the leaf node to a list of elements.
      while ( il_get( &qt->nodes, node, node_idx_fc ) != -1 ) {
         const int index      = il_get( &qt->nodes, node, node_idx_fc );
         const int next_index = il_get( &qt->enodes, index, enode_idx_next );
         const int elt        = il_get( &qt->enodes, index, enode_idx_elt );
 
         // Pop off the element node from the leaf and remove it from the qt.
         il_set( &qt->nodes, node, node_idx_fc, next_index );
         il_erase( &qt->enodes, index );
 
         // Insert element to the list.
         il_set( &elts, il_push_back( &elts ), 0, elt );
      }
 
      // Start by allocating 4 child nodes.
      fc = il_insert( &qt->nodes );
      il_insert( &qt->nodes );
      il_insert( &qt->nodes );
      il_insert( &qt->nodes );
      il_set( &qt->nodes, node, node_idx_fc, fc );
 
      // Initialize the new child nodes.
      for ( int j = 0; j < 4; ++j ) {
         il_set( &qt->nodes, fc + j, node_idx_fc, -1 );
         il_set( &qt->nodes, fc + j, node_idx_num, 0 );
      }
 
      // Transfer the elements in the former leaf node to its new children.
      il_set( &qt->nodes, node, node_idx_num, -1 );
      for ( int j = 0; j < il_size( &elts ); ++j )
         node_insert( qt, node, depth, mx, my, sx, sy, il_get( &elts, j, 0 ) );
      il_destroy( &elts );
   } else {
      // Increment the leaf element count.
      il_set( &qt->nodes, node, node_idx_num,
              il_get( &qt->nodes, node, node_idx_num ) + 1 );
   }
}
 
static void push_node( IntList *nodes, int nd_index, int nd_depth, int nd_mx,
                       int nd_my, int nd_sx, int nd_sy )
{
   const int back_idx = il_push_back( nodes );
   il_set( nodes, back_idx, nd_idx_mx, nd_mx );
   il_set( nodes, back_idx, nd_idx_my, nd_my );
   il_set( nodes, back_idx, nd_idx_sx, nd_sx );
   il_set( nodes, back_idx, nd_idx_sy, nd_sy );
   il_set( nodes, back_idx, nd_idx_index, nd_index );
   il_set( nodes, back_idx, nd_idx_depth, nd_depth );
}
 
static void find_leaves( IntList *out, const Quadtree *qt, int node, int depth,
                         int mx, int my, int sx, int sy, int lft, int top,
                         int rgt, int btm )
{
   IntList to_process = { 0 };
   il_create( &to_process, nd_num );
   push_node( &to_process, node, depth, mx, my, sx, sy );
 
   while ( il_size( &to_process ) > 0 ) {
      const int back_idx = il_size( &to_process ) - 1;
      const int nd_mx    = il_get( &to_process, back_idx, nd_idx_mx );
      const int nd_my    = il_get( &to_process, back_idx, nd_idx_my );
      const int nd_sx    = il_get( &to_process, back_idx, nd_idx_sx );
      const int nd_sy    = il_get( &to_process, back_idx, nd_idx_sy );
      const int nd_index = il_get( &to_process, back_idx, nd_idx_index );
      const int nd_depth = il_get( &to_process, back_idx, nd_idx_depth );
      il_pop_back( &to_process );
 
      // If this node is a leaf, insert it to the list.
      if ( il_get( &qt->nodes, nd_index, node_idx_num ) != -1 )
         push_node( out, nd_index, nd_depth, nd_mx, nd_my, nd_sx, nd_sy );
      else {
         // Otherwise push the children that intersect the rectangle.
         const int fc = il_get( &qt->nodes, nd_index, node_idx_fc );
         const int hx = nd_sx >> 1, hy = nd_sy >> 1;
         const int l = nd_mx - hx, t = nd_my - hy, r = nd_mx + hx,
                   b = nd_my + hy;
 
         if ( top <= nd_my ) {
            if ( lft <= nd_mx )
               push_node( &to_process, fc + 0, nd_depth + 1, l, t, hx, hy );
            if ( rgt > nd_mx )
               push_node( &to_process, fc + 1, nd_depth + 1, r, t, hx, hy );
         }
         if ( btm > nd_my ) {
            if ( lft <= nd_mx )
               push_node( &to_process, fc + 2, nd_depth + 1, l, b, hx, hy );
            if ( rgt > nd_mx )
               push_node( &to_process, fc + 3, nd_depth + 1, r, b, hx, hy );
         }
      }
   }
   il_destroy( &to_process );
}
 
static void node_insert( Quadtree *qt, int index, int depth, int mx, int my,
                         int sx, int sy, int element )
{
   // Find the leaves and insert the element to all the leaves found.
   IntList leaves = { 0 };
 
   const int lft = il_get( &qt->elts, element, elt_idx_lft );
   const int top = il_get( &qt->elts, element, elt_idx_top );
   const int rgt = il_get( &qt->elts, element, elt_idx_rgt );
   const int btm = il_get( &qt->elts, element, elt_idx_btm );
 
   il_create( &leaves, nd_num );
   find_leaves( &leaves, qt, index, depth, mx, my, sx, sy, lft, top, rgt, btm );
   for ( int j = 0; j < il_size( &leaves ); ++j ) {
      const int nd_mx    = il_get( &leaves, j, nd_idx_mx );
      const int nd_my    = il_get( &leaves, j, nd_idx_my );
      const int nd_sx    = il_get( &leaves, j, nd_idx_sx );
      const int nd_sy    = il_get( &leaves, j, nd_idx_sy );
      const int nd_index = il_get( &leaves, j, nd_idx_index );
      const int nd_depth = il_get( &leaves, j, nd_idx_depth );
      leaf_insert( qt, nd_index, nd_depth, nd_mx, nd_my, nd_sx, nd_sy,
                   element );
   }
   il_destroy( &leaves );
}
 
void qt_create( Quadtree *qt, int x1, int y1, int x2, int y2, int max_elements,
                int max_depth )
{
   qt->max_elements = max_elements;
   qt->max_depth    = max_depth;
   qt->temp         = NULL;
   qt->temp_size    = 0;
   il_create( &qt->nodes, node_num );
   il_create( &qt->elts, elt_num );
   il_create( &qt->enodes, enode_num );
 
   // Insert the root node to the qt.
   il_insert( &qt->nodes );
   il_set( &qt->nodes, 0, node_idx_fc, -1 );
   il_set( &qt->nodes, 0, node_idx_num, 0 );
 
   int half_width  = ( x2 - x1 ) >> 1;
   int half_height = ( y2 - y1 ) >> 1;
   qt->root_sx     = half_width;
   qt->root_sy     = half_height;
 
   // Center
   qt->root_mx = x1 + half_width;
   qt->root_my = y1 + half_height;
}
 
void qt_clear( Quadtree *qt )
{
   il_clear( &qt->nodes );
   il_clear( &qt->elts );
   il_clear( &qt->enodes );
 
   // Insert the root node to the qt.
   il_insert( &qt->nodes );
   il_set( &qt->nodes, 0, node_idx_fc, -1 );
   il_set( &qt->nodes, 0, node_idx_num, 0 );
}
 
void qt_destroy( Quadtree *qt )
{
   il_destroy( &qt->nodes );
   il_destroy( &qt->elts );
   il_destroy( &qt->enodes );
   free( qt->temp );
}
 
int qt_insert( Quadtree *qt, int id, int x1, int y1, int x2, int y2 )
{
   // Insert a new element.
   const int new_element = il_insert( &qt->elts );
 
   // Set the fields of the new element.
   il_set( &qt->elts, new_element, elt_idx_lft, x1 );
   il_set( &qt->elts, new_element, elt_idx_top, y1 );
   il_set( &qt->elts, new_element, elt_idx_rgt, x2 );
   il_set( &qt->elts, new_element, elt_idx_btm, y2 );
   il_set( &qt->elts, new_element, elt_idx_id, id );
 
   // Insert the element to the appropriate leaf node(s).
   node_insert( qt, 0, 0, qt->root_mx, qt->root_my, qt->root_sx, qt->root_sy,
                new_element );
   return new_element;
}
 
void qt_remove( Quadtree *qt, int element )
{
   // Find the leaves.
   IntList leaves = { 0 };
 
   const int lft = il_get( &qt->elts, element, elt_idx_lft );
   const int top = il_get( &qt->elts, element, elt_idx_top );
   const int rgt = il_get( &qt->elts, element, elt_idx_rgt );
   const int btm = il_get( &qt->elts, element, elt_idx_btm );
 
   il_create( &leaves, nd_num );
   find_leaves( &leaves, qt, 0, 0, qt->root_mx, qt->root_my, qt->root_sx,
                qt->root_sy, lft, top, rgt, btm );
 
   // For each leaf node, remove the element node.
   for ( int j = 0; j < il_size( &leaves ); ++j ) {
      const int nd_index = il_get( &leaves, j, nd_idx_index );
 
      // Walk the list until we find the element node.
      int node_index = il_get( &qt->nodes, nd_index, node_idx_fc );
      int prev_index = -1;
      while ( node_index != -1 &&
              il_get( &qt->enodes, node_index, enode_idx_elt ) != element ) {
         prev_index = node_index;
         node_index = il_get( &qt->enodes, node_index, enode_idx_next );
      }
 
      if ( node_index != -1 ) {
         // Remove the element node.
         const int next_index =
            il_get( &qt->enodes, node_index, enode_idx_next );
         if ( prev_index == -1 )
            il_set( &qt->nodes, nd_index, node_idx_fc, next_index );
         else
            il_set( &qt->enodes, prev_index, enode_idx_next, next_index );
         il_erase( &qt->enodes, node_index );
 
         // Decrement the leaf element count.
         il_set( &qt->nodes, nd_index, node_idx_num,
                 il_get( &qt->nodes, nd_index, node_idx_num ) - 1 );
      }
   }
   il_destroy( &leaves );
 
   // Remove the element.
   il_erase( &qt->elts, element );
}
 
void qt_query( Quadtree *qt, IntList *out, int qlft, int qtop, int qrgt,
               int qbtm )
{
   // Find the leaves that intersect the specified query rectangle.
   IntList   leaves  = { 0 };
   const int elt_cap = il_size( &qt->elts );
 
   if ( qt->temp_size < elt_cap ) {
      qt->temp_size = elt_cap;
      qt->temp      = realloc( qt->temp, qt->temp_size * sizeof( *qt->temp ) );
      memset( qt->temp, 0, qt->temp_size * sizeof( *qt->temp ) );
   }
 
   // For each leaf node, look for elements that intersect.
   il_create( &leaves, nd_num );
   find_leaves( &leaves, qt, 0, 0, qt->root_mx, qt->root_my, qt->root_sx,
                qt->root_sy, qlft, qtop, qrgt, qbtm );
 
   il_clear( out );
   for ( int j = 0; j < il_size( &leaves ); ++j ) {
      const int nd_index = il_get( &leaves, j, nd_idx_index );
 
      // Walk the list and add elements that intersect.
      int elt_node_index = il_get( &qt->nodes, nd_index, node_idx_fc );
      while ( elt_node_index != -1 ) {
         const int element =
            il_get( &qt->enodes, elt_node_index, enode_idx_elt );
         const int lft = il_get( &qt->elts, element, elt_idx_lft );
         const int top = il_get( &qt->elts, element, elt_idx_top );
         const int rgt = il_get( &qt->elts, element, elt_idx_rgt );
         const int btm = il_get( &qt->elts, element, elt_idx_btm );
         if ( !qt->temp[element] &&
              intersect( qlft, qtop, qrgt, qbtm, lft, top, rgt, btm ) ) {
            il_set( out, il_push_back( out ), 0, element );
            qt->temp[element] = 1;
         }
         elt_node_index = il_get( &qt->enodes, elt_node_index, enode_idx_next );
      }
   }
   il_destroy( &leaves );
 
   /* Unmark the elements that were inserted, and convert to IDs. */
   for ( int j = 0; j < il_size( out ); ++j ) {
      const int element = il_get( out, j, 0 );
      const int id      = il_get( &qt->elts, element, elt_idx_id );
      qt->temp[element] = 0;
      il_set( out, j, 0, id );
   }
}
 
void qt_cleanup( Quadtree *qt )
{
   IntList to_process = { 0 };
   il_create( &to_process, 1 );
 
   // Only process the root if it's not a leaf.
   if ( il_get( &qt->nodes, 0, node_idx_num ) == -1 ) {
      // Push the root index to the stack.
      il_set( &to_process, il_push_back( &to_process ), 0, 0 );
   }
 
   while ( il_size( &to_process ) > 0 ) {
      // Pop a node from the stack.
      const int node = il_get( &to_process, il_size( &to_process ) - 1, 0 );
      const int fc   = il_get( &qt->nodes, node, node_idx_fc );
      int       num_empty_leaves = 0;
      il_pop_back( &to_process );
 
      // Loop through the children.
      for ( int j = 0; j < 4; ++j ) {
         const int child = fc + j;
 
         // Increment empty leaf count if the child is an empty
         // leaf. Otherwise if the child is a branch, add it to
         // the stack to be processed in the next iteration.
         if ( il_get( &qt->nodes, child, node_idx_num ) == 0 )
            ++num_empty_leaves;
         else if ( il_get( &qt->nodes, child, node_idx_num ) == -1 ) {
            // Push the child index to the stack.
            il_set( &to_process, il_push_back( &to_process ), 0, child );
         }
      }
 
      // If all the children were empty leaves, remove them and
      // make this node the new empty leaf.
      if ( num_empty_leaves == 4 ) {
         // Remove all 4 children in reverse order so that they
         // can be reclaimed on subsequent insertions in proper
         // order.
         il_erase( &qt->nodes, fc + 3 );
         il_erase( &qt->nodes, fc + 2 );
         il_erase( &qt->nodes, fc + 1 );
         il_erase( &qt->nodes, fc + 0 );
 
         // Make this node the new empty leaf.
         il_set( &qt->nodes, node, node_idx_fc, -1 );
         il_set( &qt->nodes, node, node_idx_num, 0 );
      }
   }
   il_destroy( &to_process );
}
 
void qt_traverse( Quadtree *qt, void *user_data, QtNodeFunc *branch,
                  QtNodeFunc *leaf )
{
   IntList to_process = { 0 };
   il_create( &to_process, nd_num );
   push_node( &to_process, 0, 0, qt->root_mx, qt->root_my, qt->root_sx,
              qt->root_sy );
 
   while ( il_size( &to_process ) > 0 ) {
      const int back_idx = il_size( &to_process ) - 1;
      const int nd_mx    = il_get( &to_process, back_idx, nd_idx_mx );
      const int nd_my    = il_get( &to_process, back_idx, nd_idx_my );
      const int nd_sx    = il_get( &to_process, back_idx, nd_idx_sx );
      const int nd_sy    = il_get( &to_process, back_idx, nd_idx_sy );
      const int nd_index = il_get( &to_process, back_idx, nd_idx_index );
      const int nd_depth = il_get( &to_process, back_idx, nd_idx_depth );
      const int fc       = il_get( &qt->nodes, nd_index, node_idx_fc );
      il_pop_back( &to_process );
 
      if ( il_get( &qt->nodes, nd_index, node_idx_num ) == -1 ) {
         // Push the children of the branch to the stack.
         const int hx = nd_sx >> 1, hy = nd_sy >> 1;
         const int l = nd_mx - hx, t = nd_my - hy, r = nd_mx + hx,
                   b = nd_my + hy;
         push_node( &to_process, fc + 0, nd_depth + 1, l, t, hx, hy );
         push_node( &to_process, fc + 1, nd_depth + 1, r, t, hx, hy );
         push_node( &to_process, fc + 2, nd_depth + 1, l, b, hx, hy );
         push_node( &to_process, fc + 3, nd_depth + 1, r, b, hx, hy );
         if ( branch )
            branch( qt, user_data, nd_index, nd_depth, nd_mx, nd_my, nd_sx,
                    nd_sy );
      } else if ( leaf )
         leaf( qt, user_data, nd_index, nd_depth, nd_mx, nd_my, nd_sx, nd_sy );
   }
   il_destroy( &to_process );
}