0 | ||| Specialised map from `Fin`s to any type.
  1 | ||| Isomorphic to a vector of appropriate size of maybe's of appropriate type.
  2 | |||
  3 | ||| This should be a replacement for `SortedMap (Fin n) v` working better at least in elaborator scripts.
  4 | module Data.Fin.Map
  5 |
  6 | import Data.List
  7 | import Data.SortedMap
  8 | import Data.String
  9 | import Data.Vect
 10 |
 11 | %default total
 12 |
 13 | export
 14 | data FinMap : Nat -> Type -> Type where
 15 |   MkFM : Vect n (Maybe v) -> FinMap n v
 16 |
 17 | %inline
 18 | unFM : FinMap n v -> Vect n $ Maybe v
 19 | unFM (MkFM bs) = bs
 20 |
 21 | export %inline
 22 | empty : {n : _} -> FinMap n v
 23 | empty = MkFM $ replicate _ Nothing
 24 |
 25 | export %inline
 26 | lookup : Fin n -> FinMap n v -> Maybe v
 27 | lookup i = index i . unFM
 28 |
 29 | export %inline
 30 | lookup' : FinMap n v -> Fin n -> Maybe v
 31 | lookup' = flip lookup
 32 |
 33 | export %inline
 34 | insert : Fin n -> v -> FinMap n v -> FinMap n v
 35 | insert i x = MkFM . replaceAt i (Just x) . unFM
 36 |
 37 | export
 38 | insertWith : (v -> v -> v) -> Fin n -> v -> FinMap n v -> FinMap n v
 39 | insertWith f i x = MkFM . updateAt i (Just . maybe x (f x)) . unFM
 40 |
 41 | public export %inline
 42 | insert' : FinMap n v -> (Fin n, v) -> FinMap n v
 43 | insert' = flip $ uncurry insert
 44 |
 45 | export %inline
 46 | singleton : {n : _} -> Fin n -> v -> FinMap n v
 47 | singleton i x = insert i x empty
 48 |
 49 | export %inline
 50 | insertFrom : Foldable f => f (Fin n, v) -> FinMap n v -> FinMap n v
 51 | insertFrom = flip $ foldl insert'
 52 |
 53 | export %inline
 54 | insertFrom' : Foldable f => FinMap n v -> f (Fin n, v) -> FinMap n v
 55 | insertFrom' = flip insertFrom
 56 |
 57 | export %inline
 58 | insertFromWith : Foldable f => (v -> v -> v) -> f (Fin n, v) -> FinMap n v -> FinMap n v
 59 | insertFromWith = flip . foldl . flip . uncurry . insertWith
 60 |
 61 | export %inline
 62 | delete : Fin n -> FinMap n v -> FinMap n v
 63 | delete i = MkFM . replaceAt i Nothing . unFM
 64 |
 65 | export %inline
 66 | delete' : FinMap n v -> Fin n -> FinMap n v
 67 | delete' = flip delete
 68 |
 69 | export %inline
 70 | update : (Maybe v -> Maybe v) -> Fin n -> FinMap n v -> FinMap n v
 71 | update f i = MkFM . updateAt i f . unFM
 72 |
 73 | public export %inline
 74 | update' : FinMap n v -> (Maybe v -> Maybe v) -> Fin n -> FinMap n v
 75 | update' m f x = update f x m
 76 |
 77 | export %inline
 78 | updateExisting : (v -> v) -> Fin n -> FinMap n v -> FinMap n v
 79 | updateExisting f i = MkFM . updateAt i (map f) . unFM
 80 |
 81 | public export %inline
 82 | updateExisting' : FinMap n v -> (v -> v) -> Fin n -> FinMap n v
 83 | updateExisting' m f x = updateExisting f x m
 84 |
 85 | export %inline
 86 | fromFoldable : Foldable f => {n : _} -> f (Fin n, v) -> FinMap n v
 87 | fromFoldable = insertFrom' empty
 88 |
 89 | export %inline
 90 | fromFoldableWith : Foldable f => {n : _} -> (v -> v -> v) -> f (Fin n, v) -> FinMap n v
 91 | fromFoldableWith f = flip (insertFromWith f) empty
 92 |
 93 | export %inline
 94 | fromList : {n : _} -> List (Fin n, v) -> FinMap n v
 95 | fromList = insertFrom' empty
 96 |
 97 | export %inline
 98 | fromListWith : {n : _} -> (v -> v -> v) -> List (Fin n, v) -> FinMap n v
 99 | fromListWith f = flip (insertFromWith f) empty
