1-- | Various zipping and unzipping functions for chunked data structures.
2module Data.ChunkedZip where
3
4import Prelude hiding (zipWith, zipWith3)
5import Control.Arrow ((&&&), (***))
6import qualified Data.List as List
7import qualified Data.List.NonEmpty as NonEmpty
8import Data.List.NonEmpty (NonEmpty)
9import qualified Data.Vector as Vector
10-- import qualified Data.Vector.Unboxed as UVector
11import qualified Data.Sequence as Seq
12import Control.Monad.Trans.Identity
13import Control.Monad.Trans.Reader
14import qualified Data.IntMap as IntMap
15import Data.Tree
16import Data.Functor.Compose
17import Data.Foldable (toList)
18
19class Functor f => Zip f where
20    zipWith :: (a -> b -> c) -> f a -> f b -> f c
21
22    zip :: f a -> f b -> f (a, b)
23    zip = zipWith (,)
24
25    zap :: f (a -> b) -> f a -> f b
26    zap = zipWith id
27
28    unzip :: f (a, b) -> (f a, f b)
29    unzip = fmap fst &&& fmap snd
30
31instance Zip [] where
32    zip = List.zip
33    zipWith = List.zipWith
34    unzip = List.unzip
35instance Zip NonEmpty where
36    zipWith = NonEmpty.zipWith
37    zip = NonEmpty.zip
38    unzip = NonEmpty.unzip
39instance Zip Seq.Seq where
40    zip = Seq.zip
41    zipWith = Seq.zipWith
42    unzip = (Seq.fromList *** Seq.fromList) . List.unzip . toList
43instance Zip Tree where
44    zipWith f (Node a as) (Node b bs) = Node (f a b) (zipWith (zipWith f) as bs)
45instance Zip Vector.Vector where
46    zip = Vector.zip
47    unzip = Vector.unzip
48    zipWith = Vector.zipWith
49  {-
50instance Zip UVector where
51    zip = UVector.zip
52    unzip = UVector.unzip
53    zipWith = UVector.zipWith
54    -}
55
56instance Zip m => Zip (IdentityT m) where
57    zipWith f (IdentityT m) (IdentityT n) = IdentityT (zipWith f m n)
58instance Zip ((->)a) where
59    zipWith f g h a = f (g a) (h a)
60instance Zip m => Zip (ReaderT e m) where
61    zipWith f (ReaderT m) (ReaderT n) = ReaderT $ \a ->
62      zipWith f (m a) (n a)
63instance Zip IntMap.IntMap where
64    zipWith = IntMap.intersectionWith
65instance (Zip f, Zip g) => Zip (Compose f g) where
66    zipWith f (Compose a) (Compose b) = Compose $ zipWith (zipWith f) a b
67
68class Functor f => Zip3 f where
69    zipWith3 :: (a -> b -> c -> d) -> f a -> f b -> f c -> f d
70
71    zip3 :: f a -> f b -> f c -> f (a, b, c)
72    zip3 = zipWith3 (\x y z -> (x,y,z))
73
74    zap3 :: f (a -> b -> c) -> f a -> f b -> f c
75    zap3 = zipWith3 id
76
77    unzip3 :: f (a, b, c) -> (f a, f b, f c)
78    -- unzip3 = fmap (\(x,_,_)->x) &&& fmap (\(_,x,_)->x) &&& fmap (\(_,_,x)->x)
79
80instance Zip3 [] where
81    zip3 = List.zip3
82    unzip3 = List.unzip3
83    zipWith3 = List.zipWith3
84instance Zip3 Vector.Vector where
85    zip3 = Vector.zip3
86    unzip3 = Vector.unzip3
87    zipWith3 = Vector.zipWith3
88instance Zip3 Seq.Seq where
89    zip3 = Seq.zip3
90    unzip3 = (\(a, b, c) -> (Seq.fromList a, Seq.fromList b, Seq.fromList c)) . List.unzip3 . toList
91    zipWith3 = Seq.zipWith3
92
93class Functor f => Zip4 f where
94    zipWith4 :: (a -> b -> c -> d -> e) -> f a -> f b -> f c -> f d -> f e
95
96    zip4 :: f a -> f b -> f c -> f d ->  f (a, b, c, d)
97    zip4 = zipWith4 (\w x y z -> (w, x,y,z))
98
99    zap4 :: f (a -> b -> c -> d) -> f a -> f b -> f c -> f d
100    zap4 = zipWith4 id
101
102    unzip4 :: f (a, b, c, d) -> (f a, f b, f c, f d)
103
104instance Zip4 [] where
105    zip4 = List.zip4
106    unzip4 = List.unzip4
107    zipWith4 = List.zipWith4
108instance Zip4 Vector.Vector where
109    zip4 = Vector.zip4
110    unzip4 = Vector.unzip4
111    zipWith4 = Vector.zipWith4
112instance Zip4 Seq.Seq where
113    zip4 = Seq.zip4
114    unzip4 = (\(a, b, c, d) -> (Seq.fromList a, Seq.fromList b, Seq.fromList c, Seq.fromList d)) . List.unzip4 . toList
115    zipWith4 = Seq.zipWith4
116
117class Functor f => Zip5 f where
118    zipWith5 :: (a -> b -> c -> d -> e -> g) -> f a -> f b -> f c -> f d -> f e -> f g
119
120    zip5 :: f a -> f b -> f c -> f d -> f e -> f (a, b, c, d, e)
121    zip5 = zipWith5 (\v w x y z -> (v,w,x,y,z))
122
123    zap5 :: f (a -> b -> c -> d -> e) -> f a -> f b -> f c -> f d -> f e
124    zap5 = zipWith5 id
125
126    unzip5 :: f (a, b, c, d, e) -> (f a, f b, f c, f d, f e)
127
128instance Zip5 [] where
129    zip5 = List.zip5
130    unzip5 = List.unzip5
131    zipWith5 = List.zipWith5
132instance Zip5 Vector.Vector where
133    zip5 = Vector.zip5
134    unzip5 = Vector.unzip5
135    zipWith5 = Vector.zipWith5
136
137class Functor f => Zip6 f where
138    zipWith6 :: (a -> b -> c -> d -> e -> g -> h) -> f a -> f b -> f c -> f d -> f e -> f g -> f h
139
140    zip6 :: f a -> f b -> f c -> f d -> f e -> f g -> f (a, b, c, d, e, g)
141    zip6 = zipWith6 (\u v w x y z -> (u, v,w,x,y,z))
142
143    zap6 :: f (a -> b -> c -> d -> e -> g) -> f a -> f b -> f c -> f d -> f e -> f g
144    zap6 = zipWith6 id
145
146    unzip6 :: f (a, b, c, d, e, g) -> (f a, f b, f c, f d, f e, f g)
147
148instance Zip6 [] where
149    zip6 = List.zip6
150    unzip6 = List.unzip6
151    zipWith6 = List.zipWith6
152instance Zip6 Vector.Vector where
153    zip6 = Vector.zip6
154    unzip6 = Vector.unzip6
155    zipWith6 = Vector.zipWith6
156
157class Functor f => Zip7 f where
158    zipWith7 :: (a -> b -> c -> d -> e -> g -> h -> i) -> f a -> f b -> f c -> f d -> f e -> f g -> f h -> f i
159
160    zip7 :: f a -> f b -> f c -> f d -> f e -> f g -> f h -> f (a, b, c, d, e, g, h)
161    zip7 = zipWith7 (\t u v w x y z -> (t,u,v,w,x,y,z))
162
163    zap7 :: f (a -> b -> c -> d -> e -> g -> h) -> f a -> f b -> f c -> f d -> f e -> f g -> f h
164    zap7 = zipWith7 id
165
166    unzip7 :: f (a, b, c, d, e, g, h) -> (f a, f b, f c, f d, f e, f g, f h)
167    -- unzip3 = fmap (\(x,_,_)->x) &&& fmap (\(_,x,_)->x) &&& fmap (\(_,_,x)->x)
168
169instance Zip7 [] where
170    zip7 = List.zip7
171    unzip7 = List.unzip7
172    zipWith7 = List.zipWith7
173