diff --git a/basis/heaps/heaps-tests.factor b/basis/heaps/heaps-tests.factor index 703cf53080..64871a69e5 100644 --- a/basis/heaps/heaps-tests.factor +++ b/basis/heaps/heaps-tests.factor @@ -2,7 +2,7 @@ ! See http://factorcode.org/license.txt for BSD license. USING: arrays kernel math namespaces tools.test heaps heaps.private math.parser random assocs sequences sorting -accessors math.order ; +accessors math.order locals ; IN: heaps.tests [ heap-pop ] must-fail @@ -27,19 +27,31 @@ IN: heaps.tests [ 0 ] [ heap-size ] unit-test [ 1 ] [ t 1 pick heap-push heap-size ] unit-test -: heap-sort ( alist -- keys ) - [ heap-push-all ] keep heap-pop-all ; +: heap-sort ( alist heap -- keys ) + [ heap-push-all ] keep heap-pop-all ; : random-alist ( n -- alist ) iota [ drop 32 random-bits dup number>string - ] H{ } map>assoc ; + ] H{ } map>assoc >alist ; -: test-heap-sort ( n -- ? ) - random-alist dup >alist sort-keys swap heap-sort = ; +:: test-heap-sort ( n heap reverse? -- ? ) + n random-alist + [ sort-keys reverse? [ reverse ] when ] keep + heap heap-sort = ; + +: test-minheap-sort ( n -- ? ) + f test-heap-sort ; + +: test-maxheap-sort ( n -- ? ) + t test-heap-sort ; 14 [ - [ t ] swap [ 2^ test-heap-sort ] curry unit-test + [ t ] swap [ 2^ f test-heap-sort ] curry unit-test +] each-integer + +14 [ + [ t ] swap [ 2^ t test-heap-sort ] curry unit-test ] each-integer : test-entry-indices ( n -- ? ) diff --git a/basis/heaps/heaps.factor b/basis/heaps/heaps.factor index 28d18cb53a..326266773b 100644 --- a/basis/heaps/heaps.factor +++ b/basis/heaps/heaps.factor @@ -2,7 +2,7 @@ ! Slava Pestov. ! See http://factorcode.org/license.txt for BSD license. USING: kernel math sequences arrays assocs sequences.private -growable accessors math.order summary vectors ; +growable accessors math.order summary vectors fry combinators ; IN: heaps GENERIC: heap-push* ( value key heap -- entry ) @@ -58,30 +58,25 @@ M: heap heap-size ( heap -- n ) [ right ] dip data-nth ; inline : data-set-nth ( entry n heap -- ) - [ [ >>index drop ] 2keep ] dip + [ [ >>index drop ] [ ] 2bi ] dip data>> set-nth-unsafe ; inline : data-push ( entry heap -- n ) dup heap-size [ swap 2dup data>> ensure 2drop data-set-nth - ] keep ; inline - -: data-pop ( heap -- entry ) - data>> pop ; inline - -: data-pop* ( heap -- ) - data>> pop* ; inline + ] [ + ] bi ; inline : data-first ( heap -- entry ) data>> first ; inline : data-exchange ( m n heap -- ) - [ [ data-nth ] curry bi@ ] - [ [ data-set-nth ] curry bi@ ] 3bi ; inline + [ '[ _ data-nth ] bi@ ] + [ '[ _ data-set-nth ] bi@ ] 3bi ; inline -GENERIC: heap-compare ( pair1 pair2 heap -- ? ) +GENERIC: heap-compare ( entry1 entry2 heap -- ? ) -: (heap-compare) ( pair1 pair2 heap -- <=> ) +: (heap-compare) ( entry1 entry2 heap -- <=> ) drop [ key>> ] compare ; inline M: min-heap heap-compare (heap-compare) +gt+ eq? ; @@ -97,16 +92,17 @@ M: max-heap heap-compare (heap-compare) +lt+ eq? ; : right-bounds-check? ( m heap -- ? ) [ right ] dip heap-bounds-check? ; inline -: continue? ( m up[m] heap -- ? ) - [ data-nth swap ] keep [ data-nth ] keep - heap-compare ; inline +: continue? ( m n heap -- ? ) + [ data-nth nip ] + [ nip data-nth ] + [ 2nip ] 3tri heap-compare ; DEFER: up-heap : (up-heap) ( n heap -- ) [ dup up ] dip 3dup continue? [ - [ data-exchange ] 2keep up-heap + [ data-exchange ] [ up-heap ] 2bi ] [ 3drop ] if ; inline recursive @@ -115,10 +111,8 @@ DEFER: up-heap over 0 > [ (up-heap) ] [ 2drop ] if ; inline recursive : (child) ( m heap -- n ) - 2dup right-value - [ 2dup left-value ] dip - rot heap-compare - [ right ] [ left ] if ; + { [ drop ] [ left-value ] [ right-value ] [ nip ] } 2cleave + heap-compare [ right ] [ left ] if ; : child ( m heap -- n ) 2dup right-bounds-check? @@ -127,11 +121,11 @@ DEFER: up-heap DEFER: down-heap : (down-heap) ( m heap -- ) - [ child ] 2keep swapd + [ drop ] [ child ] [ nip ] 2tri 3dup continue? [ 3drop ] [ - [ data-exchange ] 2keep down-heap + [ data-exchange ] [ down-heap ] 2bi ] if ; inline recursive : down-heap ( m heap -- ) @@ -140,14 +134,14 @@ DEFER: down-heap PRIVATE> M: heap heap-push* ( value key heap -- entry ) - [ dup ] keep [ data-push ] keep up-heap ; + [ dup ] [ data-push ] [ ] tri up-heap ; : heap-push ( value key heap -- ) heap-push* drop ; : heap-push-all ( assoc heap -- ) - [ swapd heap-push ] curry assoc-each ; + '[ swap _ heap-push ] assoc-each ; -: >entry< ( entry -- key value ) +: >entry< ( entry -- value key ) [ value>> ] [ key>> ] bi ; inline M: heap heap-peek ( heap -- value key ) @@ -163,29 +157,28 @@ M: bad-heap-delete summary index>> ; M: heap heap-delete ( entry heap -- ) - [ entry>index ] keep + [ entry>index ] [ ] bi 2dup heap-size 1 - = [ - nip data-pop* + nip data>> pop* ] [ - [ nip data-pop ] 2keep - [ data-set-nth ] 2keep + [ nip data>> pop ] + [ data-set-nth ] + [ ] 2tri down-heap ] if ; M: heap heap-pop* ( heap -- ) - dup data-first swap heap-delete ; + [ data-first ] keep heap-delete ; M: heap heap-pop ( heap -- value key ) - dup data-first [ swap heap-delete ] keep >entry< ; + [ data-first ] keep + [ heap-delete ] [ drop ] 2bi >entry< ; : heap-pop-all ( heap -- alist ) [ dup heap-empty? not ] [ dup heap-pop swap 2array ] produce nip ; -: heap-values ( heap -- alist ) - data>> [ value>> ] { } map-as ; - : slurp-heap ( heap quot: ( elt -- ) -- ) over heap-empty? [ 2drop ] [ [ [ heap-pop drop ] dip call ] [ slurp-heap ] 2bi