1	# OMap: an ordered map keyed by any Data type, as a weight-balanced tree.
2	#
3	# The order is a template cmp, as in List.sort: OMap.put(~U32, ~U32,
4	# ~U32.cmp, m, 7, 70). put, get and del are O(log n) comparisons. Balance
5	# follows Hirai and Yamamoto (delta 3, gamma 2), which is correct for both
6	# insertion and deletion with one rebalance per level.
7	import Base
8	
9	type OMap<-K: Data, -V: Data> is Data:
10	  OTip{}
11	  OBin{size: U32, key: K, val: V, lo: OMap<K, V>, hi: OMap<K, V>}
12	
13	def OMap.new(-K: Data, -V: Data) -> OMap<K, V>:
14	  OTip{}
15	
16	def OMap.size(-K: Data, -V: Data, m: OMap<K, V>) -> U32:
17	  match m:
18	    case OTip{}:
19	      0
20	    case OBin{s, k, v, lo, hi}:
21	      s
22	
23	# The recursion on cmp's result: a match cannot scrutinize cmp(k, key), and
24	# a template cannot call a law filled below it, so each branch is a closure.
25	def Cmp.case(
26	  -A: Type, c: Cmp, lt: Unit -> A, eq: Unit -> A, gt: Unit -> A
27	) -> A:
28	  match c:
29	    case LT{}:
30	      lt(Unit{})
31	    case EQ{}:
32	      eq(Unit{})
33	    case GT{}:
34	      gt(Unit{})
35	
36	def OMap.bin(
37	  -K: Data, -V: Data, k: K, v: V, +lo: OMap<K, V>, +hi: OMap<K, V>
38	) -> OMap<K, V>:
39	  OBin{(1 + OMap.size(K, V, lo) + OMap.size(K, V, hi) : U32), k, v, lo, hi}
40	
41	def OMap.rot_lo.double(
42	  -K: Data, -V: Data, k: K, v: V, lo: OMap<K, V>, hk: K, hv: V,
43	  hl: OMap<K, V>, hh: OMap<K, V>
44	) -> OMap<K, V>:
45	  match hl:
46	    case OTip{}:
47	      OMap.bin(K, V, hk, hv, OMap.bin(K, V, k, v, lo, OTip{}), hh)
48	    case OBin{s, mk, mv, ml, mh}:
49	      OMap.bin(K, V, mk, mv, OMap.bin(K, V, k, v, lo, ml),
50	        OMap.bin(K, V, hk, hv, mh, hh))
51	
52	def OMap.rot_lo.pick(
53	  -K: Data, -V: Data, single: Bool, k: K, v: V, lo: OMap<K, V>, hk: K, hv: V,
54	  hl: OMap<K, V>, hh: OMap<K, V>
55	) -> OMap<K, V>:
56	  match single:
57	    case True{}:
58	      OMap.bin(K, V, hk, hv, OMap.bin(K, V, k, v, lo, hl), hh)
59	    case False{}:
60	      OMap.rot_lo.double(K, V, k, v, lo, hk, hv, hl, hh)
61	
62	# The hi side is too heavy: move weight toward lo.
63	def OMap.rot_lo(
64	  -K: Data, -V: Data, k: K, v: V, lo: OMap<K, V>, hi: OMap<K, V>
65	) -> OMap<K, V>:
66	  match hi:
67	    case OTip{}:
68	      OMap.bin(K, V, k, v, lo, OTip{})
69	    case OBin{s, hk, hv, +hl, +hh}:
70	      OMap.rot_lo.pick(K, V,
71	        U32.is_lt((OMap.size(K, V, hl) + 1 : U32),
72	          (2 * (OMap.size(K, V, hh) + 1) : U32)),
73	        k, v, lo, hk, hv, hl, hh)
74	
75	def OMap.rot_hi.double(
76	  -K: Data, -V: Data, k: K, v: V, lk: K, lv: V, ll: OMap<K, V>,
77	  lh: OMap<K, V>, hi: OMap<K, V>
78	) -> OMap<K, V>:
79	  match lh:
80	    case OTip{}:
81	      OMap.bin(K, V, lk, lv, ll, OMap.bin(K, V, k, v, OTip{}, hi))
82	    case OBin{s, mk, mv, ml, mh}:
83	      OMap.bin(K, V, mk, mv, OMap.bin(K, V, lk, lv, ll, ml),
84	        OMap.bin(K, V, k, v, mh, hi))
85	
86	def OMap.rot_hi.pick(
87	  -K: Data, -V: Data, single: Bool, k: K, v: V, lk: K, lv: V, ll: OMap<K, V>,
88	  lh: OMap<K, V>, hi: OMap<K, V>
89	) -> OMap<K, V>:
90	  match single:
91	    case True{}:
92	      OMap.bin(K, V, lk, lv, ll, OMap.bin(K, V, k, v, lh, hi))
93	    case False{}:
94	      OMap.rot_hi.double(K, V, k, v, lk, lv, ll, lh, hi)
95	
96	# The lo side is too heavy: move weight toward hi.
97	def OMap.rot_hi(
98	  -K: Data, -V: Data, k: K, v: V, lo: OMap<K, V>, hi: OMap<K, V>
99	) -> OMap<K, V>:
100	  match lo:
101	    case OTip{}:
102	      OMap.bin(K, V, k, v, OTip{}, hi)
103	    case OBin{s, lk, lv, +ll, +lh}:
104	      OMap.rot_hi.pick(K, V,
105	        U32.is_lt((OMap.size(K, V, lh) + 1 : U32),
106	          (2 * (OMap.size(K, V, ll) + 1) : U32)),
107	        k, v, lk, lv, ll, lh, hi)
108	
109	def OMap.balance.if(
110	  -K: Data, -V: Data, heavy_hi: Bool, heavy_lo: Bool, k: K, v: V,
111	  lo: OMap<K, V>, hi: OMap<K, V>
112	) -> OMap<K, V>:
113	  match heavy_hi heavy_lo:
114	    case True{} _:
115	      OMap.rot_lo(K, V, k, v, lo, hi)
116	    case False{} True{}:
117	      OMap.rot_hi(K, V, k, v, lo, hi)
118	    case False{} False{}:
119	      OMap.bin(K, V, k, v, lo, hi)
120	
121	# A node whose sides were balanced before one insertion or deletion.
122	def OMap.balance(
123	  -K: Data, -V: Data, k: K, v: V, +lo: OMap<K, V>, +hi: OMap<K, V>
124	) -> OMap<K, V>:
125	  +wl = (OMap.size(K, V, lo) + 1 : U32)
126	  +wh = (OMap.size(K, V, hi) + 1 : U32)
127	  OMap.balance.if(K, V, U32.is_gt(wh, (3 * wl : U32)),
128	    U32.is_gt(wl, (3 * wh : U32)), k, v, lo, hi)
129	
130	def OMap.put(
131	  ~K: Data, ~V: Data, ~cmp: K -> K -> Cmp, m: OMap<K, V>, +k: K, +v: V
132	) -> OMap<K, V>:
133	  match m:
134	    case OTip{}:
135	      OBin{1, k, v, OTip{}, OTip{}}
136	    case OBin{+s, +key, +val, +lo, +hi}:
137	      Cmp.case(OMap<K, V>, cmp(k, key),
138	        u => OMap.balance(K, V, key, val, OMap.put(~K, ~V, ~cmp, lo, k, v), hi),
139	        u => OBin{s, key, v, lo, hi},
140	        u => OMap.balance(K, V, key, val, lo, OMap.put(~K, ~V, ~cmp, hi, k, v)))
141	
142	def OMap.get(
143	  ~K: Data, ~V: Data, ~cmp: K -> K -> Cmp, m: OMap<K, V>, +k: K
144	) -> Maybe<&2, V>:
145	  match m:
146	    case OTip{}:
147	      None{}
148	    case OBin{s, +key, +val, +lo, +hi}:
149	      Cmp.case(Maybe<&2, V>, cmp(k, key),
150	        u => OMap.get(~K, ~V, ~cmp, lo, k),
151	        u => Some{val},
152	        u => OMap.get(~K, ~V, ~cmp, hi, k))
153	
154	def OMap.has(
155	  ~K: Data, ~V: Data, ~cmp: K -> K -> Cmp, m: OMap<K, V>, +k: K
156	) -> Bool:
157	  Maybe.is_some(&2, V, OMap.get(~K, ~V, ~cmp, m, k))
158	
159	def OMap.pop_min.fin(
160	  -K: Data, -V: Data, k: K, v: V, hi: OMap<K, V>,
161	  r: K & V & OMap<K, V>
162	) -> K & V & OMap<K, V>:
163	  (mk, mv, lo) = r
164	  (mk, mv, OMap.balance(K, V, k, v, lo, hi))
165	
166	# The least entry of the node lo < (k, v) < hi, and the node without it.
167	def OMap.pop_min(
168	  -K: Data, -V: Data, lo: OMap<K, V>, k: K, v: V, hi: OMap<K, V>
169	) -> K & V & OMap<K, V>:
170	  match lo:
171	    case OTip{}:
172	      (k, v, hi)
173	    case OBin{s, lk, lv, ll, lh}:
174	      OMap.pop_min.fin(K, V, k, v, hi, OMap.pop_min(K, V, ll, lk, lv, lh))
175	
176	def OMap.glue.fin(
177	  -K: Data, -V: Data, lo: OMap<K, V>, r: K & V & OMap<K, V>
178	) -> OMap<K, V>:
179	  (k, v, hi) = r
180	  OMap.balance(K, V, k, v, lo, hi)
181	
182	# Joins two sides of a deleted node.
183	def OMap.glue(-K: Data, -V: Data, lo: OMap<K, V>, hi: OMap<K, V>) -> OMap<K, V>:
184	  match hi:
185	    case OTip{}:
186	      lo
187	    case OBin{s, k, v, hl, hh}:
188	      OMap.glue.fin(K, V, lo, OMap.pop_min(K, V, hl, k, v, hh))
189	
190	def OMap.del(
191	  ~K: Data, ~V: Data, ~cmp: K -> K -> Cmp, m: OMap<K, V>, +k: K
192	) -> OMap<K, V>:
193	  match m:
194	    case OTip{}:
195	      OTip{}
196	    case OBin{s, +key, +val, +lo, +hi}:
197	      Cmp.case(OMap<K, V>, cmp(k, key),
198	        u => OMap.balance(K, V, key, val, OMap.del(~K, ~V, ~cmp, lo, k), hi),
199	        u => OMap.glue(K, V, lo, hi),
200	        u => OMap.balance(K, V, key, val, lo, OMap.del(~K, ~V, ~cmp, hi, k)))
201	
202	def OMap.keys.go(
203	  -K: Data, -V: Data, m: OMap<K, V>, acc: List<&2, K>
204	) -> List<&2, K>:
205	  match m:
206	    case OTip{}:
207	      acc
208	    case OBin{s, k, v, lo, hi}:
209	      OMap.keys.go(K, V, lo, k <> OMap.keys.go(K, V, hi, acc))
210	
211	# The keys in ascending order.
212	def OMap.keys(-K: Data, -V: Data, m: OMap<K, V>) -> List<&2, K>:
213	  OMap.keys.go(K, V, m, Nil{})
214	
215	def OMap.to_list.go(
216	  -K: Data, -V: Data, m: OMap<K, V>, acc: List<&1, K & V>
217	) -> List<&1, K & V>:
218	  match m:
219	    case OTip{}:
220	      acc
221	    case OBin{s, k, v, lo, hi}:
222	      OMap.to_list.go(K, V, lo, (k, v) <> OMap.to_list.go(K, V, hi, acc))
223	
224	# The entries in ascending key order.
225	def OMap.to_list(-K: Data, -V: Data, m: OMap<K, V>) -> List<&1, K & V>:
226	  OMap.to_list.go(K, V, m, Nil{})
227	
228	def OSet(-K: Data) -> Data:
229	  OMap<K, Unit>
230	
231	def OSet.new(-K: Data) -> OSet(K):
232	  OTip{}
233	
234	def OSet.size(-K: Data, s: OSet(K)) -> U32:
235	  OMap.size(K, Unit, s)
236	
237	def OSet.add(~K: Data, ~cmp: K -> K -> Cmp, s: OSet(K), +k: K) -> OSet(K):
238	  OMap.put(~K, ~Unit, ~cmp, s, k, Unit{})
239	
240	def OSet.has(~K: Data, ~cmp: K -> K -> Cmp, s: OSet(K), +k: K) -> Bool:
241	  OMap.has(~K, ~Unit, ~cmp, s, k)
242	
243	def OSet.del(~K: Data, ~cmp: K -> K -> Cmp, s: OSet(K), +k: K) -> OSet(K):
244	  OMap.del(~K, ~Unit, ~cmp, s, k)
245	
246	# The elements in ascending order.
247	def OSet.to_list(-K: Data, s: OSet(K)) -> List<&2, K>:
248	  OMap.keys(K, Unit, s)