-- | -- Module : Foundation.Collection.Zippable -- License : BSD-style -- Maintainer : foundation -- Stability : experimental -- Portability : portable -- -- Common functions (e. g. 'zip', 'zipWith') that are useful for combining -- multiple collections. -- {-# LANGUAGE FlexibleInstances #-} {-# LANGUAGE MultiParamTypeClasses #-} module Foundation.Collection.Zippable ( BoxedZippable(..) , Zippable(..) ) where import qualified Foundation.Array.Unboxed as UV import qualified Foundation.Array.Boxed as BA import qualified Foundation.String.UTF8 as S import Foundation.Collection.Element import Foundation.Collection.Sequential import Foundation.Internal.Base import qualified Prelude import GHC.ST class Sequential col => Zippable col where -- | 'zipWith' generalises 'zip' by zipping with the function given as the -- first argument, instead of a tupling function. For example, @'zipWith' (+)@ -- is applied to two collections to produce the collection of corresponding -- sums. zipWith :: (Sequential a, Sequential b) => (Element a -> Element b -> Element col) -> a -> b -> col zipWith f a b = go f (toList a, toList b) where go f' = maybe mempty (\(x, xs) -> uncurry2 f' x `cons` go f' xs) . uncons2 -- | Like 'zipWith', but works with 3 collections. zipWith3 :: (Sequential a, Sequential b, Sequential c) => (Element a -> Element b -> Element c -> Element col) -> a -> b -> c -> col zipWith3 f a b c = go f (toList a, toList b, toList c) where go f' = maybe mempty (\(x, xs) -> uncurry3 f' x `cons` go f' xs) . uncons3 -- | Like 'zipWith', but works with 4 collections. zipWith4 :: (Sequential a, Sequential b, Sequential c, Sequential d) => (Element a -> Element b -> Element c -> Element d -> Element col) -> a -> b -> c -> d -> col zipWith4 fn a b c d = go fn (toList a, toList b, toList c, toList d) where go f' = maybe mempty (\(x, xs) -> uncurry4 f' x `cons` go f' xs) . uncons4 -- | Like 'zipWith', but works with 5 collections. zipWith5 :: (Sequential a, Sequential b, Sequential c, Sequential d, Sequential e) => (Element a -> Element b -> Element c -> Element d -> Element e -> Element col) -> a -> b -> c -> d -> e -> col zipWith5 fn a b c d e = go fn (toList a, toList b, toList c, toList d, toList e) where go f' = maybe mempty (\(x, xs) -> uncurry5 f' x `cons` go f' xs) . uncons5 -- | Like 'zipWith', but works with 6 collections. zipWith6 :: ( Sequential a, Sequential b, Sequential c, Sequential d, Sequential e , Sequential f) => (Element a -> Element b -> Element c -> Element d -> Element e -> Element f -> Element col) -> a -> b -> c -> d -> e -> f -> col zipWith6 fn a b c d e f = go fn (toList a, toList b, toList c, toList d, toList e, toList f) where go f' = maybe mempty (\(x, xs) -> uncurry6 f' x `cons` go f' xs) . uncons6 -- | Like 'zipWith', but works with 7 collections. zipWith7 :: ( Sequential a, Sequential b, Sequential c, Sequential d, Sequential e , Sequential f, Sequential g ) => (Element a -> Element b -> Element c -> Element d -> Element e -> Element f -> Element g -> Element col) -> a -> b -> c -> d -> e -> f -> g -> col zipWith7 fn a b c d e f g = go fn (toList a, toList b, toList c, toList d, toList e, toList f, toList g) where go f' = maybe mempty (\(x, xs) -> uncurry7 f' x `cons` go f' xs) . uncons7 instance Zippable [c] instance UV.PrimType ty => Zippable (UV.UArray ty) where zipWith f as bs = runST $ UV.builderBuild 64 $ go f (toList as) (toList bs) where go _ [] _ = return () go _ _ [] = return () go f' (a':as') (b':bs') = UV.builderAppend (f' a' b') >> go f' as' bs' instance Zippable (BA.Array ty) where zipWith f as bs = runST $ BA.builderBuild 64 $ go f (toList as) (toList bs) where go _ [] _ = return () go _ _ [] = return () go f' (a':as') (b':bs') = BA.builderAppend (f' a' b') >> go f' as' bs' instance Zippable S.String where zipWith f as bs = runST $ S.builderBuild 64 $ go f (toList as) (toList bs) where go _ [] _ = return () go _ _ [] = return () go f' (a':as') (b':bs') = S.builderAppend (f' a' b') >> go f' as' bs' class Zippable col => BoxedZippable col where -- | 'zip' takes two collections and returns a collections of corresponding -- pairs. If one input collection is short, excess elements of the longer -- collection are discarded. zip :: ( Sequential a, Sequential b , Element col ~ (Element a, Element b) ) => a -> b -> col zip = zipWith (,) -- | Like 'zip', but works with 3 collections. zip3 :: ( Sequential a, Sequential b, Sequential c , Element col ~ (Element a, Element b, Element c) ) => a -> b -> c -> col zip3 = zipWith3 (,,) -- | Like 'zip', but works with 4 collections. zip4 :: ( Sequential a, Sequential b, Sequential c, Sequential d , Element col ~ (Element a, Element b, Element c, Element d) ) => a -> b -> c -> d -> col zip4 = zipWith4 (,,,) -- | Like 'zip', but works with 5 collections. zip5 :: ( Sequential a, Sequential b, Sequential c, Sequential d, Sequential e , Element col ~ (Element a, Element b, Element c, Element d, Element e) ) => a -> b -> c -> d -> e -> col zip5 = zipWith5 (,,,,) -- | Like 'zip', but works with 6 collections. zip6 :: ( Sequential a, Sequential b, Sequential c, Sequential d, Sequential e, Sequential f , Element col ~ (Element a, Element b, Element c, Element d, Element e, Element f) ) => a -> b -> c -> d -> e -> f -> col zip6 = zipWith6 (,,,,,) -- | Like 'zip', but works with 7 collections. zip7 :: ( Sequential a, Sequential b, Sequential c, Sequential d, Sequential e, Sequential f, Sequential g , Element col ~ (Element a, Element b, Element c, Element d, Element e, Element f, Element g) ) => a -> b -> c -> d -> e -> f -> g -> col zip7 = zipWith7 (,,,,,,) -- | 'unzip' transforms a collection of pairs into a collection of first -- components and a collection of second components. unzip :: (Sequential a, Sequential b, Element col ~ (Element a, Element b)) => col -> (a, b) unzip = go . toList where go [] = (mempty, mempty) go ((a, b):xs) = let (as, bs) = go xs in (a `cons` as, b `cons` bs) -- | Like 'unzip', but works on a collection of 3-element tuples. unzip3 :: ( Sequential a, Sequential b, Sequential c , Element col ~ (Element a, Element b, Element c) ) => col -> (a, b, c) unzip3 = go . toList where go [] = (mempty, mempty, mempty) go ((a, b, c):xs) = let (as, bs, cs) = go xs in (a `cons` as, b `cons` bs, c `cons` cs) -- | Like 'unzip', but works on a collection of 4-element tuples. unzip4 :: ( Sequential a, Sequential b, Sequential c, Sequential d , Element col ~ (Element a, Element b, Element c, Element d) ) => col -> (a, b, c, d) unzip4 = go . toList where go [] = (mempty, mempty, mempty, mempty) go ((a, b, c, d):xs) = let (as, bs, cs, ds) = go xs in (a `cons` as, b `cons` bs, c `cons` cs, d `cons` ds) -- | Like 'unzip', but works on a collection of 5-element tuples. unzip5 :: ( Sequential a, Sequential b, Sequential c, Sequential d, Sequential e , Element col ~ (Element a, Element b, Element c, Element d, Element e) ) => col -> (a, b, c, d, e) unzip5 = go . toList where go [] = (mempty, mempty, mempty, mempty, mempty) go ((a, b, c, d, e):xs) = let (as, bs, cs, ds, es) = go xs in (a `cons` as, b `cons` bs, c `cons` cs, d `cons` ds, e `cons` es) -- | Like 'unzip', but works on a collection of 6-element tuples. unzip6 :: ( Sequential a, Sequential b, Sequential c, Sequential d, Sequential e, Sequential f , Element col ~ (Element a, Element b, Element c, Element d, Element e, Element f) ) => col -> (a, b, c, d, e, f) unzip6 = go . toList where go [] = (mempty, mempty, mempty, mempty, mempty, mempty) go ((a, b, c, d, e, f):xs) = let (as, bs, cs, ds, es, fs) = go xs in (a `cons` as, b `cons` bs, c `cons` cs, d `cons` ds, e `cons` es, f `cons` fs) -- | Like 'unzip', but works on a collection of 7-element tuples. unzip7 :: ( Sequential a, Sequential b, Sequential c, Sequential d, Sequential e, Sequential f, Sequential g , Element col ~ (Element a, Element b, Element c, Element d, Element e, Element f, Element g) ) => col -> (a, b, c, d, e, f, g) unzip7 = go . toList where go [] = (mempty, mempty, mempty, mempty, mempty, mempty, mempty) go ((a, b, c, d, e, f, g):xs) = let (as, bs, cs, ds, es, fs, gs) = go xs in (a `cons` as, b `cons` bs, c `cons` cs, d `cons` ds, e `cons` es, f `cons` fs, g `cons` gs) instance BoxedZippable [a] instance BoxedZippable (BA.Array ty) -- * Tuple helper functions uncons2 :: (Sequential a, Sequential b) => (a, b) -> Maybe ((Element a, Element b), (a, b)) uncons2 xs = let (as, bs) = xs in do (a', as') <- uncons as (b', bs') <- uncons bs return ((a', b'), (as', bs')) uncons3 :: (Sequential a, Sequential b, Sequential c) => (a, b, c) -> Maybe ((Element a, Element b, Element c), (a, b, c)) uncons3 xs = let (as, bs, cs) = xs in do (a', as') <- uncons as (b', bs') <- uncons bs (c', cs') <- uncons cs return ((a', b', c'), (as', bs', cs')) uncons4 :: (Sequential a, Sequential b, Sequential c, Sequential d) => (a, b, c, d) -> Maybe ( (Element a, Element b, Element c, Element d) , (a, b, c, d) ) uncons4 xs = let (as, bs, cs, ds) = xs in do (a', as') <- uncons as (b', bs') <- uncons bs (c', cs') <- uncons cs (d', ds') <- uncons ds return ((a', b', c', d'), (as', bs', cs', ds')) uncons5 :: (Sequential a, Sequential b, Sequential c, Sequential d, Sequential e) => (a, b, c, d, e) -> Maybe ( (Element a, Element b, Element c, Element d, Element e) , (a, b, c, d, e) ) uncons5 xs = let (as, bs, cs, ds, es) = xs in do (a', as') <- uncons as (b', bs') <- uncons bs (c', cs') <- uncons cs (d', ds') <- uncons ds (e', es') <- uncons es return ((a', b', c', d', e'), (as', bs', cs', ds', es')) uncons6 :: ( Sequential a, Sequential b, Sequential c, Sequential d, Sequential e , Sequential f ) => (a, b, c, d, e, f) -> Maybe ( (Element a, Element b, Element c, Element d, Element e, Element f) , (a, b, c, d, e, f) ) uncons6 xs = let (as, bs, cs, ds, es, fs) = xs in do (a', as') <- uncons as (b', bs') <- uncons bs (c', cs') <- uncons cs (d', ds') <- uncons ds (e', es') <- uncons es (f', fs') <- uncons fs return ((a', b', c', d', e', f'), (as', bs', cs', ds', es', fs')) uncons7 :: ( Sequential a, Sequential b, Sequential c, Sequential d, Sequential e , Sequential f, Sequential g ) => (a, b, c, d, e, f, g) -> Maybe ( (Element a, Element b, Element c, Element d, Element e, Element f , Element g) , (a, b, c, d, e, f, g) ) uncons7 xs = let (as, bs, cs, ds, es, fs, gs) = xs in do (a', as') <- uncons as (b', bs') <- uncons bs (c', cs') <- uncons cs (d', ds') <- uncons ds (e', es') <- uncons es (f', fs') <- uncons fs (g', gs') <- uncons gs return ( (a', b', c', d', e', f', g') , (as', bs', cs', ds', es', fs', gs') ) uncurry2 :: (a -> b -> c) -> (a, b) -> c uncurry2 = Prelude.uncurry uncurry3 :: (a -> b -> c -> d) -> (a, b, c) -> d uncurry3 fn (a, b, c) = fn a b c uncurry4 :: (a -> b -> c -> d -> g) -> (a, b, c, d) -> g uncurry4 fn (a, b, c, d) = fn a b c d uncurry5 :: (a -> b -> c -> d -> e -> f) -> (a, b, c, d, e) -> f uncurry5 fn (a, b, c, d, e) = fn a b c d e uncurry6 :: (a -> b -> c -> d -> e -> f -> g) -> (a, b, c, d, e, f) -> g uncurry6 fn (a, b, c, d, e, f) = fn a b c d e f uncurry7 :: (a -> b -> c -> d -> e -> f -> g -> h) -> (a, b, c, d, e, f, g) -> h uncurry7 fn (a, b, c, d, e, f, g) = fn a b c d e f g