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 fry ;
6 TUPLE: quadtree { bounds rect } point value ll lr ul ur leaf? ;
8 : <quadtree> ( bounds -- quadtree )
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+ ;
18 : rect-center ( rect -- point ) [ loc>> ] [ dim>> 0.5 v*n ] bi v+ ; inline
20 : (quadrant) ( pt node -- quadrant )
21 swap [ first 0.0 < ] [ second 0.0 < ] bi
22 [ [ ll>> ] [ lr>> ] if ]
23 [ [ ul>> ] [ ur>> ] if ] if ;
25 : quadrant ( pt node -- quadrant )
26 [ bounds>> rect-center v- ] keep (quadrant) ;
28 : descend ( pt node -- pt subnode )
29 [ drop ] [ quadrant ] 2bi ; inline
31 : each-quadrant ( node quot -- )
33 [ [ ll>> ] [ call ] bi* ]
34 [ [ lr>> ] [ call ] bi* ]
35 [ [ ul>> ] [ call ] bi* ]
36 [ [ ur>> ] [ call ] bi* ]
38 : map-quadrant ( node quot: ( child-node -- x ) -- array )
39 each-quadrant 4array ; inline
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> ;
62 : {quadrants} ( node -- quadrants )
63 { [ ll>> ] [ lr>> ] [ ul>> ] [ ur>> ] } cleave 4array ;
65 : add-subnodes ( node -- node )
67 [ ll-bounds <quadtree> >>ll ]
68 [ lr-bounds <quadtree> >>lr ]
69 [ ul-bounds <quadtree> >>ul ]
70 [ ur-bounds <quadtree> >>ur ]
74 : split-leaf ( value point leaf -- )
76 [ value>> ] [ point>> ] [ ] tri
77 [ node-insert ] [ node-insert ] bi ;
79 : leaf-replaceable? ( pt leaf -- ? ) point>> { [ nip not ] [ = ] } 2|| ;
81 : leaf-insert ( value point leaf -- )
82 2dup leaf-replaceable?
83 [ [ (>>point) ] [ (>>value) ] bi ]
86 : node-insert ( value point node -- )
89 : insert ( value point tree -- )
90 dup leaf?>> [ leaf-insert ] [ node-insert ] if ;
92 : leaf-at-point ( point leaf -- value/f ? )
93 tuck point>> = [ value>> t ] [ drop f f ] if ;
95 : node-at-point ( point node -- value/f ? )
98 : at-point ( point tree -- value/f ? )
99 dup leaf?>> [ leaf-at-point ] [ node-at-point ] if ;
101 : (node-in-rect*) ( values rect node -- values )
102 2dup bounds>> contains-rect? [ in-rect* ] [ 2drop ] if ;
103 : node-in-rect* ( values rect node -- values )
104 [ (node-in-rect*) ] with each-quadrant ;
106 : leaf-in-rect* ( values rect leaf -- values )
107 tuck { [ nip point>> ] [ point>> swap contains-point? ] } 2&&
108 [ value>> over push ] [ drop ] if ;
110 : in-rect* ( values rect tree -- values )
111 dup leaf?>> [ leaf-in-rect* ] [ node-in-rect* ] if ;
113 : leaf-erase ( point leaf -- )
114 tuck point>> = [ f >>point f >>value ] when drop ;
116 : node-erase ( point node -- )
119 : erase ( point tree -- )
120 dup leaf?>> [ leaf-erase ] [ node-erase ] if ;
122 : (?leaf) ( quadrant -- {point,value}/f )
123 dup point>> [ swap value>> 2array ] [ drop f ] if* ;
124 : ?leaf ( quadrants -- {point,value}/f )
125 [ (?leaf) ] map sift dup length {
127 { 0 [ drop { f f } ] }
131 : collapseable? ( node -- {point,value}/f )
132 {quadrants} { [ [ leaf?>> ] all? ] [ ?leaf ] } 1&& ;
134 : remove-subnodes ( node -- leaf ) f >>ll f >>lr f >>ul f >>ur t >>leaf? ;
136 : collapse ( node {point,value} -- )
137 first2 [ >>point ] [ >>value ] bi* remove-subnodes drop ;
139 : node-prune ( node -- )
140 [ [ (prune) ] each-quadrant ] [ ] [ collapseable? ] tri
141 [ collapse ] [ drop ] if* ;
143 : (prune) ( tree -- )
144 dup leaf?>> [ drop ] [ node-prune ] if ;
146 : leaf>alist ( leaf -- alist )
147 dup point>> [ [ point>> ] [ value>> ] bi 2array 1array ] [ drop { } ] if ;
149 : node>alist ( node -- alist ) [ quadtree>alist ] map-quadrant concat ;
151 : quadtree>alist ( tree -- assoc )
152 dup leaf?>> [ leaf>alist ] [ node>alist ] if ;
154 : leaf-size ( leaf -- count )
155 point>> [ 1 ] [ 0 ] if ;
156 : node-size ( node -- count )
157 0 swap [ quadtree-size + ] each-quadrant ;
159 : quadtree-size ( tree -- count )
160 dup leaf?>> [ leaf-size ] [ node-size ] if ;
162 : leaf= ( a b -- ? ) [ [ point>> ] [ value>> ] bi 2array ] bi@ = ;
164 : node= ( a b -- ? ) [ {quadrants} ] bi@ = ;
166 : (tree=) ( a b -- ? ) dup leaf?>> [ leaf= ] [ node= ] if ;
169 2dup [ leaf?>> ] bi@ = [ (tree=) ] [ 2drop f ] if ;
173 : prune-quadtree ( tree -- tree ) [ (prune) ] keep ;
175 : in-rect ( tree rect -- values )
176 [ 16 <vector> ] 2dip in-rect* ;
178 M: quadtree equal? ( a b -- ? )
179 over quadtree? [ tree= ] [ 2drop f ] if ;
181 INSTANCE: quadtree assoc
183 M: quadtree at* ( key assoc -- value/f ? ) at-point ;
184 M: quadtree assoc-size ( assoc -- n ) quadtree-size ;
185 M: quadtree >alist ( assoc -- alist ) quadtree>alist ;
186 M: quadtree set-at ( value key assoc -- ) insert ;
187 M: quadtree delete-at ( key assoc -- ) erase ;
188 M: quadtree clear-assoc ( assoc -- )
194 : swizzle ( sequence quot -- sequence' )
196 [ zip ] [ rect-containing <quadtree> ] bi
197 [ '[ first2 _ set-at ] each ] [ values ] bi ; inline