0 | module Data.SortedMap
  1 |
  2 | import Data.SortedMap.Dependent
  3 | import Data.Zippable
  4 |
  5 | export
  6 | record SortedMap k v where
  7 |   constructor M
  8 |   unM : SortedDMap k $ const v
  9 |
 10 | -- Helper function
 11 | unDPair : (x : a ** const b x) -> (a, b)
 12 | unDPair (k ** v) = (k, v)
 13 |
 14 | export
 15 | empty : Ord k => SortedMap k v
 16 | empty = M empty
 17 |
 18 | export
 19 | lookup : k -> SortedMap k v -> Maybe v
 20 | lookup k = map snd . lookup k . unM
 21 |
 22 | public export %inline
 23 | lookup' : SortedMap k v -> k -> Maybe v
 24 | lookup' = flip lookup
 25 |
 26 | export
 27 | insert : k -> v -> SortedMap k v -> SortedMap k v
 28 | insert k v = M . insert k v . unM
 29 |
 30 | ||| Inserts a key value pair into a map and merges duplicated values
 31 | ||| with the given function.
 32 | export
 33 | insertWith : (v -> v -> v) -> k -> v -> SortedMap k v -> SortedMap k v
 34 | insertWith f k v xs =
 35 |   case lookup k xs of
 36 |     Just x  => insert k (f v x) xs
 37 |     Nothing => insert k v xs
 38 |
 39 | public export %inline
 40 | insert' : SortedMap k v -> (k, v) -> SortedMap k v
 41 | insert' = flip $ uncurry insert
 42 |
 43 | export
 44 | singleton : Ord k => k -> v -> SortedMap k v
 45 | singleton = M .: singleton
 46 |
 47 | export
 48 | insertFrom : Foldable f => f (k, v) -> SortedMap k v -> SortedMap k v
 49 | insertFrom = flip $ foldl insert'
 50 |
 51 | public export %inline
 52 | insertFrom' : Foldable f => SortedMap k v -> f (k, v) -> SortedMap k v
 53 | insertFrom' = flip insertFrom
 54 |
 55 | ||| Inserts any foldable of a key value pair into a map and merges duplicated
 56 | ||| values with the given function.
 57 | export
 58 | insertFromWith : Foldable f => (v -> v -> v) -> f (k, v) -> SortedMap k v -> SortedMap k v
 59 | insertFromWith f = flip $ foldl $ flip $ uncurry $ insertWith f
 60 |
 61 | export
 62 | delete : k -> SortedMap k v -> SortedMap k v
 63 | delete k = M . delete k . unM
 64 |
 65 | public export %inline
 66 | delete' : SortedMap k v -> k -> SortedMap k v
 67 | delete' = flip delete
 68 |
 69 | ||| Updates or deletes a value based on the decision function
 70 | |||
 71 | ||| The decision function takes information about the presence of the value,
 72 | ||| and the value itself, if it is present.
 73 | ||| It returns a new value or the fact that there should be no value as the result.
 74 | |||
 75 | ||| The current implementation performs up to two traversals of the original map
 76 | export
 77 | update : (Maybe v -> Maybe v) -> k -> SortedMap k v -> SortedMap k v
 78 | update f k m = case f $ lookup k m of
 79 |   Just v  => insert k v m
 80 |   Nothing => delete k m
 81 |
 82 | public export %inline
 83 | update' : SortedMap k v -> (Maybe v -> Maybe v) -> k -> SortedMap k v
 84 | update' m f x = update f x m
 85 |
 86 | ||| Updates existing value, if it is present, and does nothing otherwise
 87 | |||
 88 | ||| The current implementation performs up to two traversals of the original map
 89 | export
 90 | updateExisting : (v -> v) -> k -> SortedMap k v -> SortedMap k v
 91 | updateExisting f k m = case lookup k m of
 92 |   Just v  => insert k (f v) m
 93 |   Nothing => m
 94 |
 95 | public export %inline
 96 | updateExisting' : SortedMap k v -> (v -> v) -> k -> SortedMap k v
 97 | updateExisting' m f x = updateExisting f x m
 98 |
 99 | export
