0 | module Data.Cozippable
  1 |
  2 | import Control.Monad.Error.Interface
  3 |
  4 | import Data.Colist
  5 | import Data.Colist1
  6 | import Data.Either
  7 | import Data.List.Lazy
  8 | import Data.List1
  9 | import public Data.These
 10 | import Data.SortedSet
 11 | import Data.SortedMap
 12 | import Data.Vect
 13 |
 14 | %default total
 15 |
 16 | ||| The `Cozippable` interface describes how you can combine and split elements of a parameterised type, those elements may contain not equal amount of information.
 17 | |||
 18 | ||| The easiest example is `List`. When you `zip` them using standard `Zippable` interface, only their prefixes of commit length are taken into account,
 19 | ||| or else it is impossible to fulfill the interface.
 20 | ||| Sometimes we may want to deal with such lack of information explicitly.
 21 | ||| That's why zipping function of the `Cozippable` interface takes `These` data type, which is an inclusive "or", i.e. either one of elements, or both of them.
 22 | public export
 23 | interface Cozippable z where
 24 |
 25 |   cozipWith : (These a b -> c) -> z a -> z b -> z c
 26 |
 27 |   cozip : z a -> z b -> z (These a b)
 28 |   cozip = cozipWith id
 29 |
 30 |   uncozipWith : (a -> These b c) -> z a -> These (z b) (z c)
 31 |
 32 |   uncozip : z (These a b) -> These (z a) (z b)
 33 |   uncozip = uncozipWith id
 34 |
 35 | export infixr 6 `cozip` -- same as `zip` from `Data.Zippable`
 36 |
 37 | public export
 38 | cozipWith' : Cozippable z => Functor z => (These a b -> c) -> These (z a) (z b) -> z c
 39 | cozipWith' f (This x)   = map (f . This) x
 40 | cozipWith' f (That x)   = map (f . That) x
 41 | cozipWith' f (Both x y) = cozipWith f x y
 42 |
 43 | public export
 44 | cozip' : Cozippable z => Functor z => These (z a) (z b) -> z (These a b)
 45 | cozip' = cozipWith' id
 46 |
 47 | public export
 48 | [Compose] Cozippable f => Cozippable g => Functor g => Cozippable (f . g) where
 49 |   cozipWith f = cozipWith $ cozipWith' f
 50 |   uncozipWith f = uncozipWith $ uncozipWith f
 51 |
 52 | --- Particular case ---
 53 |
 54 | ||| Cozip two cozippables taking a value from each one, and combining corresponding values if it is present in both using given function
 55 | export %inline
 56 | comergeWith : Cozippable z => (a -> a -> a) -> z a -> z a -> z a
 57 | comergeWith = cozipWith . these id id
 58 |
 59 | public export %inline
 60 | comerge : Semigroup a => Cozippable z => z a -> z a -> z a
 61 | comerge = comergeWith (<+>)
 62 |
 63 | export infixr 6 `comerge` -- same as `cozip` above
 64 |
 65 | --- Particular implementations ---
 66 |
 67 | public export
 68 | Semigroup a => Cozippable (Pair a) where
 69 |   cozipWith f (x, y) (x', y') = (x <+> x', f $ Both y y')
 70 |   uncozipWith f (x, y) = bimap (x,) (x,) $ f y
 71 |
 72 | public export
 73 | Cozippable Maybe where
 74 |   cozipWith _ Nothing  Nothing  = Nothing
 75 |   cozipWith f Nothing  (Just y) = Just $ f $ That y
 76 |   cozipWith f (Just x) Nothing  = Just $ f $ This x
 77 |   cozipWith f (Just x) (Just y) = Just $ f $ Both x y
 78 |
 79 |   uncozipWith f Nothing  = Both Nothing Nothing
 80 |   uncozipWith f (Just x) = bimap Just Just $ f x
 81 |
 82 | -- Prefers left `Left` when both are `Left`s, as default `Zippable` implementation
 83 | public export
 84 | Cozippable (Either a) where
 85 |   cozipWith _ (Left x)  (Left _)  = Left x
 86 |   cozipWith f (Left _)  (Right y) = Right $ f $ That y
 87 |   cozipWith f (Right x) (Left _)  = Right $ f $ This x
 88 |   cozipWith f (Right x) (Right y) = Right $ f $ Both x y
 89 |
 90 |   uncozipWith f (Left x)  = Both (Left x) (Left x)
 91 |   uncozipWith f (Right x) = bimap Right Right $ f x
 92 |
 93 | -- Combines both `Left`s when there are both of them
 94 | public export
 95 | [CombineLeft] Semigroup a => Cozippable (Either a) where
 96 |   cozipWith _ (Left x)  (Left y)  = Left $ x <+> y
 97 |   cozipWith f (Left _)  (Right y) = Right $ f $ That y
 98 |   cozipWith f (Right x) (Left _)  = Right $ f $ This x
 99 |   cozipWith f (Right x) (Right y) = Right $ f $ Both x y
