
   0 | module Data.BitMap
   1 | 
   2 | %default total
   3 | 
   4 | --------------------------------------------------------------------------------
   5 | --          Utilities
   6 | --------------------------------------------------------------------------------
   7 | 
   8 | public export
   9 | Key : Type
  10 | Key = Bits32
  11 | 
  12 | %inline
  13 | mask : Bits32 -> Bits32
  14 | mask k = prim__and_Bits32 k 3
  15 | 
  16 | %inline %tcinline
  17 | shr2 : Bits32 -> Bits32
  18 | shr2 k = assert_smaller k $ prim__shr_Bits32 k 2
  19 | 
  20 | %inline
  21 | shl2 : Bits32 -> Bits32
  22 | shl2 k = prim__shl_Bits32 k 2
  23 | 
  24 | --------------------------------------------------------------------------------
  25 | --          BitMap
  26 | --------------------------------------------------------------------------------
  27 | 
  28 | export
  29 | data BitMap : Type -> Type where
  30 |   Empty  : BitMap a
  31 |   Leaf   : (v : a) -> BitMap a
  32 |   Branch : (b0,b1,b2,b3 : BitMap a) -> BitMap a
  33 | 
  34 | --------------------------------------------------------------------------------
  35 | --          Inspecting IntMaps
  36 | --------------------------------------------------------------------------------
  37 | 
  38 | export
  39 | isEmpty : BitMap v -> Bool
  40 | isEmpty Empty = True
  41 | isEmpty _     = False
  42 | 
  43 | export
  44 | nonEmpty : BitMap v -> Bool
  45 | nonEmpty Empty = False
  46 | nonEmpty _     = True
  47 | 
  48 | export
  49 | lookup : (k : Key) -> BitMap a -> Maybe a
  50 | lookup 0 (Leaf v) = Just v
  51 | lookup k (Branch b0 b1 b2 b3) = case mask k of
  52 |   0 => lookup (shr2 k) b0
  53 |   1 => lookup (shr2 k) b1
  54 |   2 => lookup (shr2 k) b2
  55 |   _ => lookup (shr2 k) b3
  56 | lookup _ _ = Nothing
  57 | 
  58 | export
  59 | pairs : BitMap v -> List (Key,v)
  60 | pairs = go 1 0
  61 | 
  62 |   where
  63 |     go : Bits32 -> Key -> BitMap v ->  List (Key,v)
  64 |     go p k Empty = []
  65 |     go p k (Leaf x) = [(k,x)]
  66 |     go p k (Branch b0 b1 b2 b3) =
  67 |       let p2 = shl2 p
  68 |        in go p2 k           b0 ++
  69 |           go p2 (k + p * 1) b1 ++
  70 |           go p2 (k + p * 2) b2 ++
  71 |           go p2 (k + p * 3) b3
  72 | 
  73 | export
  74 | foldlKV : (acc -> Key -> v -> acc) -> acc -> BitMap v -> acc
  75 | foldlKV f = go 1 0
  76 | 
  77 |   where
  78 |     go : Bits32 -> Key -> acc -> BitMap v ->  acc
  79 |     go p k x0 (Branch b0 b1 b2 b3) =
  80 |       let p2 := shl2 p
  81 |           x1 := go p2 k           x0 b0
  82 |           x2 := go p2 (k + p * 1) x1 b1
  83 |           x3 := go p2 (k + p * 2) x2 b2
  84 |        in go p2 (k + p * 3) x3 b3
  85 |     go p k x Empty    = x
  86 |     go p k x (Leaf y) = f x k y
  87 | 
  88 | export
  89 | traverseWithKey : Applicative f => (Key -> v -> f w) -> BitMap v -> f (BitMap w)
  90 | traverseWithKey g = go 1 0
  91 | 
  92 |   where
  93 |     go : Bits32 -> Key -> BitMap v -> f (BitMap w)
  94 |     go p k (Branch b0 b1 b2 b3) =
  95 |       let p2 := shl2 p
  96 |        in [| Branch
  97 |              (go p2   k b0)
  98 |              (go p2 (k + p) b1)
  99 |              (go p2 (k + p * 2) b2)
 100 |              (go p2 (k + p * 3) b3)
 101 |           |]
 102 |     go p k Empty    = pure Empty
 103 |     go p k (Leaf y) = Leaf <$> g k y
 104 | 
 105 | export
 106 | mapWithKey : (Key -> v -> w) -> BitMap v -> BitMap w
 107 | mapWithKey g = go 1 0
 108 | 
 109 |   where
 110 |     go : Bits32 -> Key -> BitMap v -> BitMap w
 111 |     go p k (Branch b0 b1 b2 b3) =
 112 |       let p2 := shl2 p
 113 |        in Branch
 114 |             (go p2   k b0)
 115 |             (go p2 (k + p) b1)
 116 |             (go p2 (k + p * 2) b2)
 117 |             (go p2 (k + p * 3) b3)
 118 |     go p k Empty    = Empty
 119 |     go p k (Leaf y) = Leaf $ g k y
 120 | 
 121 | ||| Gets the keys of the map.
 122 | export
 123 | keys : BitMap v -> List Key
 124 | keys = map fst . pairs
 125 | 
 126 | export
 127 | values : BitMap v -> List v
 128 | values = map snd . pairs
 129 | 
 130 | --------------------------------------------------------------------------------
 131 | --          Creating IntMaps
 132 | --------------------------------------------------------------------------------
 133 | 
 134 | export
 135 | empty : BitMap v
 136 | empty = Empty
 137 | 
 138 | public export
 139 | singleton : (k : Key) -> (v : a) -> BitMap a
 140 | singleton 0 v = Leaf v
 141 | singleton k v = case mask k of
 142 |   0 => Branch (singleton (shr2 k) v) Empty Empty Empty
 143 |   1 => Branch Empty (singleton (shr2 k) v) Empty Empty
 144 |   2 => Branch Empty Empty (singleton (shr2 k) v) Empty
 145 |   _ => Branch Empty Empty Empty (singleton (shr2 k) v)
 146 | 
 147 | branch : (t0,t1,t2,t3 : BitMap a) -> BitMap a
 148 | branch Empty Empty Empty Empty = Empty
 149 | branch t0    t1    t2    t3    = Branch t0 t1 t2 t3
 150 | 
 151 | export
 152 | mapMaybe : (a -> Maybe b) -> BitMap a -> BitMap b
 153 | mapMaybe f Empty    = Empty
 154 | mapMaybe f (Leaf v) = maybe Empty Leaf (f v)
 155 | mapMaybe f (Branch b0 b1 b2 b3) =
 156 |   branch (mapMaybe f b0) (mapMaybe f b1) (mapMaybe f b2) (mapMaybe f b3)
 157 | 
 158 | export
 159 | filter : (a -> Bool) -> BitMap a -> BitMap a
 160 | filter f Empty    = Empty
 161 | filter f (Leaf v) = if f v then Leaf v else Empty
 162 | filter f (Branch b0 b1 b2 b3) =
 163 |   branch (filter f b0) (filter f b1) (filter f b2) (filter f b3)
 164 | 
 165 | --------------------------------------------------------------------------------
 166 | --          Insert and Update
 167 | --------------------------------------------------------------------------------
 168 | 
 169 | insertAtLeaf : (k : Key) -> (v1,v2 : a) -> BitMap a
 170 | insertAtLeaf 0 v1 v2 = Leaf v1
 171 | insertAtLeaf k v1 v2 = case mask k of
 172 |   0 => Branch (insertAtLeaf (shr2 k) v1 v2) Empty Empty Empty
 173 |   1 => Branch (Leaf v2) (singleton (shr2 k) v1) Empty Empty
 174 |   2 => Branch (Leaf v2) Empty (singleton (shr2 k) v1) Empty
 175 |   _ => Branch (Leaf v2) Empty Empty (singleton (shr2 k) v1)
 176 | 
 177 | export
 178 | insert : (k : Key) -> (v : a) -> BitMap a -> BitMap a
 179 | insert k v (Branch b0 b1 b2 b3) = case mask k of
 180 |   0 => Branch (insert (shr2 k) v b0) b1 b2 b3
 181 |   1 => Branch b0 (insert (shr2 k) v b1) b2 b3
 182 |   2 => Branch b0 b1 (insert (shr2 k) v b2) b3
 183 |   _ => Branch b0 b1 b2 (insert (shr2 k) v b3)
 184 | insert k v (Leaf v2) = insertAtLeaf k v v2
 185 | insert k v Empty     = singleton k v
 186 | 
 187 | export
 188 | fromList : List (Key,v) -> BitMap v
 189 | fromList = foldl (\m,(k,v) => insert k v m) empty
 190 | 
 191 | insertAtLeafWith : (f : a -> a -> a) -> (k : Key) -> (v1,v2 : a) -> BitMap a
 192 | insertAtLeafWith f 0 v1 v2 = Leaf $ f v1 v2
 193 | insertAtLeafWith f k v1 v2 = case mask k of
 194 |   0 => Branch (insertAtLeafWith f (shr2 k) v1 v2) Empty Empty Empty
 195 |   1 => Branch (Leaf v2) (singleton (shr2 k) v1) Empty Empty
 196 |   2 => Branch (Leaf v2) Empty (singleton (shr2 k) v1) Empty
 197 |   _ => Branch (Leaf v2) Empty Empty (singleton (shr2 k) v1)
 198 | 
 199 | export
 200 | insertWith : (a -> a -> a) -> (k : Key) -> (v : a) -> BitMap a -> BitMap a
 201 | insertWith f k v (Branch b0 b1 b2 b3) = case mask k of
 202 |   0 => Branch (insertWith f (shr2 k) v b0) b1 b2 b3
 203 |   1 => Branch b0 (insertWith f (shr2 k) v b1) b2 b3
 204 |   2 => Branch b0 b1 (insertWith f (shr2 k) v b2) b3
 205 |   _ => Branch b0 b1 b2 (insertWith f (shr2 k) v b3)
 206 | insertWith f k v (Leaf v2) = insertAtLeafWith f k v v2
 207 | insertWith f k v Empty     = singleton k v
 208 | 
 209 | export
 210 | update : (k : Key) -> (f : a -> a) -> BitMap a -> BitMap a
 211 | update 0 f (Leaf v) = Leaf $ f v
 212 | update k f (Branch b0 b1 b2 b3) = case mask k of
 213 |   0 => Branch (update (shr2 k) f b0) b1 b2 b3
 214 |   1 => Branch b0 (update (shr2 k) f b1) b2 b3
 215 |   2 => Branch b0 b1 (update (shr2 k) f b2) b3
 216 |   _ => Branch b0 b1 b2 (update (shr2 k) f b3)
 217 | update _ _ t = t
 218 | 
 219 | union0 : a -> BitMap a -> BitMap a
 220 | union0 x Empty = Leaf x
 221 | union0 x (Leaf v) = Leaf x
 222 | union0 x (Branch b0 b1 b2 b3) = Branch (union0 x b0) b1 b2 b3
 223 | 
 224 | export
 225 | union : BitMap a -> BitMap a -> BitMap a
 226 | union (Branch b0 b1 b2 b3) (Branch c0 c1 c2 c3) =
 227 |   Branch (union b0 c0) (union b1 c1) (union b2 c2) (union b3 c3)
 228 | union (Branch b0 b1 b2 b3) (Leaf v) = Branch (union0 v b0) b1 b2 b3
 229 | union (Leaf v) (Branch b0 b1 b2 b3) = Branch (union0 v b0) b1 b2 b3
 230 | union (Leaf v) (Leaf _)             = Leaf v
 231 | union Empty y = y
 232 | union x Empty = x
 233 | 
 234 | export
 235 | delete : (k : Key) -> BitMap a -> BitMap a
 236 | delete 0 (Leaf v) = Empty
 237 | delete k (Branch b0 b1 b2 b3) = case mask k of
 238 |   0 => branch (delete (shr2 k) b0) b1 b2 b3
 239 |   1 => branch b0 (delete (shr2 k) b1) b2 b3
 240 |   2 => branch b0 b1 (delete (shr2 k) b2) b3
 241 |   _ => branch b0 b1 b2 (delete (shr2 k) b3)
 242 | delete k m = m
 243 | 
 244 | --------------------------------------------------------------------------------
 245 | --          Decomposing IntMaps
 246 | --------------------------------------------------------------------------------
 247 | 
 248 | public export
 249 | data DecompAt : Type -> Type where
 250 |   NoMatchAt : DecompAt a
 251 |   MatchAt   : (v : a) -> (rem : BitMap a) -> DecompAt a
 252 | 
 253 | export
 254 | decompAt : (k : Key) -> BitMap a -> DecompAt a
 255 | decompAt 0 (Leaf v) = MatchAt v Empty
 256 | decompAt k (Branch b0 b1 b2 b3) = case mask k of
 257 |   0 => case decompAt (shr2 k) b0 of
 258 |     MatchAt v rem => MatchAt v (branch rem b1 b2 b3)
 259 |     NoMatchAt     => NoMatchAt
 260 |   1 => case decompAt (shr2 k) b1 of
 261 |     MatchAt v rem => MatchAt v (branch b0 rem b2 b3)
 262 |     NoMatchAt     => NoMatchAt
 263 |   2 => case decompAt (shr2 k) b2 of
 264 |     MatchAt v rem => MatchAt v (branch b0 b1 rem b3)
 265 |     NoMatchAt     => NoMatchAt
 266 |   _ => case decompAt (shr2 k) b3 of
 267 |     MatchAt v rem => MatchAt v (branch b0 b1 b2 rem)
 268 |     NoMatchAt     => NoMatchAt
 269 | decompAt _ _ = NoMatchAt
 270 | 
 271 | public export
 272 | data Decomp : Type -> Type where
 273 |   NoMatch : Decomp a
 274 |   Match   : (k : Key) -> (v : a) -> (rem : BitMap a) -> Decomp a
 275 | 
 276 | export
 277 | decomp : BitMap a -> Decomp a
 278 | decomp (Leaf v) = Match 0 v Empty
 279 | decomp (Branch b0 b1 b2 b3) = case decomp b3 of
 280 |   (Match k v rem) => Match (shl2 k + 3) v (branch b0 b1 b2 rem)
 281 |   NoMatch           => case decomp b2 of
 282 |     (Match k v rem) => Match (shl2 k + 2) v (branch b0 b1 rem b3)
 283 |     NoMatch           => case decomp b1 of
 284 |       (Match k v rem) => Match (shl2 k + 1) v (branch b0 rem b2 b3)
 285 |       NoMatch           => case decomp b0 of
 286 |         (Match k v rem) => Match (shl2 k) v (branch rem b1 b2 b3)
 287 |         NoMatch           => NoMatch
 288 | 
 289 | decomp Empty = NoMatch
 290 | 
 291 | --------------------------------------------------------------------------------
 292 | --          Interfaces
 293 | --------------------------------------------------------------------------------
 294 | 
 295 | export
 296 | Eq v => Eq (BitMap v) where
 297 |   (==) = (==) `on` pairs
 298 | 
 299 | export
 300 | Show v => Show (BitMap v) where
 301 |   showPrec p m = showCon p "fromList" $ showArg (pairs m)
 302 | 
 303 | export
 304 | Functor BitMap where
 305 |   map _ Empty                = Empty
 306 |   map f (Leaf v)             = Leaf $ f v
 307 |   map f (Branch b0 b1 b2 b3) =
 308 |     Branch (map f b0) (map f b1) (map f b2) (map f b3)
 309 | 
 310 | export
 311 | Foldable BitMap where
 312 |   foldr _ acc Empty = acc
 313 |   foldr f acc (Leaf v) = f v acc
 314 |   foldr f acc (Branch b0 b1 b2 b3) =
 315 |     let acc2 := foldr f acc  b3
 316 |         acc1 := foldr f acc2 b2
 317 |         acc0 := foldr f acc1 b1
 318 |      in foldr f acc0 b0
 319 | 
 320 |   foldl _ acc Empty = acc
 321 |   foldl f acc (Leaf v) = f acc v
 322 |   foldl f acc (Branch b0 b1 b2 b3) =
 323 |     let acc1 := foldl f acc  b0
 324 |         acc2 := foldl f acc1 b1
 325 |         acc3 := foldl f acc2 b2
 326 |      in foldl f acc3 b3
 327 | 
 328 |   foldMap _ Empty = neutral
 329 |   foldMap f (Leaf v) = f v
 330 |   foldMap f (Branch b0 b1 b2 b3) =
 331 |     foldMap f b0 <+> foldMap f b1 <+> foldMap f b2 <+> foldMap f b3
 332 | 
 333 |   foldlM _ acc Empty    = pure acc
 334 |   foldlM f acc (Leaf v) = f acc v
 335 |   foldlM f acc (Branch b0 b1 b2 b3) = do
 336 |     acc1 <- foldlM f acc  b0
 337 |     acc2 <- foldlM f acc1 b1
 338 |     acc3 <- foldlM f acc2 b2
 339 |     foldlM f acc3 b3
 340 | 
 341 |   null m  = isEmpty m
 342 | 
 343 | export
 344 | Traversable BitMap where
 345 |   traverse f Empty = pure Empty
 346 |   traverse f (Leaf v) = Leaf <$> f v
 347 |   traverse f (Branch b0 b1 b2 b3) =
 348 |     [| Branch
 349 |          (traverse f b0)
 350 |          (traverse f b1)
 351 |          (traverse f b2)
 352 |          (traverse f b3)
 353 |     |]
 354 |