{-# LANGUAGE DerivingStrategies #-}
module Data.Traversable.Unzip.LazyPair (LazyPair (..)) where
import Data.Biapplicative (Bifunctor (..), Biapplicative (..))

-- | A pair with a lazier 'Functor' instance (we don't actually use), and with
-- 'Bifunctor' and 'Biapplicative' instances designed not to leak memory.
newtype LazyPair a b = LazyPair { forall a b. LazyPair a b -> (a, b)
unLazyPair :: (a, b) }

-- Getting the thunks we want when we want them is quite fragile. I found it
-- helpful to inspect the Core for unzipping Maps; those have enough strictness
-- to make GHC want to do bad things that will leak memory. It may be possible
-- to improve it, but this is the most robust approach I've found so far.

instance Functor (LazyPair a) where
  fmap :: forall a b. (a -> b) -> LazyPair a a -> LazyPair a b
fmap a -> b
f (LazyPair (a, a)
ab) = a -> a -> LazyPair a b
forall {a}. a -> a -> LazyPair a b
combine a
a a
b
    where
      ~(a
a, a
b) = (a, a)
ab
      combine :: a -> a -> LazyPair a b
combine a
x a
y = (a, b) -> LazyPair a b
forall a b. (a, b) -> LazyPair a b
LazyPair (a
x, a -> b
f a
y)
      {-# NOINLINE combine #-}

instance Bifunctor LazyPair where
  bimap :: forall a b c d.
(a -> b) -> (c -> d) -> LazyPair a c -> LazyPair b d
bimap a -> b
f c -> d
g (LazyPair (a, c)
ab) = a -> c -> LazyPair b d
combine a
a c
b
    where
      ~(a
a, c
b) = (a, c)
ab
      combine :: a -> c -> LazyPair b d
combine a
p c
q = (b, d) -> LazyPair b d
forall a b. (a, b) -> LazyPair a b
LazyPair (a -> b
f a
p, c -> d
g c
q)
      {-# NOINLINE combine #-}

instance Biapplicative LazyPair where
  bipure :: forall a b. a -> b -> LazyPair a b
bipure a
x b
y = (a, b) -> LazyPair a b
forall a b. (a, b) -> LazyPair a b
LazyPair (a
x, b
y)
  biliftA2 :: forall a b c d e f.
(a -> b -> c)
-> (d -> e -> f) -> LazyPair a d -> LazyPair b e -> LazyPair c f
biliftA2 a -> b -> c
f d -> e -> f
g (LazyPair (a, d)
x1y1) (LazyPair (b, e)
x2y2) = a -> b -> d -> e -> LazyPair c f
combine a
x1 b
x2 d
y1 e
y2
    where
      ~(a
x1, d
y1) = (a, d)
x1y1
      ~(b
x2, e
y2) = (b, e)
x2y2
      -- I worked out this "combine" trick for Data.List.transpose.
      -- By marking the combine function NOINLINE, we ensure that
      -- all four selector thunks are constructed up front. In particular,
      -- we don't let GHC do something like
      --
      -- (let {~(x1, _) = x1y1; ~(x2, _) = x2y2} in f x1 x2,
      --  let {~(_, y1) = x1y1; ~(_, y2) = x2y2} in g y1 y2)
      combine :: a -> b -> d -> e -> LazyPair c f
combine a
p b
q d
r e
s = (c, f) -> LazyPair c f
forall a b. (a, b) -> LazyPair a b
LazyPair (a -> b -> c
f a
p b
q, d -> e -> f
g d
r e
s)
      {-# NOINLINE combine #-}