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.
3	import Base
4	
5	type Entry<-K: Data, -V: Data> is Data:
6	  Entry{key: K, val: V}
7	
8	type 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	
13	def HMap.new(-K: Data, -V: Data) -> HMap<K, V>:
14	  HNil{}
15	
16	def 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	
23	def 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	
35	def 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	
47	def 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	
59	def 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	
68	def 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	
84	def 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	
110	def 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	
116	def 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	
142	def 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.
149	def 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.
157	def 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	
172	def 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	
198	def 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	
204	def 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))