basis/persistent/heaps/heaps.factor

   1 USING: kernel accessors locals combinators math arrays
   2 assocs namespaces sequences ;
   3 IN: persistent.heaps
   4 ! These are minheaps
   5
   6 <PRIVATE
   7 TUPLE: branch value prio left right ;
   8 TUPLE: empty-heap ;
   9
  10 PREDICATE: singleton-heap < branch
  11     [ left>> ] [ right>> ] bi [ empty-heap? ] both? ;
  12
  13 C: <branch> branch
  14 : >branch< ( branch -- value prio left right )
  15     { [ value>> ] [ prio>> ] [ left>> ] [ right>> ] } cleave ;
  16 PRIVATE>
  17
  18 : <persistent-heap> ( -- heap ) T{ empty-heap } ;
  19
  20 : <singleton-heap> ( value prio -- heap )
  21     <persistent-heap> <persistent-heap> <branch> ;
  22
  23 : pheap-empty? ( heap -- ? ) empty-heap? ;
  24
  25 : empty-pheap ( -- * )
  26     "Attempt to delete from an empty heap" throw ;
  27
  28 <PRIVATE
  29 : remove-left ( heap -- value prio newheap )
  30     dup [ left>> ] [ right>> ] bi [ pheap-empty? ] both?
  31     [ [ value>> ] [ prio>> ] bi <persistent-heap> ]
  32     [ >branch< swap remove-left -rot [ <branch> ] 2dip rot ] if ;
  33
  34 : both-with? ( obj a b quot -- ? )
  35    swap [ with ] dip swap both? ; inline
  36
  37 GENERIC: sift-down ( value prio left right -- heap )
  38
  39 : singleton-sift-down ( value prio singleton empty -- heap )
  40     2over prio>> <= [ <branch> ] [
  41         drop -rot [ [ value>> ] [ prio>> ] bi ] 2dip
  42         <singleton-heap> <persistent-heap> <branch>
  43     ] if ;
  44
  45 M: empty-heap sift-down
  46     over singleton-heap? [ singleton-sift-down ] [ <branch> ] if ;
  47
  48 :: reroot-left ( value prio left right -- heap )
  49     left value>> left prio>>
  50     value prio left left>> left right>> sift-down
  51     right <branch> ;
  52
  53 :: reroot-right ( value prio left right -- heap )
  54     right value>> right prio>> left
  55     value prio right left>> right right>> sift-down
  56     <branch> ;
  57
  58 M: branch sift-down ! both arguments are branches
  59     3dup [ prio>> <= ] both-with? [ <branch> ] [
  60         2dup [ prio>> ] bi@ <= [ reroot-left ] [ reroot-right ] if
  61     ] if ;
  62 PRIVATE>
  63
  64 GENERIC: pheap-peek ( heap -- value prio )
  65 M: empty-heap pheap-peek empty-pheap ;
  66 M: branch pheap-peek [ value>> ] [ prio>> ] bi ;
  67
  68 GENERIC: pheap-push ( value prio heap -- newheap )
  69
  70 M: empty-heap pheap-push
  71     drop <singleton-heap> ;
  72
  73 <PRIVATE
  74 : push-top ( value prio heap -- newheap )
  75     [ [ value>> ] [ prio>> ] [ right>> ] tri pheap-push ]
  76     [ left>> ] bi <branch> ;
  77
  78 : push-in ( value prio heap -- newheap )
  79     [ 2nip [ value>> ] [ prio>> ] bi ]
  80     [ right>> pheap-push ]
  81     [ 2nip left>> ] 3tri <branch> ;
  82 PRIVATE>
  83
  84 M: branch pheap-push
  85     2dup prio>> <= [ push-top ] [ push-in ] if ;
  86
  87 : pheap-pop* ( heap -- newheap )
  88     dup pheap-empty? [ empty-pheap ] [
  89         dup left>> pheap-empty?
  90         [ drop <persistent-heap> ]
  91         [ [ left>> remove-left ] keep right>> swap sift-down ] if
  92     ] if ;
  93
  94 : pheap-pop ( heap -- newheap value prio )
  95     [ pheap-pop* ] [ pheap-peek ] bi ;
  96
  97 : assoc>pheap ( assoc -- heap ) ! Assoc is value => prio
  98     <persistent-heap> swap [ rot pheap-push ] assoc-each ;
  99
 100 : pheap>alist ( heap -- alist )
 101     [ dup pheap-empty? not ] [ pheap-pop 2array ] produce nip ;
 102
 103 : pheap>values ( heap -- seq ) pheap>alist keys ;