extra/quadtrees/quadtrees.factor

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