1 ! Copyright (C) 2007, 2008 Ryan Murphy, Doug Coleman,
3 ! See http://factorcode.org/license.txt for BSD license.
4 USING: accessors arrays assocs combinators fry growable kernel
5 kernel.private math math.order math.private sequences
6 sequences.private summary vectors ;
10 GENERIC: heap-push* ( value key heap -- entry )
11 GENERIC: heap-peek ( heap -- value key )
12 GENERIC: heap-pop* ( heap -- )
13 GENERIC: heap-pop ( heap -- value key )
14 GENERIC: heap-delete ( entry heap -- )
15 GENERIC: heap-empty? ( heap -- ? )
16 GENERIC: heap-size ( heap -- n )
20 TUPLE: heap { data vector } ;
22 : <heap> ( class -- heap )
23 [ V{ } clone ] dip boa ; inline
25 TUPLE: entry value key heap index ;
27 : <entry> ( value key heap -- entry )
32 TUPLE: min-heap < heap ;
34 : <min-heap> ( -- min-heap ) min-heap <heap> ;
36 TUPLE: max-heap < heap ;
38 : <max-heap> ( -- max-heap ) max-heap <heap> ;
40 M: heap heap-empty? ( heap -- ? )
41 data>> empty? ; inline
43 M: heap heap-size ( heap -- n )
44 data>> length ; inline
49 { fixnum } declare 1 fixnum-shift-fast 1 fixnum+fast ; inline
52 { fixnum } declare 1 fixnum-shift-fast 2 fixnum+fast ; inline
55 { fixnum } declare 1 fixnum-fast 2/ ; inline
57 : data-nth ( n heap -- entry )
58 data>> nth-unsafe { entry } declare ; inline
60 : left-value ( n heap -- entry )
61 [ left ] dip data-nth ; inline
63 : right-value ( n heap -- entry )
64 [ right ] dip data-nth ; inline
66 : data-set-nth ( entry n heap -- )
67 [ [ swap index<< ] 2keep ] dip
68 data>> set-nth-unsafe ; inline
70 : data-push ( entry heap -- n )
73 [ data>> ensure 2drop ]
77 : data-first ( heap -- entry )
80 : data-exchange ( m n heap -- )
81 [ '[ _ data-nth ] bi@ ]
82 [ '[ _ data-set-nth ] bi@ ] 3bi ; inline
84 GENERIC: heap-compare ( entry1 entry2 heap -- ? )
86 M: min-heap heap-compare
87 drop { entry entry } declare [ key>> ] bi@ after? ; inline
89 M: max-heap heap-compare
90 drop { entry entry } declare [ key>> ] bi@ before? ; inline
92 : heap-bounds-check? ( m heap -- ? )
95 : left-bounds-check? ( m heap -- ? )
96 [ left ] dip heap-bounds-check? ; inline
98 : right-bounds-check? ( m heap -- ? )
99 [ right ] dip heap-bounds-check? ; inline
101 : continue? ( m n heap -- ? )
104 [ 2nip ] 3tri heap-compare ; inline
108 : (up-heap) ( n heap -- )
111 [ data-exchange ] [ up-heap ] 2bi
114 ] if ; inline recursive
116 : up-heap ( n heap -- )
117 over 0 > [ (up-heap) ] [ 2drop ] if ; inline recursive
119 : (child) ( m heap -- n )
120 { [ drop ] [ left-value ] [ right-value ] [ nip ] } 2cleave
121 heap-compare [ right ] [ left ] if ; inline
123 : child ( m heap -- n )
124 2dup right-bounds-check?
125 [ drop left ] [ (child) ] if ; inline
129 : (down-heap) ( m heap -- )
130 [ drop ] [ child ] [ nip ] 2tri
134 [ data-exchange ] [ down-heap ] 2bi
135 ] if ; inline recursive
137 : down-heap ( m heap -- )
138 2dup left-bounds-check?
139 [ 2drop ] [ (down-heap) ] if ; inline recursive
143 M: heap heap-push* ( value key heap -- entry )
144 [ <entry> dup ] [ data-push ] [ up-heap ] tri ;
146 : heap-push ( value key heap -- ) heap-push* drop ;
148 : heap-push-all ( assoc heap -- )
149 '[ swap _ heap-push ] assoc-each ;
151 : >entry< ( entry -- value key )
152 [ value>> ] [ key>> ] bi ; inline
154 M: heap heap-peek ( heap -- value key )
157 ERROR: bad-heap-delete ;
159 M: bad-heap-delete summary
160 drop "Invalid entry passed to heap-delete" ;
162 : entry>index ( entry heap -- n )
163 over heap>> eq? [ bad-heap-delete ] unless
164 index>> { fixnum } declare ; inline
166 M: heap heap-delete ( entry heap -- )
168 2dup heap-size 1 - = [
176 M: heap heap-pop* ( heap -- )
177 [ data-first ] keep heap-delete ;
179 M: heap heap-pop ( heap -- value key )
180 [ data-first dup ] keep heap-delete >entry< ;
182 : heap-pop-all ( heap -- alist )
183 [ dup heap-empty? not ]
184 [ dup heap-pop swap 2array ]
187 ERROR: not-a-heap obj ;
189 : check-heap ( heap -- heap )
190 dup heap? [ not-a-heap ] unless ; inline
192 : slurp-heap ( heap quot: ( elt -- ) -- )
193 [ check-heap ] dip over heap-empty? [ 2drop ] [
194 [ [ heap-pop drop ] dip call ] [ slurp-heap ] 2bi
195 ] if ; inline recursive
197 : >min-heap ( assoc -- min-heap )
198 <min-heap> [ heap-push-all ] keep ;
200 : >max-heap ( assoc -- max-heap )
201 <max-heap> [ heap-push-all ] keep ;