100 | fromList : Ord k => List (k, v) -> SortedMap k v
101 | fromList = insertFrom' empty
102 |
103 | ||| Converts a list of key-value pairs into a map and merges duplicated
104 | ||| values with the given function.
105 | export
106 | fromListWith : Ord k => (v -> v -> v) -> List (k, v) -> SortedMap k v
107 | fromListWith f = flip (insertFromWith f) empty
108 |
109 | export
110 | toList : SortedMap k v -> List (k, v)
111 | toList = map unDPair . kvList . unM
112 |
113 | ||| Returns a list of key-value pairs stored in this map
114 | public export %inline
115 | kvList : SortedMap k v -> List (k, v)
116 | kvList = toList
117 |
118 | ||| Gets the keys of the map.
119 | export
120 | keys : SortedMap k v -> List k
121 | keys = map fst . kvList
122 |
123 | ||| Gets the values of the map. Could contain duplicates.
124 | export
125 | values : SortedMap k v -> List v
126 | values = map snd . kvList
127 |
128 | export
129 | implementation Functor (SortedMap k) where
130 |   map f = M . map f . unM
131 |
132 | export
133 | implementation Foldable (SortedMap k) where
134 |   foldr f z = foldr f z . values
135 |   foldl f z = foldl f z . values
136 |
137 |   null = null . unM
138 |
139 |   foldMap f = foldMap f . values
140 |
141 | export
142 | implementation Traversable (SortedMap k) where
143 |   traverse f = map M . traverse f . unM
144 |
145 | export
146 | implementation Ord k => Zippable (SortedMap k) where
147 |   zipWith f mx my = fromList $ mapMaybe (\(k, x) => (k,) . f x <$> lookup k my) $ kvList mx
148 |   zipWith3 f mx my mz = fromList $ mapMaybe (\(k, x) => (k,) .: f x <$> lookup k my <*> lookup k mz) $ kvList mx
149 |   unzipWith f m = let m' = map f m in (map fst m', map snd m')
150 |   unzipWith3 f m = let m' = map f m in (map fst m', map (fst . snd) m', map (snd . snd) m')
151 |
152 | ||| Merge two maps. When encountering duplicate keys, using a function to combine the values.
153 | ||| Uses the ordering of the first map given.
154 | export
155 | mergeWith : (v -> v -> v) -> SortedMap k v -> SortedMap k v -> SortedMap k v
156 | mergeWith f x y = insertFrom inserted x where
157 |   inserted : List (k, v)
158 |   inserted = do
159 |     (k, v) <- kvList y
160 |     let v' = (maybe id f $ lookup k x) v
161 |     pure (k, v')
162 |
163 | ||| Merge two maps using the Semigroup (and by extension, Monoid) operation.
164 | ||| Uses mergeWith internally, so the ordering of the left map is kept.
165 | export
166 | merge : Semigroup v => SortedMap k v -> SortedMap k v -> SortedMap k v
167 | merge = mergeWith (<+>)
168 |
169 | ||| Left-biased merge, also keeps the ordering specified  by the left map.
170 | export
171 | mergeLeft : SortedMap k v -> SortedMap k v -> SortedMap k v
172 | mergeLeft = mergeWith const
173 |
174 | ||| looks up a key in map, returning the left and right closest values, so that
175 | |||  k1 <= k < k2. If at the end of the beginning and/or end of the sorted map, returns
176 | ||| nothing appropriately
177 | export
178 | lookupBetween : key -> SortedMap key val -> (Maybe (key,val), Maybe (key,val))
179 | lookupBetween k = bimap (map unDPair) (map unDPair) . lookupBetween k . unM
180 |
181 |
182 | ||| Returns the leftmost (least) key and value
183 | export
184 | leftMost : SortedMap key val -> Maybe (key,val)
185 | leftMost = map unDPair . leftMost . unM
186 |
187 |
188 | ||| Returns the rightmost (greatest) key and value
189 | export
190 | rightMost : SortedMap key val -> Maybe (key,val)
191 | rightMost = map unDPair . rightMost . unM
192 |
193 |
194 | ||| Pops the leftmost key and corresponding value from the map
195 | export
196 | pop : SortedMap k v -> Maybe ((k, v), SortedMap k v)
197 | pop = map (bimap unDPair M) . pop . unM
198 |
199 | export
200 | (Show k, Show v) => Show (SortedMap k v) where
201 |    show m = "fromList " ++ show (kvList m)
202 |
203 | export
204 | (Eq k, Eq v) => Eq (SortedMap k v) where
205 |   (==) = (==) `on` kvList
206 |
207 | -- TODO: is this the right variant of merge to use for this? I think it is, but
208 | -- I could also see the advantages of using `mergeLeft`. The current approach is
209 | -- strictly more powerful I believe, because `mergeLeft` can be emulated with
210 | -- the `First` monoid. However, this does require more code to do the same
211 | -- thing.
212 | export
213 | Semigroup v => Semigroup (SortedMap k v) where
214 |   (<+>) = merge
215 |
216 | ||| For `neutral <+> y`, y is rebuilt in `Ord k`, so this is not a "strict" Monoid.
217 | ||| However, semantically, it should be equal.
218 | export
219 | (Ord k, Semigroup v) => Monoid (SortedMap k v) where
220 |   neutral = empty
221 |
222 | export %inline
223 | Cast (SortedDMap k (const v)) (SortedMap k v) where
224 |   cast = M
225 |
226 | export %inline
227 | Cast (SortedMap k v) (SortedDMap k (const v)) where
228 |   cast = unM
229 |