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