1	# Heap: a priority queue, as a skew heap ordered by a le you pass in.
2	#
3	# The smallest element (by le) sits at the top. push and pop are
4	# amortized O(log n) over a heap used once (sharing one with + and popping
5	# both copies can cost more), peek and size are O(1), and sort is
6	# heapsort. le is a template, as in List.sort: Heap.push(~U32, ~U32.is_le,
7	# 7, h).
8	#
9	# The laws at the bottom are checked every time this file is imported.
10	import Base
11	
12	type Heap<-A: Data> is Data:
13	  HLeaf{}
14	  HNode{size: Nat, top: A, left: Heap<A>, right: Heap<A>}
15	
16	def Heap.new(-A: Data) -> Heap<A>:
17	  HLeaf{}
18	
19	def Heap.size(-A: Data, h: Heap<A>) -> Nat:
20	  match h:
21	    case HLeaf{}:
22	      0n
23	    case HNode{n, x, l, r}:
24	      n
25	
26	def Heap.peek(-A: Data, h: Heap<A>) -> Maybe<&2, A>:
27	  match h:
28	    case HLeaf{}:
29	      None{}
30	    case HNode{n, x, l, r}:
31	      Some{x}
32	
33	# The smaller root wins: its right child merges with the other heap and
34	# becomes the left, and its old left becomes the right (a skew heap's
35	# swap, which keeps merges short on average). The fuel is the size of both
36	# heaps: every step takes a root out, so it never runs short.
37	def Heap.merge.go(
38	  ~A: Data, ~le: A -> A -> Bool, fuel: Nat, a: Heap<A>, b: Heap<A>
39	) -> Heap<A>:
40	  match fuel:
41	    case 0n:
42	      a
43	    case 1n+f:
44	      match a b:
45	        case HLeaf{} HLeaf{}:
46	          HLeaf{}
47	        case HLeaf{} HNode{nb, y, bl, br}:
48	          HNode{nb, y, bl, br}
49	        case HNode{na, x, al, ar} HLeaf{}:
50	          HNode{na, x, al, ar}
51	        case HNode{+na, +x, +al, +ar} HNode{+nb, +y, +bl, +br}:
52	          +c = le(x, y)
53	          HNode{Nat.add(na, nb), Bool.pick(A, c, x, y),
54	            Heap.merge.go(~A, ~le, f,
55	              Bool.pick(Heap<A>, c, ar, br),
56	              Bool.pick(Heap<A>, c, HNode{nb, y, bl, br},
57	                HNode{na, x, al, ar})),
58	            Bool.pick(Heap<A>, c, al, bl)}
59	
60	def Heap.merge(
61	  ~A: Data, ~le: A -> A -> Bool, +a: Heap<A>, +b: Heap<A>
62	) -> Heap<A>:
63	  Heap.merge.go(~A, ~le, Nat.add(Heap.size(A, a), Heap.size(A, b)), a, b)
64	
65	def Heap.push(~A: Data, ~le: A -> A -> Bool, x: A, h: Heap<A>) -> Heap<A>:
66	  Heap.merge(~A, ~le, HNode{1n, x, HLeaf{}, HLeaf{}}, h)
67	
68	# the top, and the heap without it; None on an empty heap
69	def Heap.pop(
70	  ~A: Data, ~le: A -> A -> Bool, h: Heap<A>
71	) -> Maybe<&1, A & Heap<A>>:
72	  match h:
73	    case HLeaf{}:
74	      None{}
75	    case HNode{n, x, l, r}:
76	      Some{(x, Heap.merge(~A, ~le, l, r))}
77	
78	def Heap.from_list(~A: Data, ~le: A -> A -> Bool, xs: List<&2, A>) -> Heap<A>:
79	  match xs:
80	    case Nil{}:
81	      HLeaf{}
82	    case h <> t:
83	      Heap.push(~A, ~le, h, Heap.from_list(~A, ~le, t))
84	
85	def Heap.to_list.go(
86	  ~A: Data, ~le: A -> A -> Bool, fuel: Nat, h: Heap<A>
87	) -> List<&2, A>:
88	  match fuel:
89	    case 0n:
90	      Nil{}
91	    case 1n+f:
92	      match h:
93	        case HLeaf{}:
94	          Nil{}
95	        case HNode{n, x, l, r}:
96	          x <> Heap.to_list.go(~A, ~le, f, Heap.merge(~A, ~le, l, r))
97	
98	# the elements, smallest first
99	def Heap.to_list(~A: Data, ~le: A -> A -> Bool, +h: Heap<A>) -> List<&2, A>:
100	  Heap.to_list.go(~A, ~le, Heap.size(A, h), h)
101	
102	# heapsort
103	def Heap.sort(~A: Data, ~le: A -> A -> Bool, xs: List<&2, A>) -> List<&2, A>:
104	  Heap.to_list(~A, ~le, Heap.from_list(~A, ~le, xs))
105	
106	# Laws
107	# ----
108	
109	# LAW: popping what was pushed onto an empty heap gives it back
110	law Heap.pop.push_new:
111	  for -x: U32
112	  {Heap.pop(~U32, ~U32.is_le, Heap.push(~U32, ~U32.is_le, x, Heap.new(U32)))
113	   == Some{(x, Heap.new(U32))} : Maybe<&1, U32 & Heap<U32>>}
114	
115	def Heap.pop.push_new(x):
116	  {==}
117	
118	# LAW: a push adds one to the size
119	law Heap.size.push:
120	  for -x: U32
121	  for  h: Heap<U32>
122	  {Heap.size(U32, Heap.push(~U32, ~U32.is_le, x, h))
123	   == 1n+Heap.size(U32, h) : Nat}
124	
125	def Heap.size.push(x, h):
126	  match h:
127	    case HLeaf{}:
128	      {==}
129	    case HNode{n, y, l, r}:
130	      {==}