0 | module Libraries.Data.StringMap
  1 |
  2 | -- Hand specialised map, for efficiency...
  3 |
  4 | %default total
  5 |
  6 | Key : Type
  7 | Key = String
  8 |
  9 | -- TODO: write split
 10 |
 11 | private
 12 | data Tree : Nat -> Type -> Type where
 13 |   Leaf : Key -> v -> Tree Z v
 14 |   Branch2 : Tree n v -> Key -> Tree n v -> Tree (S n) v
 15 |   Branch3 : Tree n v -> Key -> Tree n v -> Key -> Tree n v -> Tree (S n) v
 16 |
 17 | branch4 :
 18 |   Tree n v -> Key ->
 19 |   Tree n v -> Key ->
 20 |   Tree n v -> Key ->
 21 |   Tree n v ->
 22 |   Tree (S (S n)) v
 23 | branch4 a b c d e f g =
 24 |   Branch2 (Branch2 a b c) d (Branch2 e f g)
 25 |
 26 | branch5 :
 27 |   Tree n v -> Key ->
 28 |   Tree n v -> Key ->
 29 |   Tree n v -> Key ->
 30 |   Tree n v -> Key ->
 31 |   Tree n v ->
 32 |   Tree (S (S n)) v
 33 | branch5 a b c d e f g h i =
 34 |   Branch2 (Branch2 a b c) d (Branch3 e f g h i)
 35 |
 36 | branch6 :
 37 |   Tree n v -> Key ->
 38 |   Tree n v -> Key ->
 39 |   Tree n v -> Key ->
 40 |   Tree n v -> Key ->
 41 |   Tree n v -> Key ->
 42 |   Tree n v ->
 43 |   Tree (S (S n)) v
 44 | branch6 a b c d e f g h i j k =
 45 |   Branch3 (Branch2 a b c) d (Branch2 e f g) h (Branch2 i j k)
 46 |
 47 | branch7 :
 48 |   Tree n v -> Key ->
 49 |   Tree n v -> Key ->
 50 |   Tree n v -> Key ->
 51 |   Tree n v -> Key ->
 52 |   Tree n v -> Key ->
 53 |   Tree n v -> Key ->
 54 |   Tree n v ->
 55 |   Tree (S (S n)) v
 56 | branch7 a b c d e f g h i j k l m =
 57 |   Branch3 (Branch3 a b c d e) f (Branch2 g h i) j (Branch2 k l m)
 58 |
 59 | merge1 : Tree n v -> Key -> Tree (S n) v -> Key -> Tree (S n) v -> Tree (S (S n)) v
 60 | merge1 a b (Branch2 c d e) f (Branch2 g h i) = branch5 a b c d e f g h i
 61 | merge1 a b (Branch2 c d e) f (Branch3 g h i j k) = branch6 a b c d e f g h i j k
 62 | merge1 a b (Branch3 c d e f g) h (Branch2 i j k) = branch6 a b c d e f g h i j k
 63 | merge1 a b (Branch3 c d e f g) h (Branch3 i j k l m) = branch7 a b c d e f g h i j k l m
 64 |
 65 | merge2 : Tree (S n) v -> Key -> Tree n v -> Key -> Tree (S n) v -> Tree (S (S n)) v
 66 | merge2 (Branch2 a b c) d e f (Branch2 g h i) = branch5 a b c d e f g h i
 67 | merge2 (Branch2 a b c) d e f (Branch3 g h i j k) = branch6 a b c d e f g h i j k
 68 | merge2 (Branch3 a b c d e) f g h (Branch2 i j k) = branch6 a b c d e f g h i j k
 69 | merge2 (Branch3 a b c d e) f g h (Branch3 i j k l m) = branch7 a b c d e f g h i j k l m
 70 |
 71 | merge3 : Tree (S n) v -> Key -> Tree (S n) v -> Key -> Tree n v -> Tree (S (S n)) v
 72 | merge3 (Branch2 a b c) d (Branch2 e f g) h i = branch5 a b c d e f g h i
 73 | merge3 (Branch2 a b c) d (Branch3 e f g h i) j k = branch6 a b c d e f g h i j k
 74 | merge3 (Branch3 a b c d e) f (Branch2 g h i) j k = branch6 a b c d e f g h i j k
 75 | merge3 (Branch3 a b c d e) f (Branch3 g h i j k) l m = branch7 a b c d e f g h i j k l m
 76 |
 77 | treeLookup : Key -> Tree n v -> Maybe v
 78 | treeLookup k (Leaf k' v) =
 79 |   if k == k' then
 80 |     Just v
 81 |   else
 82 |     Nothing
 83 | treeLookup k (Branch2 t1 k' t2) =
 84 |   if k <= k' then
 85 |     treeLookup k t1
 86 |   else
 87 |     treeLookup k t2
 88 | treeLookup k (Branch3 t1 k1 t2 k2 t3) =
 89 |   if k <= k1 then
 90 |     treeLookup k t1
 91 |   else if k <= k2 then
 92 |     treeLookup k t2
 93 |   else
 94 |     treeLookup k t3
 95 |
 96 | treeInsert' : Key -> v -> Tree n v -> Either (Tree n v) (Tree n v, Key, Tree n v)
 97 | treeInsert' k v (Leaf k' v') =
 98 |   case compare k k' of
 99 |     LT => Right (Leaf k v, k, Leaf k' v')
100 |     EQ => Left (Leaf k v)
101 |     GT => Right (Leaf k' v', k', Leaf k v)
102 | treeInsert' k v (Branch2 t1 k' t2) =
103 |   if k <= k' then
104 |     case treeInsert' k v t1 of
105 |       Left t1' => Left (Branch2 t1' k' t2)
106 |       Right (a, b, c) => Left (Branch3 a b c k' t2)
107 |   else
108 |     case treeInsert' k v t2 of
109 |       Left t2' => Left (Branch2 t1 k' t2')
110 |       Right (a, b, c) => Left (Branch3 t1 k' a b c)
111 | treeInsert' k v (Branch3 t1 k1 t2 k2 t3) =
112 |   if k <= k1 then
113 |     case treeInsert' k v t1 of
114 |       Left t1' => Left (Branch3 t1' k1 t2 k2 t3)
115 |       Right (a, b, c) => Right (Branch2 a b c, k1, Branch2 t2 k2 t3)
116 |   else
117 |     if k <= k2 then
118 |       case treeInsert' k v t2 of
119 |         Left t2' => Left (Branch3 t1 k1 t2' k2 t3)
120 |         Right (a, b, c) => Right (Branch2 t1 k1 a, b, Branch2 c k2 t3)
121 |     else
122 |       case treeInsert' k v t3 of
123 |         Left t3' => Left (Branch3 t1 k1 t2 k2 t3')
124 |         Right (a, b, c) => Right (Branch2 t1 k1 t2, k2, Branch2 a b c)
125 |
126 | treeInsert : Key -> v -> Tree n v -> Either (Tree n v) (Tree (S n) v)
127 | treeInsert k v t =
128 |   case treeInsert' k v t of
129 |     Left t' => Left t'
130 |     Right (a, b, c) => Right (Branch2 a b c)
131 |
132 | delType : Nat -> Type -> Type
133 | delType Z v = ()
134 | delType (S n) v = Tree n v
135 |
136 | treeDelete : {n : _} -> Key -> Tree n v -> Either (Tree n v) (delType n v)
137 | treeDelete k (Leaf k' v) =
138 |   if k == k' then
139 |     Right ()
140 |   else
141 |     Left (Leaf k' v)
142 | treeDelete {n=S Z} k (Branch2 t1 k' t2) =
143 |   if k <= k' then
144 |     case treeDelete k t1 of
145 |       Left t1' => Left (Branch2 t1' k' t2)
146 |       Right () => Right t2
147 |   else
148 |     case treeDelete k t2 of
149 |       Left t2' => Left (Branch2 t1 k' t2')
150 |       Right () => Right t1
151 | treeDelete {n=S Z} k (Branch3 t1 k1 t2 k2 t3) =
152 |   if k <= k1 then
153 |     case treeDelete k t1 of
154 |       Left t1' => Left (Branch3 t1' k1 t2 k2 t3)
155 |       Right () => Left (Branch2 t2 k2 t3)
156 |   else if k <= k2 then
157 |     case treeDelete k t2 of
158 |       Left t2' => Left (Branch3 t1 k1 t2' k2 t3)
159 |       Right () => Left (Branch2 t1 k1 t3)
160 |   else
161 |     case treeDelete k t3 of
162 |       Left t3' => Left (Branch3 t1 k1 t2 k2 t3')
163 |       Right () => Left (Branch2 t1 k1 t2)
164 | treeDelete {n=S (S _)} k (Branch2 t1 k' t2) =
165 |   if k <= k' then
166 |     case treeDelete k t1 of
167 |       Left t1' => Left (Branch2 t1' k' t2)
168 |       Right t1' =>
169 |         case t2 of
170 |           Branch2 a b c => Right (Branch3 t1' k' a b c)
171 |           Branch3 a b c d e => Left (branch4 t1' k' a b c d e)
172 |   else
173 |     case treeDelete k t2 of
174 |       Left t2' => Left (Branch2 t1 k' t2')
175 |       Right t2' =>
176 |         case t1 of
177 |           Branch2 a b c => Right (Branch3 a b c k' t2')
178 |           Branch3 a b c d e => Left (branch4 a b c d e k' t2')
179 | treeDelete {n=(S (S _))} k (Branch3 t1 k1 t2 k2 t3) =
180 |   if k <= k1 then
181 |     case treeDelete k t1 of
182 |       Left t1' => Left (Branch3 t1' k1 t2 k2 t3)
183 |       Right t1' => Left (merge1 t1' k1 t2 k2 t3)
184 |   else if k <= k2 then
185 |     case treeDelete k t2 of
186 |       Left t2' => Left (Branch3 t1 k1 t2' k2 t3)
187 |       Right t2' => Left (merge2 t1 k1 t2' k2 t3)
188 |   else
189 |     case treeDelete k t3 of
190 |       Left t3' => Left (Branch3 t1 k1 t2 k2 t3')
191 |       Right t3' => Left (merge3 t1 k1 t2 k2 t3')
192 |
193 | treeToList : Tree n v -> List (Key, v)
194 | treeToList = treeToList' (:: [])
195 |   where
196 |     treeToList' : forall n . ((Key, v) -> List (Key, v)) -> Tree n v -> List (Key, v)
197 |     treeToList' cont (Leaf k v) = cont (k, v)
198 |     treeToList' cont (Branch2 t1 _ t2) = treeToList' (:: treeToList' cont t2) t1
199 |     treeToList' cont (Branch3 t1 _ t2 _ t3) = treeToList' (:: treeToList' (:: treeToList' cont t3) t2) t1
200 |
201 | export
202 | data StringMap : Type -> Type where
203 |   Empty : StringMap v
204 |   M : (n : Nat) -> Tree n v -> StringMap v
205 |
206 | export
207 | empty : StringMap v
208 | empty = Empty
209 |
210 | export
211 | singleton : String -> v -> StringMap v
212 | singleton k v = M Z (Leaf k v)
213 |
214 | export
215 | lookup : String -> StringMap v -> Maybe v
216 | lookup _ Empty = Nothing
217 | lookup k (M _ t) = treeLookup k t
218 |
219 | export
220 | insert : String -> v -> StringMap v -> StringMap v
221 | insert k v Empty = M Z (Leaf k v)
222 | insert k v (M _ t) =
223 |   case treeInsert k v t of
224 |     Left t' => (M _ t')
225 |     Right t' => (M _ t')
226 |
227 | export
228 | insertFrom : List (String, v) -> StringMap v -> StringMap v
229 | insertFrom = flip $ foldl $ flip $ uncurry insert
230 |
231 | export
232 | delete : String -> StringMap v -> StringMap v
233 | delete _ Empty = Empty
234 | delete k (M Z t) =
235 |   case treeDelete k t of
236 |     Left t' => (M _ t')
237 |     Right () => Empty
238 | delete k (M (S _) t) =
239 |   case treeDelete k t of
240 |     Left t' => (M _ t')
241 |     Right t' => (M _ t')
242 |
243 | export
244 | fromList : List (String, v) -> StringMap v
245 | fromList l = foldl (flip (uncurry insert)) empty l
246 |
247 | export
248 | toList : StringMap v -> List (String, v)
249 | toList Empty = []
250 | toList (M _ t) = treeToList t
251 |
252 | ||| Gets the Keys of the map.
253 | export
254 | keys : StringMap v -> List String
255 | keys = map fst . StringMap.toList
256 |
257 | ||| Gets the values of the map. Could contain duplicates.
258 | export
259 | values : StringMap v -> List v
260 | values = map snd . StringMap.toList
261 |
262 | treeMap : (a -> b) -> Tree n a -> Tree n b
263 | treeMap f (Leaf k v) = Leaf k (f v)
264 | treeMap f (Branch2 t1 k t2) = Branch2 (treeMap f t1) k (treeMap f t2)
265 | treeMap f (Branch3 t1 k1 t2 k2 t3)
266 |     = Branch3 (treeMap f t1) k1 (treeMap f t2) k2 (treeMap f t3)
267 |
268 | export
269 | implementation Functor StringMap where
270 |   map _ Empty = Empty
271 |   map f (M n t) = M _ (treeMap f t)
272 |
273 | ||| Return a map of pairs of values that happen to share the same keys
274 | export
275 | sharedSupport : StringMap v -> StringMap v -> StringMap (v, v)
276 | sharedSupport l
277 |   = foldl (\acc, (k, v) => case lookup k l of
278 |              Just v' => insert k (v, v') acc
279 |              Nothing => acc)
280 |       empty . StringMap.toList
281 |
282 | ||| Merge two maps. When encountering duplicate keys, using a function to combine the values.
283 | ||| Uses the ordering of the first map given.
284 | export
285 | mergeWith : (v -> v -> v) -> StringMap v -> StringMap v -> StringMap v
286 | mergeWith f x y = insertFrom inserted x where
287 |   inserted : List (Key, v)
288 |   inserted = do
289 |     (k, v) <- StringMap.toList y
290 |     let v' = (maybe id f $ lookup k x) v
291 |     pure (k, v')
292 |
293 | ||| Merge two maps using the Semigroup (and by extension, Monoid) operation.
294 | ||| Uses mergeWith internally, so the ordering of the left map is kept.
295 | export
296 | merge : Semigroup v => StringMap v -> StringMap v -> StringMap v
297 | merge = mergeWith (<+>)
298 |
299 | ||| Left-biased merge, also keeps the ordering specified by the left map.
300 | export
301 | mergeLeft : StringMap v -> StringMap v -> StringMap v
302 | mergeLeft x y = mergeWith const x y
303 |
304 | export
305 | adjust : String -> (v -> v) -> StringMap v -> StringMap v
306 | adjust k f m =
307 |   case lookup k m of
308 |     Nothing => m
309 |     Just v => insert k (f v) m
310 |
311 | export
312 | Show v => Show (StringMap v) where
313 |   show m = show $ map {b=String} (\(k,v) => k ++ "->" ++ show v) $ StringMap.toList m
314 |
315 | -- TODO: is this the right variant of merge to use for this? I think it is, but
316 | -- I could also see the advantages of using `mergeLeft`. The current approach is
317 | -- strictly more powerful I believe, because `mergeLeft` can be emulated with
318 | -- the `First` monoid. However, this does require more code to do the same
319 | -- thing.
320 | export
321 | Semigroup v => Semigroup (StringMap v) where
322 |   (<+>) = merge
323 |
324 | export
325 | (Semigroup v) => Monoid (StringMap v) where
326 |   neutral = empty
327 |