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