100 |
101 |   uncozipWith f (Left x)  = Both (Left x) (Left x)
102 |   uncozipWith f (Right x) = bimap Right Right $ f x
103 |
104 | public export
105 | Semigroup a => Cozippable (These a) where
106 |   cozipWith f (This x)   (This y)   = This $ x <+> y
107 |   cozipWith f (This x)   (That y)   = That $ f $ That y
108 |   cozipWith f (This x)   (Both y z) = Both (x <+> y) $ f $ That z
109 |   cozipWith f (That x)   (This y)   = Both y $ f $ This x
110 |   cozipWith f (That x)   (That y)   = That $ f $ Both x y
111 |   cozipWith f (That x)   (Both y z) = Both y $ f $ Both x z
112 |   cozipWith f (Both x z) (This y)   = Both (x <+> y) $ f $ This z
113 |   cozipWith f (Both x z) (That y)   = Both x $ f $ Both z y
114 |   cozipWith f (Both x z) (Both y w) = Both (x <+> y) $ f $ Both z w
115 |
116 |   uncozipWith f (This x)   = Both (This x) (This x)
117 |   uncozipWith f (That x)   = bimap That That $ f x
118 |   uncozipWith f (Both x y) = bimap (Both x) (Both x) $ f y
119 |
120 | public export
121 | Cozippable List where
122 |   cozipWith f []      []      = []
123 |   cozipWith f []      (y::ys) = f (That y) :: cozipWith f [] ys
124 |   cozipWith f (x::xs) []      = f (This x) :: cozipWith f xs []
125 |   cozipWith f (x::xs) (y::ys) = f (Both x y) :: cozipWith f xs ys
126 |
127 |   uncozipWith f []      = Both [] []
128 |   uncozipWith f (x::xs) = do
129 |     let sub = uncozipWith f xs
130 |     case f x of
131 |       This y   => mapFst (y::) sub
132 |       That y   => mapSnd (y::) sub
133 |       Both y z => bimap (y::) (z::) sub
134 |
135 | public export
136 | Cozippable SnocList where
137 |   cozipWith f [<]     [<]     = [<]
138 |   cozipWith f [<]     (sy:<y) = cozipWith f [<] sy :< f (That y)
139 |   cozipWith f (sx:<x) [<]     = cozipWith f sx [<] :< f (This x)
140 |   cozipWith f (sx:<x) (sy:<y) = cozipWith f sx sy :< f (Both x y)
141 |
142 |   uncozipWith f [<]     = Both [<] [<]
143 |   uncozipWith f (sx:<x) = do
144 |     let sub = uncozipWith f sx
145 |     case f x of
146 |       This y   => mapFst (:<y) sub
147 |       That y   => mapSnd (:<y) sub
148 |       Both y z => bimap (:<y) (:<z) sub
149 |
150 | public export
151 | Cozippable List1 where
152 |   cozipWith f (x:::xs) (y:::ys) = f (Both x y) ::: cozipWith f xs ys
153 |   uncozipWith f (x:::xs) = case (f x, uncozipWith f xs) of
154 |     (This y  , This zs)         => This $ y:::zs
155 |     (This y  , That [])         => This $ singleton y
156 |     (This y  , That $ z::zs)    => Both (singleton y) (z:::zs)
157 |     (This y  , Both zs [])      => This $ y:::zs
158 |     (This y  , Both zs (w::ws)) => Both (y:::zs) (w:::ws)
159 |     (That y  , This [])         => That $ singleton y
160 |     (That y  , This (z::zs))    => Both (z:::zs) (singleton y)
161 |     (That y  , That zs)         => That $ y:::zs
162 |     (That y  , Both [] ws)      => That $ y:::ws
163 |     (That y  , Both (z::zs) ws) => Both (z:::zs) (y:::ws)
164 |     (Both y w, This zs)         => Both (y:::zs) (singleton w)
165 |     (Both y w, That zs)         => Both (singleton y) (w:::zs)
166 |     (Both y w, Both zs vs)      => Both (y:::zs) (w:::vs)
167 |
168 | public export
169 | [VectMaybe] Cozippable (Vect n . Maybe) where
170 |   cozipWith f []      []      = []
171 |   cozipWith f (x::xs) (y::ys) = cozipWith f x y :: cozipWith @{VectMaybe} f xs ys
172 |
173 |   uncozipWith f []      = Both [] []
174 |   uncozipWith f (x::xs) = do
175 |     let (l, r) = fromBoth Nothing Nothing $ uncozipWith f x
176 |     bimap (l::) (r::) $ uncozipWith @{VectMaybe} f xs
177 |
178 | public export
179 | Cozippable LazyList where
180 |   cozipWith f []      []      = []
181 |   cozipWith f []      (y::ys) = f (That y) :: cozipWith f [] ys
182 |   cozipWith f (x::xs) []      = f (This x) :: cozipWith f xs []
183 |   cozipWith f (x::xs) (y::ys) = f (Both x y) :: cozipWith f xs ys
184 |
185 |   uncozipWith f []      = Both [] []
186 |   uncozipWith f (x::xs) = do
187 |     let left  : Lazy (LazyList b) = fromMaybe [] $ fromThis $ uncozipWith f xs
188 |     let right : Lazy (LazyList c) = fromMaybe [] $ fromThat $ uncozipWith f xs
189 |     case f x of
190 |       This y   => Both (y::left) right
191 |       That y   => Both left (y::right)
192 |       Both y z => Both (y::left) (z::right)
193 |
194 | public export
195 | Ord k => Cozippable (SortedMap k) where
196 |   cozipWith f mx my = SortedMap.fromList $ merge (SortedMap.toList mx) (SortedMap.toList my) where
197 |     merge : List (k, a) -> List (k, b) -> List (k, c)
198 |     merge [] [] = []
199 |     merge [] ys = mapSnd (f . That) <$> ys
200 |     merge xs [] = mapSnd (f . This) <$> xs
201 |     merge xxs@((kx,x)::xs) yys@((ky,y)::ys) =
202 |       if kx < ky then (kx, f $ This x) :: merge xs yys
203 |       else if kx == ky then (kx, f $ Both x y) :: merge xs ys
204 |       else (ky, f $ That y) :: merge xxs ys
205 |
206 |   uncozipWith f mx = do
207 |     let xs = mapSnd f <$> SortedMap.toList mx
208 |     let ls = flip mapMaybe xs $ \kbc => (Builtin.fst kbc,) <$> fromThis (Builtin.snd kbc)
209 |     let rs = flip mapMaybe xs $ \kbc => (Builtin.fst kbc,) <$> fromThat (Builtin.snd kbc)
210 |     Both (fromList ls) (fromList rs)
211 |
212 | public export
213 | [MonadError] MonadError e m => Monoid e => Cozippable m where
214 |   cozipWith f ex ey = liftEither =<< cozipWith f <$> tryError ex <*> tryError ey
215 |
216 |   uncozipWith f ex = do
217 |     let left  = liftEither . maybeToEither (neutral {ty=e}) . fromThis . f =<< ex
218 |     let right = liftEither . maybeToEither (neutral {ty=e}) . fromThat . f =<< ex
219 |     Both left right
220 |