-- |
-- 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 #-}
{-# LANGUAGE StandaloneDeriving #-}
{-# LANGUAGE GeneralizedNewtypeDeriving #-}
module Foundation.Collection.Zippable
( BoxedZippable(..)
, Zippable(..)
) where
import qualified Basement.UArray as UV
import qualified Basement.BoxedArray as BA
import qualified Basement.String as S
import Foundation.Collection.Element
import Foundation.Collection.Sequential
import Basement.Compat.Base
import Basement.Types.AsciiString(AsciiString(..))
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'
deriving instance Zippable AsciiString
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