omap.bend

bend-kit-collections@0.1.4.1 · 248 lines · 32 definitions

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.
7import Base
8
9type 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
13def OMap.new(-K: Data, -V: Data) -> OMap<K, V>:
14  OTip{}
15
16def 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.
25def 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
36def 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
41def 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
52def 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.
63def 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
75def 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
86def 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.
97def 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
109def 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.
122def 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
130def 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
142def 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
154def 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
159def 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.
167def 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
176def 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.
183def 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
190def 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
202def 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.
212def OMap.keys(-K: Data, -V: Data, m: OMap<K, V>) -> List<&2, K>:
213  OMap.keys.go(K, V, m, Nil{})
214
215def 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.
225def 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
228def OSet(-K: Data) -> Data:
229  OMap<K, Unit>
230
231def OSet.new(-K: Data) -> OSet(K):
232  OTip{}
233
234def OSet.size(-K: Data, s: OSet(K)) -> U32:
235  OMap.size(K, Unit, s)
236
237def 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
240def OSet.has(~K: Data, ~cmp: K -> K -> Cmp, s: OSet(K), +k: K) -> Bool:
241  OMap.has(~K, ~Unit, ~cmp, s, k)
242
243def 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.
247def OSet.to_list(-K: Data, s: OSet(K)) -> List<&2, K>:
248  OMap.keys(K, Unit, s)