100 |
101 | export -- I do this here just not to depend on libraries like `collection-utils`
102 | iList : Vect n a -> List (Fin n, a)
103 | iList [] = []
104 | iList (x::xs) = (FZ, x) :: (mapFst FS <$> iList xs)
105 |
106 | export %inline
107 | kvList : FinMap n v -> List (Fin n, v)
108 | kvList = mapMaybe (\(i, mv) => (i,) <$> mv) . iList . unFM
109 |
110 | public export %inline
111 | (.asKVList) : FinMap n v -> List (Fin n, v)
112 | (.asKVList) = kvList
113 |
114 | export
115 | keys : FinMap n v -> List $ Fin n
116 | keys = map fst . kvList
117 |
118 | export
119 | values : FinMap n v -> List v
120 | values = map snd . kvList
121 |
122 | export
123 | Functor (FinMap n) where
124 |   map f = MkFM . map @{Compose} f . unFM
125 |
126 | export
127 | mapWithKey : (Fin n -> v -> w) -> FinMap n v -> FinMap n w
128 | mapWithKey f = MkFM . vmapI f . unFM where
129 |   vmapI : forall n. (Fin n -> v -> w) -> Vect n (Maybe v) -> Vect n (Maybe w)
130 |   vmapI f []      = []
131 |   vmapI f (x::xs) = map (f FZ) x :: vmapI (f . FS) xs
132 |
133 | export %inline
134 | mapWithKey' : FinMap n v -> (Fin n -> v -> w) -> FinMap n w
135 | mapWithKey' = flip mapWithKey
136 |
137 | export
138 | Foldable (FinMap n) where
139 |   foldr f z = foldr f z . values
140 |   foldl f z = foldl f z . values
141 |   null      = all isNothing . unFM
142 |   foldMap f = foldMap f . values
143 |
144 | export
145 | Traversable (FinMap n) where
146 |   traverse f = map MkFM . traverse @{Compose} f . unFM
147 |
148 | export
149 | traverseWithKey : Applicative m => (Fin n -> v -> m w) -> FinMap n v -> m $ FinMap n w
150 | traverseWithKey f = map MkFM . vmapI f . unFM where
151 |   vmapI : forall n. (Fin n -> v -> m w) -> Vect n (Maybe v) -> m $ Vect n (Maybe w)
152 |   vmapI f []      = [| [] |]
153 |   vmapI f (x::xs) = [| traverse (f FZ) x :: vmapI (f . FS) xs |]
154 |
155 | public export %inline
156 | forWithKey : Applicative m => FinMap n v -> (Fin n -> v -> m w) -> m $ FinMap n w
157 | forWithKey = flip traverseWithKey
158 |
159 | export
160 | Zippable (FinMap n) where
161 |   zipWith f mx my = MkFM $ zipWith (\x, y => f <$> x <*> y) (unFM mx) (unFM my)
162 |   zipWith3 f mx my mz = MkFM $ zipWith3 (\x, y, z => f <$> x <*> y <*> z) (unFM mx) (unFM my) (unFM mz)
163 |   unzipWith f m = let m' = map f m in (map fst m', map snd m')
164 |   unzipWith3 f m = let m' = map f m in (map fst m', map (fst . snd) m', map (snd . snd) m')
165 |
166 | ||| Merge two maps, when encountering duplicate keys, uses a function to combine the values.
167 | export
168 | mergeWith : (v -> v -> v) -> FinMap n v -> FinMap n v -> FinMap n v
169 | mergeWith f = MkFM .: zipWith fm `on` unFM where
170 |   fm : Maybe v -> Maybe v -> Maybe v
171 |   fm Nothing  Nothing  = Nothing
172 |   fm (Just x) Nothing  = Just x
173 |   fm Nothing  (Just y) = Just y
174 |   fm (Just x) (Just y) = Just $ f x y
175 |
176 | public export %inline
177 | merge : Semigroup v => FinMap n v -> FinMap n v -> FinMap n v
178 | merge = mergeWith (<+>)
179 |
180 | public export %inline
181 | mergeLeft : FinMap n v -> FinMap n v -> FinMap n v
182 | mergeLeft = mergeWith const
183 |
184 | export %inline
185 | leftMost : FinMap n v -> Maybe (Fin n, v)
186 | leftMost = head' . kvList
187 |
188 | export %inline
189 | rightMost : FinMap n v -> Maybe (Fin n, v)
190 | rightMost = last' . kvList
191 |
192 | export
193 | Interpolation v => Interpolation (Fin n) => Interpolation (FinMap n v) where
194 |   interpolate = ("{" ++) . (++ "}") . joinBy ", " . map (\(i, v) => "\{i} -> \{v}") . kvList
195 |
196 | export
197 | Eq v => Eq (FinMap n v) where
198 |   (==) = (==) `on` unFM
199 |
200 | export
201 | Semigroup v => Semigroup (FinMap n v) where
202 |   (<+>) = merge
203 |
204 | export
205 | {n : _} -> Semigroup v => Monoid (FinMap n v) where
206 |   neutral = empty
207 |
208 | public export %inline
209 | size : FinMap n v -> Nat
210 | size = length . kvList -- this implementation is to make `asVect` work seamlessly
211 |
212 | public export %inline
213 | (.size) : FinMap n v -> Nat
214 | (.size) = size
215 |
216 | public export %inline
217 | (.asKVVect) : (fs : FinMap n v) -> Vect fs.size (Fin n, v)
218 | (.asKVVect) fs = fromList $ kvList fs
219 |
220 | export
221 | fromSortedMap : {n : _} -> SortedMap (Fin n) v -> FinMap n v
222 | fromSortedMap = fromList . SortedMap.toList
223 | -- we can employ the fact that the returned list must be sorted to make this faster
224 |
225 | export
226 | weakenToSuper : {n : _} -> {i : Fin n} -> FinMap (finToNat i) v -> FinMap n v
227 | weakenToSuper = MkFM . go . unFM where
228 |   go : {n : _} -> {i : Fin n} -> Vect (finToNat i) (Maybe v) -> Vect n (Maybe v)
229 |   go {i=FZ}   []      = replicate _ Nothing
230 |   go {i=FS i} (x::xs) = x :: go xs
231 |