extra/quadtrees/quadtrees.factor

   1 ! (c) 2009 Joe Groff, see BSD license
   2 USING: assocs kernel math.geometry.rect combinators accessors
   3 math.vectors vectors sequences math math.points math.geometry
   4 combinators.short-circuit arrays fry locals ;
   5 IN: quadtrees
   6
   7 TUPLE: quadtree { bounds rect } point value ll lr ul ur leaf? ;
   8
   9 : <quadtree> ( bounds -- quadtree ) f f f f f f t quadtree boa ;
  10
  11 : rect-ll ( rect -- point ) loc>> ;
  12 : rect-lr ( rect -- point ) [ loc>> ] [ width  ] bi v+x ;
  13 : rect-ul ( rect -- point ) [ loc>> ] [ height ] bi v+y ;
  14 : rect-ur ( rect -- point ) [ loc>> ] [ dim>>  ] bi v+  ;
  15
  16 : rect-center ( rect -- point ) [ loc>> ] [ dim>> 0.5 v*n ] bi v+ ; inline
  17
  18 : (quadrant) ( pt node -- quadrant )
  19     swap [ first 0.0 < ] [ second 0.0 < ] bi
  20     [ [ ll>> ] [ lr>> ] if ]
  21     [ [ ul>> ] [ ur>> ] if ] if ;
  22
  23 : quadrant ( pt node -- quadrant )
  24     [ bounds>> rect-center v- ] keep (quadrant) ;
  25
  26 : descend ( pt node -- pt subnode )
  27     [ drop ] [ quadrant ] 2bi ; inline
  28
  29 :: each-quadrant ( node quot -- )
  30     node ll>> quot call
  31     node lr>> quot call
  32     node ul>> quot call
  33     node ur>> quot call ; inline
  34 : map-quadrant ( node quot: ( child-node -- x ) -- array )
  35     each-quadrant 4array ; inline
  36
  37 <PRIVATE
  38
  39 DEFER: (prune)
  40 DEFER: insert
  41 DEFER: erase
  42 DEFER: at-point
  43 DEFER: quadtree>alist
  44 DEFER: quadtree-size
  45 DEFER: node-insert
  46 DEFER: in-rect*
  47
  48 : child-dim ( rect -- dim/2 ) dim>> 0.5 v*n ; inline
  49 : ll-bounds ( rect -- rect' )
  50     [   loc>>                                  ] [ child-dim ] bi <rect> ;
  51 : lr-bounds ( rect -- rect' )
  52     [ [ loc>> ] [ dim>> { 0.5 0.0 } v* ] bi v+ ] [ child-dim ] bi <rect> ;
  53 : ul-bounds ( rect -- rect' )
  54     [ [ loc>> ] [ dim>> { 0.0 0.5 } v* ] bi v+ ] [ child-dim ] bi <rect> ;
  55 : ur-bounds ( rect -- rect' )
  56     [ [ loc>> ] [ dim>> { 0.5 0.5 } v* ] bi v+ ] [ child-dim ] bi <rect> ;
  57
  58 : {quadrants} ( node -- quadrants )
  59     { [ ll>> ] [ lr>> ] [ ul>> ] [ ur>> ] } cleave 4array ;
  60
  61 : add-subnodes ( node -- node )
  62     dup bounds>> {
  63         [ ll-bounds <quadtree> >>ll ]
  64         [ lr-bounds <quadtree> >>lr ]
  65         [ ul-bounds <quadtree> >>ul ]
  66         [ ur-bounds <quadtree> >>ur ]
  67     } cleave
  68     f >>leaf? ;
  69
  70 : split-leaf ( value point leaf -- )
  71     add-subnodes
  72     [ value>> ] [ point>> ] [ ] tri
  73     [ node-insert ] [ node-insert ] bi ;
  74
  75 : leaf-replaceable? ( pt leaf -- ? ) point>> { [ nip not ] [ = ] } 2|| ;
  76 : leaf-insert ( value point leaf -- )
  77     2dup leaf-replaceable?
  78     [ [ (>>point) ] [ (>>value) ] bi ]
  79     [ split-leaf ] if ;
  80
  81 : node-insert ( value point node -- )
  82     descend insert ;
  83
  84 : insert ( value point tree -- )
  85     dup leaf?>> [ leaf-insert ] [ node-insert ] if ;
  86
  87 : leaf-at-point ( point leaf -- value/f ? )
  88     tuck point>> = [ value>> t ] [ drop f f ] if ;
  89
  90 : node-at-point ( point node -- value/f ? )
  91     descend at-point ;
  92
  93 : at-point ( point tree -- value/f ? )
  94     dup leaf?>> [ leaf-at-point ] [ node-at-point ] if ;
  95
  96 : (node-in-rect*) ( values rect node -- values )
  97     2dup bounds>> intersects? [ in-rect* ] [ 2drop ] if ;
  98 : node-in-rect* ( values rect node -- values )
  99     [ (node-in-rect*) ] with each-quadrant ;
 100
 101 : leaf-in-rect* ( values rect leaf -- values )
 102     tuck { [ nip point>> ] [ point>> swap intersects? ] } 2&&
 103     [ value>> over push ] [ drop ] if ;
 104
 105 : in-rect* ( values rect tree -- values )
 106     dup leaf?>> [ leaf-in-rect* ] [ node-in-rect* ] if ;
 107
 108 : leaf-erase ( point leaf -- )
 109     tuck point>> = [ f >>point f >>value ] when drop ;
 110
 111 : node-erase ( point node -- )
 112     descend erase ;
 113
 114 : erase ( point tree -- )
 115     dup leaf?>> [ leaf-erase ] [ node-erase ] if ;
 116
 117 : (?leaf) ( quadrant -- {point,value}/f )
 118     dup point>> [ swap value>> 2array ] [ drop f ] if* ;
 119 : ?leaf ( quadrants -- {point,value}/f )
 120     [ (?leaf) ] map sift dup length {
 121         { 1 [ first ] }
 122         { 0 [ drop { f f } ] }
 123         [ 2drop f ]
 124     } case ;
 125
 126 : collapseable? ( node -- {point,value}/f )
 127     {quadrants} { [ [ leaf?>> ] all? ] [ ?leaf ] } 1&& ;
 128
 129 : remove-subnodes ( node -- leaf ) f >>ll f >>lr f >>ul f >>ur t >>leaf? ;
 130
 131 : collapse ( node {point,value} -- )
 132     first2 [ >>point ] [ >>value ] bi* remove-subnodes drop ;
 133
 134 : node-prune ( node -- )
 135     [ [ (prune) ] each-quadrant ] [ ] [ collapseable? ] tri
 136     [ collapse ] [ drop ] if* ;
 137
 138 : (prune) ( tree -- )
 139     dup leaf?>> [ drop ] [ node-prune ] if ;
 140
 141 : leaf>alist ( leaf -- alist )
 142     dup point>> [ [ point>> ] [ value>> ] bi 2array 1array ] [ drop { } ] if ;
 143
 144 : node>alist ( node -- alist ) [ quadtree>alist ] map-quadrant concat ;
 145
 146 : quadtree>alist ( tree -- assoc )
 147     dup leaf?>> [ leaf>alist ] [ node>alist ] if ;
 148
 149 : leaf-size ( leaf -- count )
 150     point>> [ 1 ] [ 0 ] if ;
 151 : node-size ( node -- count )
 152     0 swap [ quadtree-size + ] each-quadrant ;
 153
 154 : quadtree-size ( tree -- count )
 155     dup leaf?>> [ leaf-size ] [ node-size ] if ;
 156
 157 : leaf= ( a b -- ? ) [ [ point>> ] [ value>> ] bi 2array ] bi@ = ;
 158
 159 : node= ( a b -- ? ) [ {quadrants} ] bi@ = ;
 160
 161 : (tree=) ( a b -- ? ) dup leaf?>> [ leaf= ] [ node= ] if ;
 162
 163 : tree= ( a b -- ? )
 164     2dup [ leaf?>> ] bi@ = [ (tree=) ] [ 2drop f ] if ;
 165
 166 PRIVATE>
 167
 168 : prune ( tree -- tree ) [ (prune) ] keep ;
 169
 170 : in-rect ( tree rect -- values )
 171     [ 16 <vector> ] 2dip in-rect* ;
 172
 173 M: quadtree equal? ( a b -- ? )
 174     over quadtree? [ tree= ] [ 2drop f ] if ;
 175
 176 INSTANCE: quadtree assoc
 177
 178 M: quadtree at* ( key assoc -- value/f ? ) at-point ;
 179 M: quadtree assoc-size ( assoc -- n ) quadtree-size ;
 180 M: quadtree >alist ( assoc -- alist ) quadtree>alist ;
 181 M: quadtree set-at ( value key assoc -- ) insert ;
 182 M: quadtree delete-at ( key assoc -- ) erase ;
 183 M: quadtree clear-assoc ( assoc -- )
 184     t >>leaf?
 185     f >>point
 186     f >>value
 187     drop ;
 188