0 | module Data.HashMap.SparseArray
  1 |
  2 | import Data.HashMap.Bits
  3 | import public Data.HashMap.Array
  4 | import Data.List
  5 |
  6 | %default total
  7 |
  8 | ||| An array storing up to 64 elements
  9 | public export
 10 | record SparseArray a where
 11 |     constructor MkSparseArray
 12 |     bitmap : Bits64
 13 |     array : Array a
 14 |
 15 | %name SparseArray arr
 16 |
 17 | export
 18 | empty : SparseArray a
 19 | empty = MkSparseArray 0 empty
 20 |
 21 | bits32SortNub : List (Bits32, a) -> List (Bits32, a) -> List (Bits32, a)
 22 | bits32SortNub [] acc = acc
 23 | bits32SortNub lst@(x :: xs) acc =
 24 |     let (lt, gt) = seperate (fst x) xs [] []
 25 |      in bits32SortNub (assert_smaller lst lt) $ x :: bits32SortNub (assert_smaller lst gt) acc
 26 |   where
 27 |     seperate : Bits32 -> List (Bits32, a) -> (lt : List (Bits32, a)) -> (gt : List (Bits32, a)) -> (List (Bits32, a), List (Bits32, a))
 28 |     seperate x [] lt gt = (lt, gt)
 29 |     seperate x (y :: xs) lt gt = case compare (fst y) x of
 30 |         LT => seperate x xs (y :: lt) gt
 31 |         EQ => seperate x xs lt gt
 32 |         GT => seperate x xs lt (y :: gt)
 33 |
 34 | export %inline
 35 | singleton : (Bits32, a) -> SparseArray a
 36 | singleton (k, v) = MkSparseArray (bit k) (Array.singleton v)
 37 |
 38 | export
 39 | doubleton : (Bits32, a) -> (Bits32, a) -> SparseArray a
 40 | doubleton (k0, v0) (k1, v1) = case compare k0 k1 of
 41 |     LT => MkSparseArray (bit k0 .|. bit k1) (fromList [v0, v1])
 42 |     EQ => singleton (k1, v1)
 43 |     GT => MkSparseArray (bit k1 .|. bit k0) (fromList [v1, v0])
 44 |
 45 | export
 46 | fromList : List (Bits32, a) -> SparseArray a
 47 | fromList xs =
 48 |     let xs = bits32SortNub xs []
 49 |         bitmap = foldl (\acc, (idx, _) => setBit acc idx) 0 xs
 50 |         arr = Array.fromList (map snd xs)
 51 |      in MkSparseArray bitmap arr
 52 |
 53 | export
 54 | hasEntry : Bits32 -> SparseArray a -> Bool
 55 | hasEntry idx arr = testBit arr.bitmap idx
 56 |
 57 | export
 58 | findIndex : Bits32 -> Bits64 -> Bits32
 59 | findIndex idx bitmap =
 60 |     let mask = oneBits `shiftR` cast (64 - idx)
 61 |      in cast $ popCount $ bitmap .&. mask
 62 |
 63 | export
 64 | index : (idx : Bits32) -> (arr : SparseArray a) -> Maybe a
 65 | index idx arr = if hasEntry idx arr
 66 |     then
 67 |         let arrIdx = findIndex idx arr.bitmap
 68 |          in Array.index arr.array arrIdx
 69 |     else Nothing
 70 |
 71 | export
 72 | set : (idx : Bits32) -> (val : a) -> (arr : SparseArray a) -> SparseArray a
 73 | set sparseIdx val arr =
 74 |     if hasEntry sparseIdx arr 
 75 |         then
 76 |             let arrIdx = findIndex sparseIdx arr.bitmap
 77 |              in MkSparseArray
 78 |                 { bitmap = arr.bitmap
 79 |                 , array = update arr.array [(arrIdx, val)]
 80 |                 }
 81 |         else
 82 |             let bitmap = setBit arr.bitmap sparseIdx
 83 |                 arrIdx = findIndex sparseIdx bitmap
 84 |              in MkSparseArray
 85 |                 { bitmap
 86 |                 , array = insert arrIdx val arr.array
 87 |                 }
 88 |
 89 | export
 90 | delete : (idx : Bits32) -> (arr : SparseArray a) -> SparseArray a
 91 | delete idx arr = if hasEntry idx arr
 92 |     then
 93 |         let arrIdx = findIndex idx arr.bitmap
 94 |          in MkSparseArray
 95 |             { bitmap = clearBit arr.bitmap idx
 96 |             , array = delete arrIdx arr.array
 97 |             }
 98 |     else arr
 99 |
100 | export
101 | length : SparseArray a -> Nat
102 | length arr = cast $ popCount arr.bitmap
103 |
104 | export
105 | indexes : SparseArray a -> List Bits32
106 | indexes arr = filter (\idx => hasEntry idx arr) [0..63]
107 |
108 | export
109 | Functor SparseArray where
110 |     map f arr = MkSparseArray arr.bitmap (map f arr.array)
111 |
112 | export
113 | Foldable SparseArray where
114 |     foldr f z arr = foldr f z arr.array
115 |     foldl f z arr = foldl f z arr.array
116 |     null arr = arr.bitmap == 0
117 |     toList arr = toList arr.array
118 |     foldMap f arr = foldMap f arr.array
119 |
120 | export
121 | toList : SparseArray a -> List (Bits32, a)
122 | toList arr = zip (indexes arr) (toList arr)
123 |
124 | export
125 | Show a => Show (SparseArray a) where
126 |     show arr = "[" ++ fastConcat (intersperse ", " (map show (SparseArray.toList arr))) ++ "]"
127 |
128 | export
129 | Eq a => Eq (SparseArray a) where
130 |     x == y =
131 |         length x == length y
132 |         && all (\idx => index idx x == index idx y) (indexes x)
133 |