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.split.side(
48	  -K: Data, -V: Data, left: Bool, old: HMap<K, V>, fresh: HMap<K, V>
49	) -> HMap<K, V>:
50	  match left:
51	    case True{}:
52	      HBranch{old, fresh}
53	    case False{}:
54	      HBranch{fresh, old}
55	
56	def HMap.split(
57	  ~K: Data, ~V: Data, ~eq: K -> K -> Bool,
58	  +n: Nat, +mask: U32, +oldhash: U32, +entries: List<&2, Entry<K, V>>,
59	  +hash: U32, +k: K, +v: V
60	) -> HMap<K, V>:
61	  match n:
62	    case 0n:
63	      HLeaf{oldhash, HMap.entries.put(~K, ~V, ~eq, entries, k, v)}
64	    case 1n+p:
65	      HMap.bool(HMap<K, V>, U32.is_eq((oldhash .&. mask : U32), (hash .&. mask : U32)),
66	        u => HMap.bool(HMap<K, V>, U32.is_eq((oldhash .&. mask : U32), 0),
67	          x => HBranch{HMap.split(~K, ~V, ~eq, p, (mask << 1n : U32), oldhash, entries, hash, k, v), HNil{}},
68	          x => HBranch{HNil{}, HMap.split(~K, ~V, ~eq, p, (mask << 1n : U32), oldhash, entries, hash, k, v)}),
69	        u => HMap.split.side(K, V, U32.is_eq((oldhash .&. mask : U32), 0),
70	          HLeaf{oldhash, entries}, HLeaf{hash, [Entry{k, v}]}))
71	
72	def HMap.put.go(
73	  ~K: Data, ~V: Data, ~eq: K -> K -> Bool,
74	  +n: Nat, +mask: U32, m: HMap<K, V>, +hash: U32, +k: K, +v: V
75	) -> HMap<K, V>:
76	  match n:
77	    case 0n:
78	      match m:
79	        case HNil{}:
80	          HLeaf{hash, [Entry{k, v}]}
81	        case HLeaf{oldhash, entries}:
82	          HLeaf{oldhash, HMap.entries.put(~K, ~V, ~eq, entries, k, v)}
83	        case HBranch{lo, hi}:
84	          HBranch{lo, hi}
85	    case 1n+p:
86	      match m:
87	        case HNil{}:
88	          HLeaf{hash, [Entry{k, v}]}
89	        case HLeaf{+oldhash, +entries}:
90	          HMap.bool(HMap<K, V>, U32.is_eq(oldhash, hash),
91	            u => HLeaf{oldhash, HMap.entries.put(~K, ~V, ~eq, entries, k, v)},
92	            u => HMap.split(~K, ~V, ~eq, n, mask, oldhash, entries, hash, k, v))
93	        case HBranch{+lo, +hi}:
94	          HMap.bool(HMap<K, V>, U32.is_eq((hash .&. mask : U32), 0),
95	            u => HBranch{HMap.put.go(~K, ~V, ~eq, p, (mask << 1n : U32), lo, hash, k, v), hi},
96	            u => HBranch{lo, HMap.put.go(~K, ~V, ~eq, p, (mask << 1n : U32), hi, hash, k, v)})
97	
98	def HMap.put(
99	  ~K: Data, ~V: Data, ~hash: K -> U32, ~eq: K -> K -> Bool,
100	  m: HMap<K, V>, +k: K, v: V
101	) -> HMap<K, V>:
102	  HMap.put.go(~K, ~V, ~eq, 32n, 1, m, hash(k), k, v)
103	
104	def HMap.get.go(
105	  ~K: Data, ~V: Data, ~eq: K -> K -> Bool,
106	  +n: Nat, +mask: U32, m: HMap<K, V>, +hash: U32, +k: K
107	) -> Maybe<&2, V>:
108	  match n:
109	    case 0n:
110	      match m:
111	        case HNil{}:
112	          None{}
113	        case HLeaf{oldhash, entries}:
114	          HMap.entries.get(~K, ~V, ~eq, entries, k)
115	        case HBranch{lo, hi}:
116	          None{}
117	    case 1n+p:
118	      match m:
119	        case HNil{}:
120	          None{}
121	        case HLeaf{oldhash, entries}:
122	          HMap.bool(Maybe<&2, V>, U32.is_eq(oldhash, hash),
123	            u => HMap.entries.get(~K, ~V, ~eq, entries, k),
124	            u => None{})
125	        case HBranch{+lo, +hi}:
126	          HMap.bool(Maybe<&2, V>, U32.is_eq((hash .&. mask : U32), 0),
127	            u => HMap.get.go(~K, ~V, ~eq, p, (mask << 1n : U32), lo, hash, k),
128	            u => HMap.get.go(~K, ~V, ~eq, p, (mask << 1n : U32), hi, hash, k))
129	
130	def HMap.get(
131	  ~K: Data, ~V: Data, ~hash: K -> U32, ~eq: K -> K -> Bool,
132	  m: HMap<K, V>, +k: K
133	) -> Maybe<&2, V>:
134	  HMap.get.go(~K, ~V, ~eq, 32n, 1, m, hash(k), k)
135	
136	def HMap.size(-K: Data, -V: Data, m: HMap<K, V>) -> Nat:
137	  match m:
138	    case HNil{}:
139	      0n
140	    case HLeaf{hash, entries}:
141	      List.length(&2, Entry<K, V>, entries)
142	    case HBranch{lo, hi}:
143	      Nat.add(HMap.size(K, V, lo), HMap.size(K, V, hi))