Clean up heaps implementation, test max-heaps.

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
 
 [ <min-heap> heap-pop ] must-fail
@@ -27,19 +27,31 @@ IN: heaps.tests
 [ 0 ] [ <max-heap> heap-size ] unit-test
 [ 1 ] [ <max-heap> t 1 pick heap-push heap-size ] unit-test
 
-: heap-sort ( alist -- keys )
-    <min-heap> [ 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 -- ? )
+    <min-heap> f test-heap-sort ;
+
+: test-maxheap-sort ( n -- ? )
+    <max-heap> t test-heap-sort ;
 
 14 [
-    [ t ] swap [ 2^ test-heap-sort ] curry unit-test
+    [ t ] swap [ 2^ <min-heap> f test-heap-sort ] curry unit-test
+] each-integer
+
+14 [
+    [ t ] swap [ 2^ <max-heap> t test-heap-sort ] curry unit-test
 ] each-integer
 
 : test-entry-indices ( n -- ? )

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 )
-    [ <entry> dup ] keep [ data-push ] keep up-heap ;
+    [ <entry> 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