]> gitweb.factorcode.org Git - factor.git/blobdiff - basis/heaps/heaps.factor
basis: use lint.vocabs tool to trim using lists
[factor.git] / basis / heaps / heaps.factor
index b92bfe36687087cc668f8ca19fe2f674acb2e038..2c6d4c8102b06ebf247a540cdfb92b7d2ba6a31d 100644 (file)
@@ -1,7 +1,7 @@
 ! Copyright (C) 2007, 2008 Ryan Murphy, Doug Coleman,
 ! Slava Pestov.
 ! See http://factorcode.org/license.txt for BSD license.
-USING: accessors arrays assocs fry kernel kernel.private locals
+USING: accessors arrays assocs fry kernel kernel.private
 math math.order math.private sequences sequences.private summary
 vectors ;
 IN: heaps
@@ -9,7 +9,6 @@ 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 )
@@ -21,11 +20,6 @@ TUPLE: heap { data vector } ;
 : <heap> ( class -- heap )
     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 ;
 
 : <entry> ( value key heap -- entry )
@@ -41,11 +35,9 @@ TUPLE: max-heap < heap ;
 
 : <max-heap> ( -- max-heap ) max-heap <heap> ;
 
-M: heap heap-empty? ( heap -- ? )
-    data>> empty? ; inline
+M: heap heap-empty? data>> empty? ; inline
 
-M: heap heap-size ( heap -- n )
-    data>> length ; inline
+M: heap heap-size data>> length ; inline
 
 <PRIVATE
 
@@ -58,18 +50,15 @@ M: heap heap-size ( heap -- n )
 : up ( n -- m )
     { fixnum } declare 1 fixnum-fast 2/ ; inline
 
-: data-nth ( n heap -- entry )
-    data>> nth-unsafe { entry } declare ; inline
-
-: data-first ( heap -- entry )
-    data>> first ; inline
+: data-nth ( n data -- entry )
+    nth-unsafe { entry } declare ; inline
 
-: data-set-nth ( entry n heap -- )
-    [ [ >>index ] keep ] dip data>> set-nth-unsafe ; inline
+: data-set-nth ( entry n data -- )
+    [ [ >>index ] keep ] dip set-nth-unsafe ; inline
 
-: data-push ( entry heap -- n )
-    [ heap-size [ >>index ] keep ]
-    [ data>> [ set-nth ] 2keep drop ] bi ; inline
+: data-push ( entry data -- n )
+    [ length [ >>index ] keep ]
+    [ [ set-nth ] keepd ] bi ; inline
 
 GENERIC: heap-compare ( entry1 entry2 heap -- ? )
 
@@ -79,38 +68,40 @@ M: min-heap heap-compare
 M: max-heap heap-compare
     drop { entry entry } declare [ key>> ] bi@ before? ; inline
 
-: data-compare ( m n heap -- ? )
-    [ '[ _ data-nth ] bi@ ] [ heap-compare ] bi ; inline
-
 PRIVATE>
 
 : >entry< ( entry -- value key )
     [ value>> ] [ key>> ] bi ; inline
 
-M: heap heap-peek ( heap -- value key )
-    data-first >entry< ;
+M: heap heap-peek
+    data>> first >entry< ;
 
 <PRIVATE
 
+! names and optimizations inspired by cpython/Lib/heapq.py,
+! refer to it from in depth explanations of the optimizations.
+
+! called bubble-up in the literature... but we keep cpython's name.
 :: sift-down ( heap from to -- )
-    to heap data-nth :> tmp
+    heap data>>      :> data
+    to data data-nth :> tmp
 
     to t [ over from > and ] [
         dup up
-        dup heap data-nth
+        dup data data-nth
         dup tmp heap heap-compare [
-            rot heap data-set-nth t
+            rot data data-set-nth t
         ] [
             2drop f
         ] if
     ] while
 
-    tmp swap heap data-set-nth ; inline
+    tmp swap data data-set-nth ; inline
 
 PRIVATE>
 
 M: heap heap-push*
-    [ <entry> dup ] [ data-push ] [ 0 rot sift-down ] tri ;
+    [ <entry> dup ] [ data>> data-push ] [ 0 rot sift-down ] tri ;
 
 : heap-push ( value key heap -- )
     heap-push* drop ;
@@ -120,35 +111,39 @@ M: heap heap-push*
 
 <PRIVATE
 
+! 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 heap-size  :> end
-    n heap data-nth :> tmp
+    heap data>>     :> data
+    data length     :> end
+    n data data-nth :> tmp
 
     n dup left [ dup end < ] [
         dup 1 fixnum+fast
-        dup end < [ 2dup heap data-compare ] [ f ] if
+        dup end < [
+            2dup [ data data-nth ] bi@ heap heap-compare
+        ] [ f ] if
         [ nip ] [ drop ] if
-        [ heap data-nth swap heap data-set-nth ]
+        [ data data-nth swap data data-set-nth ]
         [ dup left ] bi
     ] while drop
 
-    tmp over heap data-set-nth
+    tmp over data data-set-nth
     heap n rot sift-down ; inline
 
 PRIVATE>
 
 M: heap heap-pop*
-    dup data>> dup length 1 > [
-        [ pop ] [ set-first ] bi 0 sift-up
-    ] [
-        pop* drop
-    ] if ; inline
+    dup data>> f over first index<< [ pop ] keep
+    [ 2drop ] [ set-first 0 sift-up ] if-empty ;
 
-M: heap heap-pop
-    [ data-first >entry< ] [ heap-pop* ] bi ;
+: heap-pop ( heap -- value key )
+    [ heap-peek ] [ heap-pop* ] bi ;
 
 : slurp-heap ( ... heap quot: ( ... value key -- ... ) -- ... )
-    [ check-heap ] dip
     [ drop '[ _ heap-empty? ] ]
     [ '[ _ heap-pop @ ] until ] 2bi ; inline
 
@@ -165,19 +160,21 @@ M: bad-heap-delete summary
 
 : entry>index ( entry heap -- n )
     over heap>> eq? [ bad-heap-delete ] unless
-    index>> { fixnum } declare ; inline
+    index>> dup [ bad-heap-delete ] unless
+    { fixnum } declare ; inline
 
 PRIVATE>
 
-M: heap heap-delete
-    [ entry>index ] keep
-    2dup heap-size 1 - = [
-        nip data>> pop*
-    ] [
-        [ nip data>> pop ]
-        [ data-set-nth ]
-        [ swap sift-up ] 2tri
-    ] if ;
+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