! Copyright (C) 2007, 2008 Ryan Murphy, Doug Coleman,
! Slava Pestov.
! See http://factorcode.org/license.txt for BSD license.
-USING: kernel math sequences arrays assocs sequences.private
-growable accessors math.order summary ;
+USING: accessors arrays assocs fry kernel kernel.private
+math math.order math.private sequences sequences.private summary
+vectors ;
IN: heaps
GENERIC: heap-push* ( value key heap -- entry )
GENERIC: heap-peek ( heap -- value key )
GENERIC: heap-pop* ( heap -- )
-GENERIC: heap-pop ( heap -- value key )
GENERIC: heap-delete ( entry heap -- )
GENERIC: heap-empty? ( heap -- ? )
GENERIC: heap-size ( heap -- n )
<PRIVATE
-TUPLE: heap data ;
+TUPLE: heap { data vector } ;
: <heap> ( class -- heap )
- [ V{ } clone ] dip boa ; inline
+ V{ } clone swap boa ; inline
TUPLE: entry value key heap index ;
-: <entry> ( value key heap -- entry ) f entry boa ;
+: <entry> ( value key heap -- entry )
+ f entry boa ; inline
PRIVATE>
: <max-heap> ( -- max-heap ) max-heap <heap> ;
-M: heap heap-empty? ( heap -- ? )
- data>> empty? ;
+M: heap heap-empty? data>> empty? ; inline
-M: heap heap-size ( heap -- n )
- data>> length ;
+M: heap heap-size data>> length ; inline
<PRIVATE
-: left ( n -- m ) 1 shift 1 + ; inline
+: left ( n -- m )
+ { fixnum } declare 1 fixnum-shift-fast 1 fixnum+fast ; inline
-: right ( n -- m ) 1 shift 2 + ; inline
+: right ( n -- m )
+ { fixnum } declare 1 fixnum-shift-fast 2 fixnum+fast ; inline
-: up ( n -- m ) 1- 2/ ; inline
+: up ( n -- m )
+ { fixnum } declare 1 fixnum-fast 2/ ; inline
-: data-nth ( n heap -- entry )
- data>> nth-unsafe ; inline
+: data-nth ( n data -- entry )
+ nth-unsafe { entry } declare ; inline
-: up-value ( n heap -- entry )
- [ up ] dip data-nth ; inline
+: data-set-nth ( entry n data -- )
+ [ [ >>index ] keep ] dip set-nth-unsafe ; inline
-: left-value ( n heap -- entry )
- [ left ] dip data-nth ; inline
+: data-push ( entry data -- n )
+ [ length [ >>index ] keep ]
+ [ [ set-nth ] keepd ] bi ; inline
-: right-value ( n heap -- entry )
- [ right ] dip data-nth ; inline
+GENERIC: heap-compare ( entry1 entry2 heap -- ? )
-: data-set-nth ( entry n heap -- )
- [ [ >>index drop ] 2keep ] dip
- data>> set-nth-unsafe ; inline
+M: min-heap heap-compare
+ drop { entry entry } declare [ key>> ] bi@ after? ; inline
-: data-push ( entry heap -- n )
- dup heap-size [
- swap 2dup data>> ensure 2drop data-set-nth
- ] keep ; inline
+M: max-heap heap-compare
+ drop { entry entry } declare [ key>> ] bi@ before? ; inline
-: data-pop ( heap -- entry )
- data>> pop ; inline
-
-: data-pop* ( heap -- )
- data>> pop* ; inline
-
-: data-peek ( heap -- entry )
- data>> peek ; inline
-
-: data-first ( heap -- entry )
- data>> first ; inline
-
-: data-exchange ( m n heap -- )
- [ tuck data-nth [ data-nth ] dip ] 3keep
- tuck [ data-set-nth ] 2dip data-set-nth ; inline
-
-GENERIC: heap-compare ( pair1 pair2 heap -- ? )
-
-: (heap-compare) ( pair1 pair2 heap -- <=> )
- drop [ key>> ] compare ; inline
-
-M: min-heap heap-compare (heap-compare) +gt+ eq? ;
-
-M: max-heap heap-compare (heap-compare) +lt+ eq? ;
+PRIVATE>
-: heap-bounds-check? ( m heap -- ? )
- heap-size >= ; inline
+: >entry< ( entry -- value key )
+ [ value>> ] [ key>> ] bi ; inline
-: left-bounds-check? ( m heap -- ? )
- [ left ] dip heap-bounds-check? ; inline
+M: heap heap-peek
+ data>> first >entry< ;
-: right-bounds-check? ( m heap -- ? )
- [ right ] dip heap-bounds-check? ; inline
+<PRIVATE
-: continue? ( m up[m] heap -- ? )
- [ data-nth swap ] keep [ data-nth ] keep
- heap-compare ; inline
+! names and optimizations inspired by cpython/Lib/heapq.py,
+! refer to it from in depth explanations of the optimizations.
-DEFER: up-heap
+! called bubble-up in the literature... but we keep cpython's name.
+:: sift-down ( heap from to -- )
+ heap data>> :> data
+ to data data-nth :> tmp
-: (up-heap) ( n heap -- )
- [ dup up ] dip
- 3dup continue? [
- [ data-exchange ] 2keep up-heap
- ] [
- 3drop
- ] if ;
+ to t [ over from > and ] [
+ dup up
+ dup data data-nth
+ dup tmp heap heap-compare [
+ rot data data-set-nth t
+ ] [
+ 2drop f
+ ] if
+ ] while
-: up-heap ( n heap -- )
- over 0 > [ (up-heap) ] [ 2drop ] if ;
+ tmp swap data data-set-nth ; inline
-: (child) ( m heap -- n )
- 2dup right-value
- [ 2dup left-value ] dip
- rot heap-compare
- [ right ] [ left ] if ;
+PRIVATE>
-: child ( m heap -- n )
- 2dup right-bounds-check?
- [ drop left ] [ (child) ] if ;
+M: heap heap-push*
+ [ <entry> dup ] [ data>> data-push ] [ 0 rot sift-down ] tri ;
-: swap-down ( m heap -- )
- [ child ] 2keep data-exchange ;
+: heap-push ( value key heap -- )
+ heap-push* drop ;
-DEFER: down-heap
+: heap-push-all ( assoc heap -- )
+ '[ swap _ heap-push ] assoc-each ;
-: (down-heap) ( m heap -- )
- [ child ] 2keep swapd
- 3dup continue? [
- 3drop
- ] [
- [ data-exchange ] 2keep down-heap
- ] if ;
+<PRIVATE
-: down-heap ( m heap -- )
- 2dup left-bounds-check? [ 2drop ] [ (down-heap) ] if ;
+! called bubble-down in the literature... but we keep cpython's name.
+! A quote from cpython's implementation:
+! > We *could* break out of the loop as soon as we find a pos where newitem <=
+! > both its children, but turns out that's not a good idea [...]
+! Indeed the code is 33% slower if we remove this optimization.
+:: sift-up ( heap n -- )
+ heap data>> :> data
+ data length :> end
+ n data data-nth :> tmp
+
+ n dup left [ dup end < ] [
+ dup 1 fixnum+fast
+ dup end < [
+ 2dup [ data data-nth ] bi@ heap heap-compare
+ ] [ f ] if
+ [ nip ] [ drop ] if
+ [ data data-nth swap data data-set-nth ]
+ [ dup left ] bi
+ ] while drop
+
+ tmp over data data-set-nth
+ heap n rot sift-down ; inline
PRIVATE>
-M: heap heap-push* ( value key heap -- entry )
- [ <entry> dup ] keep [ data-push ] keep up-heap ;
-
-: heap-push ( value key heap -- ) heap-push* drop ;
+M: heap heap-pop*
+ dup data>> f over first index<< [ pop ] keep
+ [ 2drop ] [ set-first 0 sift-up ] if-empty ;
-: heap-push-all ( assoc heap -- )
- [ swapd heap-push ] curry assoc-each ;
+: heap-pop ( heap -- value key )
+ [ heap-peek ] [ heap-pop* ] bi ;
-: >entry< ( entry -- key value )
- [ value>> ] [ key>> ] bi ;
+: slurp-heap ( ... heap quot: ( ... value key -- ... ) -- ... )
+ [ drop '[ _ heap-empty? ] ]
+ [ '[ _ heap-pop @ ] until ] 2bi ; inline
-M: heap heap-peek ( heap -- value key )
- data-first >entry< ;
+: heap-pop-all ( heap -- alist )
+ [ heap-size <vector> ] keep
+ [ swap 2array suffix! ] slurp-heap { } like ;
ERROR: bad-heap-delete ;
-M: bad-heap-delete summary
+M: bad-heap-delete summary
drop "Invalid entry passed to heap-delete" ;
-
-: entry>index ( entry heap -- n )
- over heap>> eq? [ bad-heap-delete ] unless
- index>> ;
-M: heap heap-delete ( entry heap -- )
- [ entry>index ] keep
- 2dup heap-size 1- = [
- nip data-pop*
- ] [
- [ nip data-pop ] 2keep
- [ data-set-nth ] 2keep
- down-heap
- ] if ;
+<PRIVATE
-M: heap heap-pop* ( heap -- )
- dup data-first swap heap-delete ;
+: entry>index ( entry heap -- n )
+ over heap>> eq? [ bad-heap-delete ] unless
+ index>> dup [ bad-heap-delete ] unless
+ { fixnum } declare ; inline
-M: heap heap-pop ( heap -- value key )
- dup data-first [ swap heap-delete ] keep >entry< ;
+PRIVATE>
-: heap-pop-all ( heap -- alist )
- [ dup heap-empty? not ]
- [ dup heap-pop swap 2array ]
- produce nip ;
-
-: slurp-heap ( heap quot: ( elt -- ) -- )
- over heap-empty? [ 2drop ] [
- [ [ heap-pop drop ] dip call ] [ slurp-heap ] 2bi
- ] if ; inline recursive
+M:: heap heap-delete ( entry heap -- )
+ entry heap entry>index :> n
+ heap data>> :> data
+ data pop :> nth-entry
+ f entry index<<
+ n data length = [
+ nth-entry n data data-set-nth
+ n 0 = [ t ] [ nth-entry n up data data-nth heap heap-compare ] if
+ [ heap n sift-up ] [ heap 0 n sift-down ] if
+ ] unless ;
+
+: >min-heap ( assoc -- min-heap )
+ dup assoc-size <vector> min-heap boa
+ [ heap-push-all ] keep ;
+
+: >max-heap ( assoc -- max-heap )
+ dup assoc-size <vector> max-heap boa
+ [ heap-push-all ] keep ;
projects / factor.git / blobdiff