1 USING: kernel accessors locals combinators math arrays
2 assocs namespaces sequences ;
7 TUPLE: branch value prio left right ;
10 PREDICATE: singleton-heap < branch
11 [ left>> ] [ right>> ] bi [ empty-heap? ] both? ;
14 : >branch< ( branch -- value prio left right )
15 { [ value>> ] [ prio>> ] [ left>> ] [ right>> ] } cleave ;
18 : <persistent-heap> ( -- heap ) T{ empty-heap } ;
20 : <singleton-heap> ( value prio -- heap )
21 <persistent-heap> <persistent-heap> <branch> ;
23 : pheap-empty? ( heap -- ? ) empty-heap? ;
25 : empty-pheap ( -- * )
26 "Attempt to delete from an empty heap" throw ;
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 ;
34 : both-with? ( obj a b quot -- ? )
35 swap >r with r> swap both? ; inline
37 GENERIC: sift-down ( value prio left right -- heap )
39 : singleton-sift-down ( value prio singleton empty -- heap )
40 3dup drop prio>> <= [ <branch> ] [
41 drop -rot [ [ value>> ] [ prio>> ] bi ] 2dip
42 <singleton-heap> <persistent-heap> <branch>
45 M: empty-heap sift-down
46 over singleton-heap? [ singleton-sift-down ] [ <branch> ] if ;
48 :: reroot-left ( value prio left right -- heap )
49 left value>> left prio>>
50 value prio left left>> left right>> sift-down
53 :: reroot-right ( value prio left right -- heap )
54 right value>> right prio>> left
55 value prio right left>> right right>> sift-down
58 M: branch sift-down ! both arguments are branches
59 3dup [ prio>> <= ] both-with? [ <branch> ] [
60 2dup [ prio>> ] bi@ <= [ reroot-left ] [ reroot-right ] if
64 GENERIC: pheap-peek ( heap -- value prio )
65 M: empty-heap pheap-peek empty-pheap ;
66 M: branch pheap-peek [ value>> ] [ prio>> ] bi ;
68 GENERIC: pheap-push ( value prio heap -- newheap )
70 M: empty-heap pheap-push
71 drop <singleton-heap> ;
74 : push-top ( value prio heap -- newheap )
75 [ [ value>> ] [ prio>> ] [ right>> ] tri pheap-push ]
76 [ left>> ] bi <branch> ;
78 : push-in ( value prio heap -- newheap )
79 [ 2nip [ value>> ] [ prio>> ] bi ]
80 [ right>> pheap-push ]
81 [ 2nip left>> ] 3tri <branch> ;
85 2dup prio>> <= [ push-top ] [ push-in ] if ;
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
94 : pheap-pop ( heap -- newheap value prio )
95 [ pheap-pop* ] [ pheap-peek ] bi ;
97 : assoc>pheap ( assoc -- heap ) ! Assoc is value => prio
98 <persistent-heap> swap [ rot pheap-push ] assoc-each ;
100 : pheap>alist ( heap -- alist )
101 [ dup pheap-empty? not ] [ pheap-pop 2array ] [ ] produce nip ;
103 : pheap>values ( heap -- seq ) pheap>alist keys ;