
   0 | module Libraries.Data.IntMap
   1 | 
   2 | -- Hand specialised map, for efficiency...
   3 | 
   4 | %default total
   5 | 
   6 | Key : Type
   7 | Key = Int
   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 . List (Key, v) -> Tree n v -> List (Key, v)
 197 |     treeToList' rest (Leaf k v) = (k, v) :: rest
 198 |     treeToList' rest (Branch2 t1 _ t2)
 199 |         = treeToList' (treeToList' rest t2) t1
 200 |     treeToList' rest (Branch3 t1 _ t2 _ t3)
 201 |         = treeToList' (treeToList' (treeToList' rest t3) t2) t1
 202 | 
 203 | export
 204 | data IntMap : Type -> Type where
 205 |   Empty : IntMap v
 206 |   M : (n : Nat) -> Tree n v -> IntMap v
 207 | 
 208 | export
 209 | empty : IntMap v
 210 | empty = Empty
 211 | 
 212 | export
 213 | lookup : Int -> IntMap v -> Maybe v
 214 | lookup _ Empty = Nothing
 215 | lookup k (M _ t) = treeLookup k t
 216 | 
 217 | export
 218 | insert : Int -> v -> IntMap v -> IntMap v
 219 | insert k v Empty = M Z (Leaf k v)
 220 | insert k v (M _ t) =
 221 |   case treeInsert k v t of
 222 |     Left t' => (M _ t')
 223 |     Right t' => (M _ t')
 224 | 
 225 | export
 226 | singleton : Int -> v -> IntMap v
 227 | singleton k v = insert k v empty
 228 | 
 229 | export
 230 | insertFrom : List (Int, v) -> IntMap v -> IntMap v
 231 | insertFrom = flip $ foldl $ flip $ uncurry insert
 232 | 
 233 | export
 234 | delete : Int -> IntMap v -> IntMap v
 235 | delete _ Empty = Empty
 236 | delete k (M Z t) =
 237 |   case treeDelete k t of
 238 |     Left t' => (M _ t')
 239 |     Right () => Empty
 240 | delete k (M (S _) t) =
 241 |   case treeDelete k t of
 242 |     Left t' => (M _ t')
 243 |     Right t' => (M _ t')
 244 | 
 245 | export
 246 | fromList : List (Int, v) -> IntMap v
 247 | fromList l = foldl (flip (uncurry insert)) empty l
 248 | 
 249 | export
 250 | toList : IntMap v -> List (Int, v)
 251 | toList Empty = []
 252 | toList (M _ t) = treeToList t
 253 | 
 254 | ||| Gets the Keys of the map.
 255 | export
 256 | keys : IntMap v -> List Int
 257 | keys = map fst . toList
 258 | 
 259 | ||| Gets the values of the map. Could contain duplicates.
 260 | export
 261 | values : IntMap v -> List v
 262 | values = map snd . toList
 263 | 
 264 | treeMap : (a -> b) -> Tree n a -> Tree n b
 265 | treeMap f (Leaf k v) = Leaf k (f v)
 266 | treeMap f (Branch2 t1 k t2) = Branch2 (treeMap f t1) k (treeMap f t2)
 267 | treeMap f (Branch3 t1 k1 t2 k2 t3)
 268 |     = Branch3 (treeMap f t1) k1 (treeMap f t2) k2 (treeMap f t3)
 269 | 
 270 | export
 271 | implementation Functor IntMap where
 272 |   map _ Empty = Empty
 273 |   map f (M n t) = M _ (treeMap f t)
 274 | 
 275 | ||| Merge two maps. When encountering duplicate keys, using a function to combine the values.
 276 | ||| Uses the ordering of the first map given.
 277 | export
 278 | mergeWith : (v -> v -> v) -> IntMap v -> IntMap v -> IntMap v
 279 | mergeWith f x y = insertFrom inserted x where
 280 |   inserted : List (Key, v)
 281 |   inserted = do
 282 |     (k, v) <- toList y
 283 |     let v' = (maybe id f $ lookup k x) v
 284 |     pure (k, v')
 285 | 
 286 | ||| Merge two maps using the Semigroup (and by extension, Monoid) operation.
 287 | ||| Uses mergeWith internally, so the ordering of the left map is kept.
 288 | export
 289 | merge : Semigroup v => IntMap v -> IntMap v -> IntMap v
 290 | merge = mergeWith (<+>)
 291 | 
 292 | ||| Left-biased merge, also keeps the ordering specified by the left map.
 293 | export
 294 | mergeLeft : IntMap v -> IntMap v -> IntMap v
 295 | mergeLeft = mergeWith const
 296 | 
 297 | -- TODO: is this the right variant of merge to use for this? I think it is, but
 298 | -- I could also see the advantages of using `mergeLeft`. The current approach is
 299 | -- strictly more powerful I believe, because `mergeLeft` can be emulated with
 300 | -- the `First` monoid. However, this does require more code to do the same
 301 | -- thing.
 302 | export
 303 | Semigroup v => Semigroup (IntMap v) where
 304 |   (<+>) = merge
 305 | 
 306 | export
 307 | (Semigroup v) => Monoid (IntMap v) where
 308 |   neutral = empty
 309 |