hmap.bend

bend-kit-collections@0.1.4.1 · 211 lines · 18 definitions

1# HMap: a hash trie keyed by any Data type, with equality-checked collision buckets.
2# Equal keys must have equal hashes; different keys with one hash share a bucket.
3import Base
4
5type Entry<-K: Data, -V: Data> is Data:
6  Entry{key: K, val: V}
7
8type HMap<-K: Data, -V: Data> is Data:
9  HNil{}
10  HLeaf{hash: U32, entries: List<&2, Entry<K, V>>}
11  HBranch{lo: HMap<K, V>, hi: HMap<K, V>}
12
13def HMap.new(-K: Data, -V: Data) -> HMap<K, V>:
14  HNil{}
15
16def HMap.bool(-A: Type, b: Bool, yes: Unit -> A, no: Unit -> A) -> A:
17  match b:
18    case True{}:
19      yes(Unit{})
20    case False{}:
21      no(Unit{})
22
23def HMap.entries.put(
24  ~K: Data, ~V: Data, ~eq: K -> K -> Bool,
25  xs: List<&2, Entry<K, V>>, +k: K, +v: V
26) -> List<&2, Entry<K, V>>:
27  match xs:
28    case Nil{}:
29      [Entry{k, v}]
30    case Con{Entry{+key, +val}, +rest}:
31      HMap.bool(List<&2, Entry<K, V>>, eq(k, key),
32        u => Entry{key, v} <> rest,
33        u => Entry{key, val} <> HMap.entries.put(~K, ~V, ~eq, rest, k, v))
34
35def HMap.entries.get(
36  ~K: Data, ~V: Data, ~eq: K -> K -> Bool,
37  xs: List<&2, Entry<K, V>>, +k: K
38) -> Maybe<&2, V>:
39  match xs:
40    case Nil{}:
41      None{}
42    case Con{Entry{+key, +val}, +rest}:
43      HMap.bool(Maybe<&2, V>, eq(k, key),
44        u => Some{val},
45        u => HMap.entries.get(~K, ~V, ~eq, rest, k))
46
47def HMap.entries.del(
48  ~K: Data, ~V: Data, ~eq: K -> K -> Bool,
49  xs: List<&2, Entry<K, V>>, +k: K
50) -> List<&2, Entry<K, V>>:
51  match xs:
52    case Nil{}:
53      Nil{}
54    case Con{Entry{+key, +val}, +rest}:
55      HMap.bool(List<&2, Entry<K, V>>, eq(k, key),
56        u => rest,
57        u => Entry{key, val} <> HMap.entries.del(~K, ~V, ~eq, rest, k))
58
59def HMap.split.side(
60  -K: Data, -V: Data, left: Bool, old: HMap<K, V>, fresh: HMap<K, V>
61) -> HMap<K, V>:
62  match left:
63    case True{}:
64      HBranch{old, fresh}
65    case False{}:
66      HBranch{fresh, old}
67
68def HMap.split(
69  ~K: Data, ~V: Data, ~eq: K -> K -> Bool,
70  +n: Nat, +mask: U32, +oldhash: U32, +entries: List<&2, Entry<K, V>>,
71  +hash: U32, +k: K, +v: V
72) -> HMap<K, V>:
73  match n:
74    case 0n:
75      HLeaf{oldhash, HMap.entries.put(~K, ~V, ~eq, entries, k, v)}
76    case 1n+p:
77      HMap.bool(HMap<K, V>, U32.is_eq((oldhash .&. mask : U32), (hash .&. mask : U32)),
78        u => HMap.bool(HMap<K, V>, U32.is_eq((oldhash .&. mask : U32), 0),
79          x => HBranch{HMap.split(~K, ~V, ~eq, p, (mask << 1n : U32), oldhash, entries, hash, k, v), HNil{}},
80          x => HBranch{HNil{}, HMap.split(~K, ~V, ~eq, p, (mask << 1n : U32), oldhash, entries, hash, k, v)}),
81        u => HMap.split.side(K, V, U32.is_eq((oldhash .&. mask : U32), 0),
82          HLeaf{oldhash, entries}, HLeaf{hash, [Entry{k, v}]}))
83
84def HMap.put.go(
85  ~K: Data, ~V: Data, ~eq: K -> K -> Bool,
86  +n: Nat, +mask: U32, m: HMap<K, V>, +hash: U32, +k: K, +v: V
87) -> HMap<K, V>:
88  match n:
89    case 0n:
90      match m:
91        case HNil{}:
92          HLeaf{hash, [Entry{k, v}]}
93        case HLeaf{oldhash, entries}:
94          HLeaf{oldhash, HMap.entries.put(~K, ~V, ~eq, entries, k, v)}
95        case HBranch{lo, hi}:
96          HBranch{lo, hi}
97    case 1n+p:
98      match m:
99        case HNil{}:
100          HLeaf{hash, [Entry{k, v}]}
101        case HLeaf{+oldhash, +entries}:
102          HMap.bool(HMap<K, V>, U32.is_eq(oldhash, hash),
103            u => HLeaf{oldhash, HMap.entries.put(~K, ~V, ~eq, entries, k, v)},
104            u => HMap.split(~K, ~V, ~eq, n, mask, oldhash, entries, hash, k, v))
105        case HBranch{+lo, +hi}:
106          HMap.bool(HMap<K, V>, U32.is_eq((hash .&. mask : U32), 0),
107            u => HBranch{HMap.put.go(~K, ~V, ~eq, p, (mask << 1n : U32), lo, hash, k, v), hi},
108            u => HBranch{lo, HMap.put.go(~K, ~V, ~eq, p, (mask << 1n : U32), hi, hash, k, v)})
109
110def HMap.put(
111  ~K: Data, ~V: Data, ~hash: K -> U32, ~eq: K -> K -> Bool,
112  m: HMap<K, V>, +k: K, v: V
113) -> HMap<K, V>:
114  HMap.put.go(~K, ~V, ~eq, 32n, 1, m, hash(k), k, v)
115
116def HMap.get.go(
117  ~K: Data, ~V: Data, ~eq: K -> K -> Bool,
118  +n: Nat, +mask: U32, m: HMap<K, V>, +hash: U32, +k: K
119) -> Maybe<&2, V>:
120  match n:
121    case 0n:
122      match m:
123        case HNil{}:
124          None{}
125        case HLeaf{oldhash, entries}:
126          HMap.entries.get(~K, ~V, ~eq, entries, k)
127        case HBranch{lo, hi}:
128          None{}
129    case 1n+p:
130      match m:
131        case HNil{}:
132          None{}
133        case HLeaf{oldhash, entries}:
134          HMap.bool(Maybe<&2, V>, U32.is_eq(oldhash, hash),
135            u => HMap.entries.get(~K, ~V, ~eq, entries, k),
136            u => None{})
137        case HBranch{+lo, +hi}:
138          HMap.bool(Maybe<&2, V>, U32.is_eq((hash .&. mask : U32), 0),
139            u => HMap.get.go(~K, ~V, ~eq, p, (mask << 1n : U32), lo, hash, k),
140            u => HMap.get.go(~K, ~V, ~eq, p, (mask << 1n : U32), hi, hash, k))
141
142def HMap.get(
143  ~K: Data, ~V: Data, ~hash: K -> U32, ~eq: K -> K -> Bool,
144  m: HMap<K, V>, +k: K
145) -> Maybe<&2, V>:
146  HMap.get.go(~K, ~V, ~eq, 32n, 1, m, hash(k), k)
147
148# An emptied bucket becomes HNil.
149def HMap.leaf(-K: Data, -V: Data, +hash: U32, xs: List<&2, Entry<K, V>>) -> HMap<K, V>:
150  match xs:
151    case Nil{}:
152      HNil{}
153    case Con{e, rest}:
154      HLeaf{hash, e <> rest}
155
156# Rebuilds a branch; one with two empty subtries becomes HNil.
157def HMap.join(-K: Data, -V: Data, lo: HMap<K, V>, hi: HMap<K, V>) -> HMap<K, V>:
158  match lo:
159    case HNil{}:
160      match hi:
161        case HNil{}:
162          HNil{}
163        case HLeaf{h, es}:
164          HBranch{HNil{}, HLeaf{h, es}}
165        case HBranch{a, b}:
166          HBranch{HNil{}, HBranch{a, b}}
167    case HLeaf{h, es}:
168      HBranch{HLeaf{h, es}, hi}
169    case HBranch{a, b}:
170      HBranch{HBranch{a, b}, hi}
171
172def HMap.del.go(
173  ~K: Data, ~V: Data, ~eq: K -> K -> Bool,
174  +n: Nat, +mask: U32, m: HMap<K, V>, +hash: U32, +k: K
175) -> HMap<K, V>:
176  match n:
177    case 0n:
178      match m:
179        case HNil{}:
180          HNil{}
181        case HLeaf{oldhash, entries}:
182          HMap.leaf(K, V, oldhash, HMap.entries.del(~K, ~V, ~eq, entries, k))
183        case HBranch{lo, hi}:
184          HBranch{lo, hi}
185    case 1n+p:
186      match m:
187        case HNil{}:
188          HNil{}
189        case HLeaf{+oldhash, +entries}:
190          HMap.bool(HMap<K, V>, U32.is_eq(oldhash, hash),
191            u => HMap.leaf(K, V, oldhash, HMap.entries.del(~K, ~V, ~eq, entries, k)),
192            u => HLeaf{oldhash, entries})
193        case HBranch{+lo, +hi}:
194          HMap.bool(HMap<K, V>, U32.is_eq((hash .&. mask : U32), 0),
195            u => HMap.join(K, V, HMap.del.go(~K, ~V, ~eq, p, (mask << 1n : U32), lo, hash, k), hi),
196            u => HMap.join(K, V, lo, HMap.del.go(~K, ~V, ~eq, p, (mask << 1n : U32), hi, hash, k)))
197
198def HMap.del(
199  ~K: Data, ~V: Data, ~hash: K -> U32, ~eq: K -> K -> Bool,
200  m: HMap<K, V>, +k: K
201) -> HMap<K, V>:
202  HMap.del.go(~K, ~V, ~eq, 32n, 1, m, hash(k), k)
203
204def HMap.size(-K: Data, -V: Data, m: HMap<K, V>) -> Nat:
205  match m:
206    case HNil{}:
207      0n
208    case HLeaf{hash, entries}:
209      List.length(&2, Entry<K, V>, entries)
210    case HBranch{lo, hi}:
211      Nat.add(HMap.size(K, V, lo), HMap.size(K, V, hi))