heaps: faster heaps, simpler heapsort.

factor · Dec 1, 2014 · 3aa14d7 · 3aa14d7
1 parent 937f575
commit 3aa14d7
Show file tree

Hide file tree

Showing 2 changed files with 89 additions and 115 deletions.
diff --git a/basis/heaps/heaps.factor b/basis/heaps/heaps.factor
@@ -1,10 +1,9 @@
 ! Copyright (C) 2007, 2008 Ryan Murphy, Doug Coleman,
 ! Slava Pestov.
 ! See http://factorcode.org/license.txt for BSD license.
-USING: accessors arrays assocs combinators fry growable kernel
-kernel.private math math.order math.private sequences
-sequences.private summary vectors ;
-
+USING: accessors arrays assocs fry kernel kernel.private locals
+math math.order math.private sequences sequences.private summary
+vectors ;
 IN: heaps
 
 GENERIC: heap-push* ( value key heap -- entry )
@@ -20,7 +19,12 @@ GENERIC: heap-size ( heap -- n )
 TUPLE: heap { data vector } ;
 
 : <heap> ( class -- heap )
-    [ V{ } clone ] dip boa ; inline
+    V{ } clone swap boa ; inline
+
+ERROR: not-a-heap object ;
+
+: check-heap ( heap -- heap )
+    dup heap? [ not-a-heap ] unless ; inline
 
 TUPLE: entry value key heap index ;
 
@@ -57,29 +61,15 @@ M: heap heap-size ( heap -- n )
 : data-nth ( n heap -- entry )
     data>> nth-unsafe { entry } declare ; inline
 
-: left-value ( n heap -- entry )
-    [ left ] dip data-nth ; inline
-
-: right-value ( n heap -- entry )
-    [ right ] dip data-nth ; inline
+: data-first ( heap -- entry )
+    0 swap data-nth ; inline
 
 : data-set-nth ( entry n heap -- )
-    [ [ swap index<< ] 2keep ] dip
-    data>> set-nth-unsafe ; inline
+    [ [ >>index ] keep ] dip data>> set-nth-unsafe ; inline
 
 : data-push ( entry heap -- n )
