{-# OPTIONS_GHC -Wall -fwarn-tabs #-}
----------------------------------------------------------------
--                                                  ~ 2021.10.17
-- |
-- Module      :  Data.List.Extras.Pair
-- Copyright   :  Copyright (c) 2007--2021 wren gayle romano
-- License     :  BSD3
-- Maintainer  :  wren@cpan.org
-- Stability   :  stable
-- Portability :  Haskell98
--
-- This module provides safe zipping functions which will fail
-- (return 'Nothing') on uneven length lists.
----------------------------------------------------------------

module Data.List.Extras.Pair
    (
    -- * Safe functions for zipping lists
      pairWithBy, pairWith, pairBy, pair

    -- * Special safe zipping functions
    , biject, biject'

    -- * New (unsafe) zipping functions
    , zipWithBy, zipBy
    ) where


----------------------------------------------------------------
----------------------------------------------------------------
-- TODO: benchmark fusion performance of:
--
--     foldr cons nil .: zipWith (,)
--     zipWithBy (,) cons nil
--
-- ...That is, the latter is a manual fusion of the former, but
-- does zip/zipWith have a special ability to fuse with the incoming
-- lists? Or can foldr fuse with consumers in ways zipWithBy can't?

-- | An unsafe variant of 'pairWithBy' to fill out the interface.
zipWithBy :: (a -> b -> c)       -- tuple homomorphism
          -> (c -> d -> d) -> d  -- list  homomorphism
          -> [a] -> [b] -> d     -- a @zip@ function
{-# INLINE zipWithBy #-}
-- We use the explicit lambda in order to improve inlining in ghc-7.
zipWithBy :: (a -> b -> c) -> (c -> d -> d) -> d -> [a] -> [b] -> d
zipWithBy a -> b -> c
k c -> d -> d
f d
z = \[a]
xs [b]
ys -> [a] -> [b] -> (d -> d) -> d
forall p. [a] -> [b] -> (d -> p) -> p
zipWB [a]
xs [b]
ys d -> d
forall a. a -> a
id
    where
    zipWB :: [a] -> [b] -> (d -> p) -> p
zipWB (a
x:[a]
xs) (b
y:[b]
ys) d -> p
cc = [a] -> [b] -> (d -> p) -> p
zipWB [a]
xs [b]
ys (d -> p
cc (d -> p) -> (d -> d) -> d -> p
forall b c a. (b -> c) -> (a -> b) -> a -> c
. c -> d -> d
f (a -> b -> c
k a
x b
y))
    zipWB [a]
_      [b]
_      d -> p
cc = d -> p
cc d
z


-- | A version of 'zip' that uses a user-defined list homomorphism.
zipBy :: ((a,b) -> d -> d) -> d -> [a] -> [b] -> d
{-# INLINE zipBy #-}
zipBy :: ((a, b) -> d -> d) -> d -> [a] -> [b] -> d
zipBy = (a -> b -> (a, b)) -> ((a, b) -> d -> d) -> d -> [a] -> [b] -> d
forall a b c d.
(a -> b -> c) -> (c -> d -> d) -> d -> [a] -> [b] -> d
zipWithBy (,)


----------------------------------------------------------------
----------------------------------------------------------------
-- | A generic version of 'pair'. The first argument is a tuple
-- homomorphism (i.e. a function for how to combine values from the
-- two lists), the second two arguments form a list homomorphism
-- (i.e. so you can 'foldr' the @[c]@ list directly without actually
-- constructing it).
--
-- In order to evaluate to WHNF 'pairWithBy' is strict in both list
-- arguments, as it must be, to determine that the lists are of the
-- same length. This means it can survive one infinite list (yielding
-- 'Nothing') but that it can't survive two. The implementation is
-- very efficient and uses a tight tail-recursive loop, however
-- with extremely long lists it will be churning through heap and
-- that tightness can make it hard to interrupt (lists of 1 million
-- elements return in 1~2 seconds, but lists of 10 million can lock
-- your system up).

pairWithBy :: (a -> b -> c)          -- @(,)@ tuple homomorphism
           -> (c -> d -> d)          -- @(:)@ list  homomorphism, pt. 1
           -> d                      -- @[]@  list  homomorphism, pt. 2
           -> [a] -> [b] -> Maybe d  -- a safer @zip@ function
{-# INLINE pairWithBy #-}
-- We use the explicit lambda in order to improve inlining in ghc-7.
pairWithBy :: (a -> b -> c) -> (c -> d -> d) -> d -> [a] -> [b] -> Maybe d
pairWithBy a -> b -> c
k c -> d -> d
f d
z = \[a]
xs [b]
ys -> [a] -> [b] -> (d -> d) -> Maybe d
forall a. [a] -> [b] -> (d -> a) -> Maybe a
pairWB [a]
xs [b]
ys d -> d
forall a. a -> a
id
    where
    -- N.B. Strict accumulators are usually awesome, but don't
    -- even consider it when doing CPS! Making @cc@ strict degrades
    -- performance significantly; it takes twice as long and twice
    -- as much heap just to get to WHNF. After evaluating the spine
    -- of the resulting list from 'pair' that drops to +10% time
    -- and +25% heap, which is still much worse.

    pairWB :: [a] -> [b] -> (d -> a) -> Maybe a
pairWB (a
x:[a]
xs) (b
y:[b]
ys) d -> a
cc = [a] -> [b] -> (d -> a) -> Maybe a
pairWB [a]
xs [b]
ys (d -> a
cc (d -> a) -> (d -> d) -> d -> a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. c -> d -> d
f (a -> b -> c
k a
x b
y))
    pairWB []     []     d -> a
cc = a -> Maybe a
forall a. a -> Maybe a
Just (d -> a
cc d
z)
    pairWB [a]
_      [b]
_      d -> a
_  = Maybe a
forall a. Maybe a
Nothing

-- TODO: we could make this more general still by fusing @f@ and @k@, which we'd often want to do anyways if we're using this full form.

----------------------------------------------------------------

-- | A safe version of 'zipWith'.
pairWith :: (a -> b -> c) -> [a] -> [b] -> Maybe [c]
{-# INLINE pairWith #-}
pairWith :: (a -> b -> c) -> [a] -> [b] -> Maybe [c]
pairWith a -> b -> c
f = (a -> b -> c)
-> (c -> [c] -> [c]) -> [c] -> [a] -> [b] -> Maybe [c]
forall a b c d.
(a -> b -> c) -> (c -> d -> d) -> d -> [a] -> [b] -> Maybe d
pairWithBy a -> b -> c
f (:) []


-- | A safe version of 'zip' that uses a user-defined list homomorphism.
pairBy :: ((a,b) -> d -> d) -> d -> [a] -> [b] -> Maybe d
{-# INLINE pairBy #-}
pairBy :: ((a, b) -> d -> d) -> d -> [a] -> [b] -> Maybe d
pairBy = (a -> b -> (a, b))
-> ((a, b) -> d -> d) -> d -> [a] -> [b] -> Maybe d
forall a b c d.
(a -> b -> c) -> (c -> d -> d) -> d -> [a] -> [b] -> Maybe d
pairWithBy (,)


-- | A safe version of 'zip'.
pair :: [a] -> [b] -> Maybe [(a,b)]
{-# INLINE pair #-}
pair :: [a] -> [b] -> Maybe [(a, b)]
pair = (a -> b -> (a, b))
-> ((a, b) -> [(a, b)] -> [(a, b)])
-> [(a, b)]
-> [a]
-> [b]
-> Maybe [(a, b)]
forall a b c d.
(a -> b -> c) -> (c -> d -> d) -> d -> [a] -> [b] -> Maybe d
pairWithBy (,) (:) []


----------------------------------------------------------------
-- These two are just here because they're often requested, and
-- besides they're kinda cute.

-- | A bijection from a list of functions and a list of arguments
-- to a list of results of applying the functions bijectively.
biject :: [a -> b] -> [a] -> Maybe [b]
{-# INLINE biject #-}
biject :: [a -> b] -> [a] -> Maybe [b]
biject = ((a -> b) -> a -> b) -> [a -> b] -> [a] -> Maybe [b]
forall a b c. (a -> b -> c) -> [a] -> [b] -> Maybe [c]
pairWith (a -> b) -> a -> b
forall a b. (a -> b) -> a -> b
($) -- 'id' also works


-- | A version of 'biject' that applies functions strictly. N.B.
-- the list is still lazily evaluated, this just makes the functions
-- strict in their argument.
biject' :: [a -> b] -> [a] -> Maybe [b]
{-# INLINE biject' #-}
biject' :: [a -> b] -> [a] -> Maybe [b]
biject' = ((a -> b) -> a -> b) -> [a -> b] -> [a] -> Maybe [b]
forall a b c. (a -> b -> c) -> [a] -> [b] -> Maybe [c]
pairWith (a -> b) -> a -> b
forall a b. (a -> b) -> a -> b
($!)

----------------------------------------------------------------
----------------------------------------------------------- fin.