-    dup heap-size [
-        swap
-        [ data>> ensure 2drop ]
-        [ data-set-nth ] 2bi
-    ] keep ; inline
-
-: data-first ( heap -- entry )
-    data>> first ; inline
-
-: data-exchange ( m n heap -- )
-    [ '[ _ data-nth ] bi@ ]
-    [ '[ _ data-set-nth ] bi@ ] 3bi ; inline
+    [ heap-size [ >>index ] keep ]
+    [ data>> [ set-nth ] 2keep drop ] bi ; inline
 
 GENERIC: heap-compare ( entry1 entry2 heap -- ? )
 
@@ -89,111 +79,105 @@ M: min-heap heap-compare
 M: max-heap heap-compare
     drop { entry entry } declare [ key>> ] bi@ before? ; inline
 
-: heap-bounds-check? ( m heap -- ? )
-    heap-size >= ; inline
+: data-compare ( m n heap -- ? )
+    [ '[ _ data-nth ] bi@ ] [ heap-compare ] bi ; inline
 
-: left-bounds-check? ( m heap -- ? )
-    [ left ] dip heap-bounds-check? ; inline
+PRIVATE>
 
-: right-bounds-check? ( m heap -- ? )
-    [ right ] dip heap-bounds-check? ; inline
+: >entry< ( entry -- value key )
+    [ value>> ] [ key>> ] bi ; inline
 
-: continue? ( m n heap -- ? )
-    [ data-nth nip ]
-    [ nip data-nth ]
-    [ 2nip ] 3tri heap-compare ; inline
+M: heap heap-peek ( heap -- value key )
+    data-first >entry< ;
 
-DEFER: up-heap
+<PRIVATE
 
-: (up-heap) ( n heap -- )
-    [ dup up ] dip
-    3dup continue? [
-        [ data-exchange ] [ up-heap ] 2bi
-    ] [
-        3drop
-    ] if ; inline recursive
+:: sift-down ( heap from to -- )
+    to heap data-nth :> tmp
 
-: up-heap ( n heap -- )
-    over 0 > [ (up-heap) ] [ 2drop ] if ; inline recursive
+    to t [ over from > and ] [
+        dup up
+        dup heap data-nth
+        dup tmp heap heap-compare [
+            rot heap data-set-nth t
+        ] [
+            2drop f
+        ] if
+    ] while
 
-: (child) ( m heap -- n )
-    { [ drop ] [ left-value ] [ right-value ] [ nip ] } 2cleave
-    heap-compare [ right ] [ left ] if ; inline
+    tmp swap heap data-set-nth ; inline
 
-: child ( m heap -- n )
-    2dup right-bounds-check?
-    [ drop left ] [ (child) ] if ; inline
+PRIVATE>
 
-DEFER: down-heap
+M: heap heap-push*
+    [ <entry> dup ] [ data-push ] [ 0 rot sift-down ] tri ;
 
-: (down-heap) ( m heap -- )
-    [ drop ] [ child ] [ nip ] 2tri
-    3dup continue? [
-        3drop
-    ] [
-        [ data-exchange ] [ down-heap ] 2bi
-    ] if ; inline recursive
+: heap-push ( value key heap -- )
+    heap-push* drop ;
 
-: down-heap ( m heap -- )
-    2dup left-bounds-check?
-    [ 2drop ] [ (down-heap) ] if ; inline recursive
+: heap-push-all ( assoc heap -- )
+    '[ swap _ heap-push ] assoc-each ;
 
-PRIVATE>
+<PRIVATE
 
-M: heap heap-push* ( value key heap -- entry )
-    [ <entry> dup ] [ data-push ] [ up-heap ] tri ;
+:: sift-up ( heap n -- )
+    heap heap-size  :> end
+    n heap data-nth :> tmp
 
-: heap-push ( value key heap -- ) heap-push* drop ;
+    n dup left [ dup end < ] [
+        dup 1 fixnum+fast
+        dup end < [ 2dup heap data-compare ] [ f ] if
+        [ nip ] [ drop ] if
+        [ heap data-nth swap heap data-set-nth ]
+        [ dup left ] bi
+    ] while drop
 
-: heap-push-all ( assoc heap -- )
-    '[ swap _ heap-push ] assoc-each ;
+    tmp over heap data-set-nth
+    heap n rot sift-down ; inline
 
-: >entry< ( entry -- value key )
-    [ value>> ] [ key>> ] bi ; inline
+PRIVATE>
 
-M: heap heap-peek ( heap -- value key )
-    data-first >entry< ;
+M: heap heap-pop*
+    dup data>> dup length 1 > [
+        [ pop ] [ set-first ] bi 0 sift-up
+    ] [
+        pop* drop
+    ] if ; inline
+
+M: heap heap-pop
+    [ data-first >entry< ] [ heap-pop* ] bi ;
+
+: heap-pop-all ( heap -- alist )
+    check-heap [ heap-size ] keep
+    '[ _ heap-pop swap 2array ] replicate ;
+
+: slurp-heap ( heap quot: ( value key -- ) -- )
+    [ check-heap [ heap-size ] keep ] dip
+    '[ _ heap-pop @ ] times ; inline
 
 ERROR: bad-heap-delete ;
 
 M: bad-heap-delete summary
     drop "Invalid entry passed to heap-delete" ;
 
+<PRIVATE
+
 : entry>index ( entry heap -- n )
     over heap>> eq? [ bad-heap-delete ] unless
     index>> { fixnum } declare ; inline
 
-M: heap heap-delete ( entry heap -- )
+PRIVATE>
+
+M: heap heap-delete
     [ entry>index ] keep
     2dup heap-size 1 - = [
         nip data>> pop*
     ] [
         [ nip data>> pop ]
         [ data-set-nth ]
-        [ down-heap ] 2tri
+        [ swap sift-up ] 2tri
     ] if ;
 
-M: heap heap-pop* ( heap -- )
-    [ data-first ] keep heap-delete ;
-
-M: heap heap-pop ( heap -- value key )
-    [ data-first dup ] keep heap-delete >entry< ;
-
-: heap-pop-all ( heap -- alist )
-    [ dup heap-empty? not ]
-    [ dup heap-pop swap 2array ]
-    produce nip ;
-
-ERROR: not-a-heap object ;
-
-: check-heap ( heap -- heap )
-    dup heap? [ not-a-heap ] unless ; inline
-
-: slurp-heap ( heap quot: ( value key -- ) -- )
-    [ check-heap ] dip over heap-empty? [ 2drop ] [
-        [ [ heap-pop ] dip call ] [ slurp-heap ] 2bi
-    ] if ; inline recursive
-
 : >min-heap ( assoc -- min-heap )
     <min-heap> [ heap-push-all ] keep ;
 

diff --git a/extra/sorting/heap/heap.factor b/extra/sorting/heap/heap.factor
@@ -1,29 +1,19 @@
 ! Copyright (C) 2014 John Benediktsson
 ! See http://factorcode.org/license.txt for BSD license
 
-USING: assocs heaps kernel sequences ;
+USING: assocs fry heaps kernel sequences vectors ;
 
 IN: sorting.heap
 
-<PRIVATE
-
-: (heapsort) ( alist accum -- sorted-seq )
-    [ >min-heap ] [ [ [ nip push ] curry slurp-heap ] keep ] bi* ; inline
-
-PRIVATE>
-
-: heapsort ( seq -- sorted-seq )
-    [
-        [ dup zip ]
-        [ length ]
-        [ new-resizable ] tri
-        (heapsort)
-    ] [ like ] bi ;
-
 : heapsort-with ( seq quot: ( elt -- key ) -- sorted-seq )
     [
-        [ keep ] curry [ { } map>assoc ] curry
-        [ length ]
-        [ new-resizable ] tri
-        (heapsort)
-    ] 2keep drop like ; inline
+        over length <vector> min-heap boa
+        [ '[ dup @ _ heap-push ] each ] keep
+    ] [
+        drop [ length ] keep new-resizable
+        [ '[ drop _ push ] slurp-heap ] keep
+    ] [
+        drop like
+    ] 2tri ; inline
+
+: heapsort ( seq -- sorted-seq ) [ ] heapsort-with ;