{-# LANGUAGE DerivingVia #-}
{-# LANGUAGE ExistentialQuantification #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE GADTs #-}
{-# LANGUAGE RankNTypes #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE StandaloneDeriving #-}

-- | 'Applicative' functor transformers, like monad transformers, for free.
module Control.Applicative.Trans.FreeAp (
  ApT (..),
  toFree,
  fromFree,
  transApT,
  hoistApT,
  liftF,
  liftT,
  appendApT,
  foldApT,
  foldApT_,
  fjoinApTLeft,
  fjoinApTRight,
  ApIx (..),
  fromIx, indices,
  reconstruct
) where

import Control.Applicative
import qualified Control.Applicative.Free as Free
import Data.Foldable (toList)
import Data.Functor.Identity
import Data.Traversable (mapAccumL)

import qualified GHC.Arr as Arr
import GHC.Stack (HasCallStack)

import Data.Functor.Classes

{- | @'ApT' f@ is a \"free\" \"applicative transformer\", in the same sense
  @'Control.Monad.Trans.Free.FreeT' f@ is a free monad transformer.

  ==== \"Applicative transformer\"

  Being an \"applicative transformer\" means these two things:

  * Applying @ApT f@ to an applicative functor @g@ constructs a new applicative
    functor @ApT f g@.

  * Using 'liftT', you can lift an action of @g@ to the action of @ApT f g@.

      > liftT :: g x -> ApT f g x

      'liftT' is an applicative transformation. In other words, @liftT@ preserves
      'pure' and @'<*>'@:

      > liftT (pure x) = pure x
      > liftT (x <*> y) = liftT x <*> liftT y

  ==== \"Free\" applicative transformer

  It's the \"free\" applicative transformer. It means @ApT f g@ is the special, the most universal
  one among various applicative functors which can lift @f@ and @g@ into them.

  * @ApT f g@ has a way to lift any value of @f a@ into an action of @ApT f g a@.

      > liftF :: (Applicative g) => f a -> ApT f g a

      Because @ApT f g@ is also an applicative transformer on @g@, it has a way to lift @g@ too.

      > liftT :: g x -> ApT f g x

  * Suppose another applicative functor @h@ is capable of lifting both @f@ and @g@ to @h@.

      > fh :: f a -> h a
      > gh :: g a -> h a

      @ApT f g@ is the universal applicative among them. There's 'foldApT' to construct
      the applicative transformation from @ApT f g@ to @h@, without losing how to lift @f@ and @g@.

      > foldApT :: forall f g h x. Applicative h => (forall a. f a -> h a) -> (forall a. g a -> h a) -> ApT f g x -> h x
      >
      > foldApT fh gh :: forall x. ApT f g x -> h x
      >
      > foldApT fh gh . liftF = fh
      > foldApT fh gh . liftT = gh

  * @ApT f g@ contains no extra data that are not from lifting @f@ and/or @g@ then combining them together
    by @Applicative@ operation '<*>'.

      It means any applicative transformation @run :: forall a. ApT f g a -> h a@ which satisfies @run . liftF = fh@ and @run . liftT = gh@
      is equivalent to @foldApT fh gh@.
-}
data ApT f g x
  = PureT (g x)
  | forall a b c. ApT (a -> b -> c -> x) (g a) (f b) (ApT f g c)

instance Functor g => Functor (ApT f g) where
  fmap :: forall a b. (a -> b) -> ApT f g a -> ApT f g b
fmap a -> b
h (PureT g a
gx) = forall (f :: * -> *) (g :: * -> *) x. g x -> ApT f g x
PureT forall a b. (a -> b) -> a -> b
$ forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap a -> b
h g a
gx
  fmap a -> b
h (ApT a -> b -> c -> a
x g a
ga f b
fb ApT f g c
rc) = forall (f :: * -> *) (g :: * -> *) x a b c.
(a -> b -> c -> x) -> g a -> f b -> ApT f g c -> ApT f g x
ApT (\a
a b
b c
c -> a -> b
h (a -> b -> c -> a
x a
a b
b c
c)) g a
ga f b
fb ApT f g c
rc

  a
x <$ :: forall a b. a -> ApT f g b -> ApT f g a
<$ PureT g b
gx = forall (f :: * -> *) (g :: * -> *) x. g x -> ApT f g x
PureT (a
x forall (f :: * -> *) a b. Functor f => a -> f b -> f a
<$ g b
gx)
  a
x <$ ApT a -> b -> c -> b
_ g a
ga f b
fb ApT f g c
rc = forall (f :: * -> *) (g :: * -> *) x a b c.
(a -> b -> c -> x) -> g a -> f b -> ApT f g c -> ApT f g x
ApT (\a
_ b
_ c
_ -> a
x) g a
ga f b
fb ApT f g c
rc

instance Applicative g => Applicative (ApT f g) where
  pure :: forall a. a -> ApT f g a
pure = forall (f :: * -> *) (g :: * -> *) x. g x -> ApT f g x
PureT forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (f :: * -> *) a. Applicative f => a -> f a
pure
  PureT g (a -> b)
gx <*> :: forall a b. ApT f g (a -> b) -> ApT f g a -> ApT f g b
<*> PureT g a
gy = forall (f :: * -> *) (g :: * -> *) x. g x -> ApT f g x
PureT (g (a -> b)
gx forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> g a
gy)
  PureT g (a -> b)
gx <*> ApT a -> b -> c -> a
y g a
ga f b
fb ApT f g c
rc = forall (f :: * -> *) (g :: * -> *) x a b c.
(a -> b -> c -> x) -> g a -> f b -> ApT f g c -> ApT f g x
ApT (\ ~(a -> b
x, a
a) b
b c
c -> a -> b
x (a -> b -> c -> a
y a
a b
b c
c)) (forall (f :: * -> *) a b c.
Applicative f =>
(a -> b -> c) -> f a -> f b -> f c
liftA2 (,) g (a -> b)
gx g a
ga) f b
fb ApT f g c
rc
  ApT a -> b -> c -> a -> b
x g a
ga f b
fb ApT f g c
rc <*> ApT f g a
rest = forall (f :: * -> *) (g :: * -> *) x a b c.
(a -> b -> c -> x) -> g a -> f b -> ApT f g c -> ApT f g x
ApT (\a
a b
b ~(c
c, a
y) -> a -> b -> c -> a -> b
x a
a b
b c
c a
y) g a
ga f b
fb (forall (f :: * -> *) a b c.
Applicative f =>
(a -> b -> c) -> f a -> f b -> f c
liftA2 (,) ApT f g c
rc ApT f g a
rest)

  PureT g a
gx *> :: forall a b. ApT f g a -> ApT f g b -> ApT f g b
*> PureT g b
gy = forall (f :: * -> *) (g :: * -> *) x. g x -> ApT f g x
PureT (g a
gx forall (f :: * -> *) a b. Applicative f => f a -> f b -> f b
*> g b
gy)
  PureT g a
gx *> ApT a -> b -> c -> b
y g a
ga f b
fb ApT f g c
rc = forall (f :: * -> *) (g :: * -> *) x a b c.
(a -> b -> c -> x) -> g a -> f b -> ApT f g c -> ApT f g x
ApT a -> b -> c -> b
y (g a
gx forall (f :: * -> *) a b. Applicative f => f a -> f b -> f b
*> g a
ga) f b
fb ApT f g c
rc
  ApT a -> b -> c -> a
_ g a
ga f b
fb ApT f g c
rc *> ApT f g b
rest = forall (f :: * -> *) (g :: * -> *) x a b c.
(a -> b -> c -> x) -> g a -> f b -> ApT f g c -> ApT f g x
ApT (\a
_ b
_ b
y -> b
y) g a
ga f b
fb (ApT f g c
rc forall (f :: * -> *) a b. Applicative f => f a -> f b -> f b
*> ApT f g b
rest)

  PureT g a
gx <* :: forall a b. ApT f g a -> ApT f g b -> ApT f g a
<* PureT g b
gy = forall (f :: * -> *) (g :: * -> *) x. g x -> ApT f g x
PureT (g a
gx forall (f :: * -> *) a b. Applicative f => f a -> f b -> f a
<* g b
gy)
  PureT g a
gx <* ApT a -> b -> c -> b
_ g a
ga f b
fb ApT f g c
rc = forall (f :: * -> *) (g :: * -> *) x a b c.
(a -> b -> c -> x) -> g a -> f b -> ApT f g c -> ApT f g x
ApT (\a
x b
_ c
_ -> a
x) (g a
gx forall (f :: * -> *) a b. Applicative f => f a -> f b -> f a
<* g a
ga) f b
fb ApT f g c
rc
  ApT a -> b -> c -> a
x g a
ga f b
fb ApT f g c
rc <* ApT f g b
rest = forall (f :: * -> *) (g :: * -> *) x a b c.
(a -> b -> c -> x) -> g a -> f b -> ApT f g c -> ApT f g x
ApT a -> b -> c -> a
x g a
ga f b
fb (ApT f g c
rc forall (f :: * -> *) a b. Applicative f => f a -> f b -> f a
<* ApT f g b
rest)

-- | When the base applicative is 'Identity', @ApT f Identity@ is the free applicative 'Free.Ap'.
toFree :: ApT f Identity a -> Free.Ap f a
toFree :: forall (f :: * -> *) a. ApT f Identity a -> Ap f a
toFree = forall a b (f :: * -> *). (a -> b) -> ApT f Identity a -> Ap f b
toFreeAux forall a. a -> a
id

toFreeAux :: (a -> b) -> ApT f Identity a -> Free.Ap f b
toFreeAux :: forall a b (f :: * -> *). (a -> b) -> ApT f Identity a -> Ap f b
toFreeAux a -> b
k (PureT (Identity a
a)) = forall a (f :: * -> *). a -> Ap f a
Free.Pure (a -> b
k a
a)
toFreeAux a -> b
k (ApT a -> b -> c -> a
x (Identity a
a) f b
fb ApT f Identity c
rc) = forall (f :: * -> *) a1 a. f a1 -> Ap f (a1 -> a) -> Ap f a
Free.Ap f b
fb (forall a b (f :: * -> *). (a -> b) -> ApT f Identity a -> Ap f b
toFreeAux (\c
c b
b -> a -> b
k (a -> b -> c -> a
x a
a b
b c
c)) ApT f Identity c
rc)

-- | Inverse of @toFree@.
fromFree :: Free.Ap f a -> ApT f Identity a
fromFree :: forall (f :: * -> *) a. Ap f a -> ApT f Identity a
fromFree (Free.Pure a
a) = forall (f :: * -> *) (g :: * -> *) x. g x -> ApT f g x
PureT (forall a. a -> Identity a
Identity a
a)
fromFree (Free.Ap f a1
fb Ap f (a1 -> a)
rest) = forall (f :: * -> *) (g :: * -> *) x a b c.
(a -> b -> c -> x) -> g a -> f b -> ApT f g c -> ApT f g x
ApT forall a b c. (a -> b -> c) -> b -> a -> c
flip (forall a. a -> Identity a
Identity forall a. a -> a
id) f a1
fb (forall (f :: * -> *) a. Ap f a -> ApT f Identity a
fromFree Ap f (a1 -> a)
rest)

{- | Lift an applicative transformation @(forall a. g a -> g' a)@ to
  an applicative transformation @(forall b. ApT f g b -> ApT f g' b)@.
-}
hoistApT :: (forall a. g a -> g' a) -> ApT f g b -> ApT f g' b
hoistApT :: forall (g :: * -> *) (g' :: * -> *) (f :: * -> *) b.
(forall a. g a -> g' a) -> ApT f g b -> ApT f g' b
hoistApT forall a. g a -> g' a
phi (PureT g b
gx) = forall (f :: * -> *) (g :: * -> *) x. g x -> ApT f g x
PureT (forall a. g a -> g' a
phi g b
gx)
hoistApT forall a. g a -> g' a
phi (ApT a -> b -> c -> b
x g a
ga f b
fb ApT f g c
rc) = forall (f :: * -> *) (g :: * -> *) x a b c.
(a -> b -> c -> x) -> g a -> f b -> ApT f g c -> ApT f g x
ApT a -> b -> c -> b
x (forall a. g a -> g' a
phi g a
ga) f b
fb (forall (g :: * -> *) (g' :: * -> *) (f :: * -> *) b.
(forall a. g a -> g' a) -> ApT f g b -> ApT f g' b
hoistApT forall a. g a -> g' a
phi ApT f g c
rc)

{- | Lift any natural transformation @(forall a. f a -> f' a)@ to
  an applicative transformation @(forall b. ApT f g b -> ApT f' g b)@.
-}
transApT :: (forall a. f a -> f' a) -> ApT f g b -> ApT f' g b
transApT :: forall (f :: * -> *) (f' :: * -> *) (g :: * -> *) b.
(forall a. f a -> f' a) -> ApT f g b -> ApT f' g b
transApT forall a. f a -> f' a
_ (PureT g b
gx) = forall (f :: * -> *) (g :: * -> *) x. g x -> ApT f g x
PureT g b
gx
transApT forall a. f a -> f' a
phi (ApT a -> b -> c -> b
x g a
ga f b
fb ApT f g c
rc) = forall (f :: * -> *) (g :: * -> *) x a b c.
(a -> b -> c -> x) -> g a -> f b -> ApT f g c -> ApT f g x
ApT a -> b -> c -> b
x g a
ga (forall a. f a -> f' a
phi f b
fb) (forall (f :: * -> *) (f' :: * -> *) (g :: * -> *) b.
(forall a. f a -> f' a) -> ApT f g b -> ApT f' g b
transApT forall a. f a -> f' a
phi ApT f g c
rc)

-- | Lift an applicative action @g x@ to @ApT f g x@
liftT :: g x -> ApT f g x
liftT :: forall (g :: * -> *) x (f :: * -> *). g x -> ApT f g x
liftT = forall (f :: * -> *) (g :: * -> *) x. g x -> ApT f g x
PureT

-- | Lift an uninterpreted action @f x@ to @ApT f g x@
liftF :: Applicative g => f x -> ApT f g x
liftF :: forall (g :: * -> *) (f :: * -> *) x.
Applicative g =>
f x -> ApT f g x
liftF f x
fx = forall (f :: * -> *) (g :: * -> *) x a b c.
(a -> b -> c -> x) -> g a -> f b -> ApT f g c -> ApT f g x
ApT (\()
_ x
x ()
_ -> x
x) (forall (f :: * -> *) a. Applicative f => a -> f a
pure ()) f x
fx (forall (f :: * -> *) a. Applicative f => a -> f a
pure ())

{- | Equivalent to the following definition, but is faster and doesn't require @Applicative g@ constraint.

  @appendApT x prefix fb postfix = x \<$\> prefix \<*\> liftF fb \<*\> postfix@
-}
appendApT :: (a -> b -> c -> x) -> ApT f g a -> f b -> ApT f g c -> ApT f g x
appendApT :: forall a b c x (f :: * -> *) (g :: * -> *).
(a -> b -> c -> x) -> ApT f g a -> f b -> ApT f g c -> ApT f g x
appendApT a -> b -> c -> x
x ApT f g a
prefix f b
fb ApT f g c
postfix = case ApT f g a
prefix of
  PureT g a
ga -> forall (f :: * -> *) (g :: * -> *) x a b c.
(a -> b -> c -> x) -> g a -> f b -> ApT f g c -> ApT f g x
ApT a -> b -> c -> x
x g a
ga f b
fb ApT f g c
postfix
  ApT a -> b -> c -> a
a g a
ga' f b
fb' ApT f g c
prefix' -> forall (f :: * -> *) (g :: * -> *) x a b c.
(a -> b -> c -> x) -> g a -> f b -> ApT f g c -> ApT f g x
ApT (\a
a' b
b' ~(c
c', b
b, c
c) -> a -> b -> c -> x
x (a -> b -> c -> a
a a
a' b
b' c
c') b
b c
c) g a
ga' f b
fb' (forall a b c x (f :: * -> *) (g :: * -> *).
(a -> b -> c -> x) -> ApT f g a -> f b -> ApT f g c -> ApT f g x
appendApT (,,) ApT f g c
prefix' f b
fb ApT f g c
postfix)

{- | Interpret @ApT f g@ into an applicative @h@.

  When @g@ is an @Applicative@ and @gh :: forall a. g a -> h a@ is an applicative transformation,
  @'foldApT' fh gh@ is an applicative transformation too.

  @foldApT@ satisfy the following equations with 'liftF' and 'liftT'.

  > foldApT fh gh . liftF = fh
  > foldApT fh gh . liftT = gh
-}
foldApT :: forall f g h x. Applicative h => (forall a. f a -> h a) -> (forall a. g a -> h a) -> ApT f g x -> h x
foldApT :: forall (f :: * -> *) (g :: * -> *) (h :: * -> *) x.
Applicative h =>
(forall a. f a -> h a)
-> (forall a. g a -> h a) -> ApT f g x -> h x
foldApT forall a. f a -> h a
f2h forall a. g a -> h a
g2h = forall y. ApT f g y -> h y
go
 where
  go :: forall y. ApT f g y -> h y
  go :: forall y. ApT f g y -> h y
go (PureT g y
gx) = forall a. g a -> h a
g2h g y
gx
  go (ApT a -> b -> c -> y
x g a
ga f b
fb ApT f g c
rc) = forall (f :: * -> *) a b c d.
Applicative f =>
(a -> b -> c -> d) -> f a -> f b -> f c -> f d
liftA3 a -> b -> c -> y
x (forall a. g a -> h a
g2h g a
ga) (forall a. f a -> h a
f2h f b
fb) (forall y. ApT f g y -> h y
go ApT f g c
rc)

{- | Perform a monoidal analysis over @ApT f g@ value.

  This is equivalent to use @foldApT@ with the applicative @'Control.Applicative.Const' m@,
  except @m@ doesn't need to be a @Monoid@ but just a @Semigroup@.
-}
foldApT_ :: forall f g m x. Semigroup m => (forall a. f a -> m) -> (forall a. g a -> m) -> ApT f g x -> m
foldApT_ :: forall (f :: * -> *) (g :: * -> *) m x.
Semigroup m =>
(forall a. f a -> m) -> (forall a. g a -> m) -> ApT f g x -> m
foldApT_ forall a. f a -> m
f2m forall a. g a -> m
g2m = forall y. ApT f g y -> m
go
 where
  go :: forall y. ApT f g y -> m
  go :: forall y. ApT f g y -> m
go (PureT g y
gx) = forall a. g a -> m
g2m g y
gx
  go (ApT a -> b -> c -> y
_ g a
ga f b
fb ApT f g c
rc) = forall a. g a -> m
g2m g a
ga forall a. Semigroup a => a -> a -> a
<> forall a. f a -> m
f2m f b
fb forall a. Semigroup a => a -> a -> a
<> forall y. ApT f g y -> m
go ApT f g c
rc

-- | Collapsing @ApT@ nested left-to-right.
fjoinApTLeft :: forall f g x. ApT f (ApT f g) x -> ApT f g x
fjoinApTLeft :: forall (f :: * -> *) (g :: * -> *) x.
ApT f (ApT f g) x -> ApT f g x
fjoinApTLeft = forall y. ApT f (ApT f g) y -> ApT f g y
go
 where
  go :: forall y. ApT f (ApT f g) y -> ApT f g y
  go :: forall y. ApT f (ApT f g) y -> ApT f g y
go (PureT ApT f g y
inner) = ApT f g y
inner
  go (ApT a -> b -> c -> y
y ApT f g a
inner f b
fb ApT f (ApT f g) c
rest) = forall a b c x (f :: * -> *) (g :: * -> *).
(a -> b -> c -> x) -> ApT f g a -> f b -> ApT f g c -> ApT f g x
appendApT a -> b -> c -> y
y ApT f g a
inner f b
fb (forall y. ApT f (ApT f g) y -> ApT f g y
go ApT f (ApT f g) c
rest)

-- | Collapsing @ApT@ nested right-to-left.
fjoinApTRight :: Applicative g => ApT (ApT f g) g x -> ApT f g x
fjoinApTRight :: forall (g :: * -> *) (f :: * -> *) x.
Applicative g =>
ApT (ApT f g) g x -> ApT f g x
fjoinApTRight = forall (f :: * -> *) (g :: * -> *) (h :: * -> *) x.
Applicative h =>
(forall a. f a -> h a)
-> (forall a. g a -> h a) -> ApT f g x -> h x
foldApT forall a. a -> a
id forall (g :: * -> *) x (f :: * -> *). g x -> ApT f g x
liftT

-------------
instance (Foldable f, Foldable g) => Foldable (ApT f g) where
  foldMap :: forall m a. Monoid m => (a -> m) -> ApT f g a -> m
foldMap a -> m
f (PureT g a
gx) = forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap a -> m
f g a
gx
  foldMap a -> m
f (ApT a -> b -> c -> a
x g a
ga f b
fb ApT f g c
rc) =
    forall (f :: * -> *) m a.
(Foldable f, Monoid m) =>
f a -> (a -> m) -> m
foldFor g a
ga forall a b. (a -> b) -> a -> b
$ \a
a ->
      forall (f :: * -> *) m a.
(Foldable f, Monoid m) =>
f a -> (a -> m) -> m
foldFor f b
fb forall a b. (a -> b) -> a -> b
$ \b
b ->
        forall (f :: * -> *) m a.
(Foldable f, Monoid m) =>
f a -> (a -> m) -> m
foldFor ApT f g c
rc forall a b. (a -> b) -> a -> b
$ \c
c -> a -> m
f (a -> b -> c -> a
x a
a b
b c
c)
  length :: forall a. ApT f g a -> Int
length = forall any. Int -> ApT f g any -> Int
go Int
1
   where
    go :: forall any. Int -> ApT f g any -> Int
    go :: forall any. Int -> ApT f g any -> Int
go Int
0 ApT f g any
_ = Int
0
    go Int
n (PureT g any
gx) = Int
n forall a. Num a => a -> a -> a
* forall (t :: * -> *) a. Foldable t => t a -> Int
length g any
gx
    go Int
n (ApT a -> b -> c -> any
_ g a
f f b
g ApT f g c
r) = forall any. Int -> ApT f g any -> Int
go (forall (t :: * -> *) a. Foldable t => t a -> Int
length g a
f forall a. Num a => a -> a -> a
* forall (t :: * -> *) a. Foldable t => t a -> Int
length f b
g forall a. Num a => a -> a -> a
* Int
n) ApT f g c
r
  null :: forall a. ApT f g a -> Bool
null (PureT g a
gx) = forall (t :: * -> *) a. Foldable t => t a -> Bool
null g a
gx
  null (ApT a -> b -> c -> a
_ g a
ga f b
fb ApT f g c
rc) = forall (t :: * -> *) a. Foldable t => t a -> Bool
null g a
ga Bool -> Bool -> Bool
|| forall (t :: * -> *) a. Foldable t => t a -> Bool
null f b
fb Bool -> Bool -> Bool
|| forall (t :: * -> *) a. Foldable t => t a -> Bool
null ApT f g c
rc

foldFor :: (Foldable f, Monoid m) => f a -> (a -> m) -> m
foldFor :: forall (f :: * -> *) m a.
(Foldable f, Monoid m) =>
f a -> (a -> m) -> m
foldFor = forall a b c. (a -> b -> c) -> b -> a -> c
flip forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap

{- | Printable value indicating \"shape\" of @ApT f g@ functor.
  If you forget the data of elements from @ApT f g x@, and leave numbers indicating
  which index these data was in the @ApT f g@, that is @ApIx f g@.

    >>> xFn = (\a b c -> if a then show b else c)
    >>> x = ApT xFn [True, False] [10, 20] (PureT ["Moo"])
    >>> toList x
    ["10", "20", "Moo", "Moo"]

A value of type @ApIx [] []@ corresponding to @x@ represents it was made from the three lists
of length @2,2,1@ each. In @ApIx f g@ values, instead of having the original contents, they contain
@Int@ values to conveniently calculate the indices of the value in @toList x@.

    >>> indices x
    ApIx [0, 2] [0, 1] (PureIx [0])

-}
data ApIx f g where
  PureIx :: g Int -> ApIx f g
  ApIx :: g Int -> f Int -> ApIx f g -> ApIx f g

deriving stock instance (Show (f Int), Show (g Int)) => Show (ApIx f g)
deriving stock instance (Eq (f Int), Eq (g Int)) => Eq (ApIx f g)
deriving stock instance (Ord (f Int), Ord (g Int)) => Ord (ApIx f g)

space :: ShowS
space :: ShowS
space = Char -> ShowS
showChar Char
' '

-- | Turn a shape value @ApIx f g@ to the actual @ApT f g Int@ value
--   containing indices.
fromIx :: Functor g => ApIx f g -> ApT f g Int
fromIx :: forall (g :: * -> *) (f :: * -> *).
Functor g =>
ApIx f g -> ApT f g Int
fromIx (PureIx g Int
gi) = forall (f :: * -> *) (g :: * -> *) x. g x -> ApT f g x
PureT g Int
gi
fromIx (ApIx g Int
gi f Int
fj ApIx f g
r) = forall (f :: * -> *) (g :: * -> *) x a b c.
(a -> b -> c -> x) -> g a -> f b -> ApT f g c -> ApT f g x
ApT (\Int
i Int
j Int
k -> Int
i forall a. Num a => a -> a -> a
+ Int
j forall a. Num a => a -> a -> a
+ Int
k) g Int
gi f Int
fj (forall (g :: * -> *) (f :: * -> *).
Functor g =>
ApIx f g -> ApT f g Int
fromIx ApIx f g
r)

lengthIx :: (Foldable f, Foldable g) => ApIx f g -> Int
lengthIx :: forall (f :: * -> *) (g :: * -> *).
(Foldable f, Foldable g) =>
ApIx f g -> Int
lengthIx = forall {t :: * -> *} {t :: * -> *}.
(Foldable t, Foldable t) =>
Int -> ApIx t t -> Int
go Int
1
 where
  go :: Int -> ApIx t t -> Int
go Int
0 ApIx t t
_ = Int
0
  go Int
n (PureIx t Int
g) = Int
n forall a. Num a => a -> a -> a
* forall (t :: * -> *) a. Foldable t => t a -> Int
length t Int
g
  go Int
n (ApIx t Int
g t Int
f ApIx t t
r) = Int -> ApIx t t -> Int
go (Int
n forall a. Num a => a -> a -> a
* forall (t :: * -> *) a. Foldable t => t a -> Int
length t Int
g forall a. Num a => a -> a -> a
* forall (t :: * -> *) a. Foldable t => t a -> Int
length t Int
f) ApIx t t
r

indicesF :: (Traversable f) => f a -> f Int
indicesF :: forall (f :: * -> *) a. Traversable f => f a -> f Int
indicesF = forall a b. (a, b) -> b
snd forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (t :: * -> *) s a b.
Traversable t =>
(s -> a -> (s, b)) -> s -> t a -> (s, t b)
mapAccumL (\Int
n a
_ -> Int
n seq :: forall a b. a -> b -> b
`seq` (Int
n forall a. Num a => a -> a -> a
+ Int
1, Int
n)) Int
0

-- | Extract only a shape from @ApT f g x@ and make it @ApIx f g@.
indices :: forall f g x. (Traversable f, Traversable g) => ApT f g x -> ApIx f g
indices :: forall (f :: * -> *) (g :: * -> *) x.
(Traversable f, Traversable g) =>
ApT f g x -> ApIx f g
indices ApT f g x
u
  | forall (t :: * -> *) a. Foldable t => t a -> Bool
null ApT f g x
u    = forall z. ApT f g z -> ApIx f g
ripoff ApT f g x
u
  | Bool
otherwise = forall a b. (a, b) -> b
snd (forall y. ApT f g y -> (Int, ApIx f g)
go ApT f g x
u)
  where
    ripoff :: ApT f g z -> ApIx f g
    ripoff :: forall z. ApT f g z -> ApIx f g
ripoff (PureT g z
gx) = forall (g :: * -> *) (f :: * -> *). g Int -> ApIx f g
PureIx (Int
0 forall (f :: * -> *) a b. Functor f => a -> f b -> f a
<$ g z
gx)
    ripoff (ApT a -> b -> c -> z
_ g a
ga f b
fb ApT f g c
rc) = forall (g :: * -> *) (f :: * -> *).
g Int -> f Int -> ApIx f g -> ApIx f g
ApIx (Int
0 forall (f :: * -> *) a b. Functor f => a -> f b -> f a
<$ g a
ga) (Int
0 forall (f :: * -> *) a b. Functor f => a -> f b -> f a
<$ f b
fb) (forall z. ApT f g z -> ApIx f g
ripoff ApT f g c
rc)
    
    go :: forall y. ApT f g y -> (Int, ApIx f g)
    go :: forall y. ApT f g y -> (Int, ApIx f g)
go (PureT g y
gx) = (forall (t :: * -> *) a. Foldable t => t a -> Int
length g y
gx, forall (g :: * -> *) (f :: * -> *). g Int -> ApIx f g
PureIx (forall (f :: * -> *) a. Traversable f => f a -> f Int
indicesF g y
gx))
    go (ApT a -> b -> c -> y
_ g a
ga f b
fb ApT f g c
rc) =
      let lenG :: Int
lenG = forall (t :: * -> *) a. Foldable t => t a -> Int
length g a
ga
          lenF :: Int
lenF = forall (t :: * -> *) a. Foldable t => t a -> Int
length f b
fb
          (Int
lenR, ApIx f g
rk) = forall y. ApT f g y -> (Int, ApIx f g)
go ApT f g c
rc
          gi' :: g Int
gi' = (Int
lenF forall a. Num a => a -> a -> a
* Int
lenR forall a. Num a => a -> a -> a
*) forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall (f :: * -> *) a. Traversable f => f a -> f Int
indicesF g a
ga
          fj' :: f Int
fj' = (Int
lenR forall a. Num a => a -> a -> a
*) forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall (f :: * -> *) a. Traversable f => f a -> f Int
indicesF f b
fb
          len :: Int
len = Int
lenG forall a. Num a => a -> a -> a
* Int
lenF forall a. Num a => a -> a -> a
* Int
lenR
      in Int
len seq :: forall a b. a -> b -> b
`seq` (Int
len, forall (g :: * -> *) (f :: * -> *).
g Int -> f Int -> ApIx f g -> ApIx f g
ApIx g Int
gi' f Int
fj' ApIx f g
rk)

-- | Construct an @ApT f g x@ value from a shape @ApIx f g@ and a list of values.
--
--   For any @u :: ApT f g x@, the following property holds.
--
--   > reconstruct (indices u) (toList u) == u
--   
--   @reconstruct shape xs@ raises 'error' if the length of list @xs@ does not match
--   the length calculated from @shape@.
reconstruct :: (HasCallStack, Foldable f, Foldable g, Functor g) => ApIx f g -> [x] -> ApT f g x
reconstruct :: forall (f :: * -> *) (g :: * -> *) x.
(HasCallStack, Foldable f, Foldable g, Functor g) =>
ApIx f g -> [x] -> ApT f g x
reconstruct ApIx f g
shape [x]
xs
  | forall (t :: * -> *) a. Foldable t => t a -> Int
length [x]
xs forall a. Eq a => a -> a -> Bool
== Int
n = (Array Int x
xsArr forall i e. Ix i => Array i e -> i -> e
Arr.!) forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall (g :: * -> *) (f :: * -> *).
Functor g =>
ApIx f g -> ApT f g Int
fromIx ApIx f g
shape
  | Bool
otherwise = forall a. HasCallStack => String -> a
error String
"Wrong number of elements in the table"
  where
    n :: Int
n = forall (f :: * -> *) (g :: * -> *).
(Foldable f, Foldable g) =>
ApIx f g -> Int
lengthIx ApIx f g
shape
    xsArr :: Array Int x
xsArr = forall i e. Ix i => (i, i) -> [e] -> Array i e
Arr.listArray (Int
0, Int
n forall a. Num a => a -> a -> a
- Int
1) [x]
xs

instance (Traversable f, Traversable g) => Traversable (ApT f g) where
  traverse :: forall (f :: * -> *) a b.
Applicative f =>
(a -> f b) -> ApT f g a -> f (ApT f g b)
traverse a -> f b
f ApT f g a
u = forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (forall (f :: * -> *) (g :: * -> *) x.
(HasCallStack, Foldable f, Foldable g, Functor g) =>
ApIx f g -> [x] -> ApT f g x
reconstruct (forall (f :: * -> *) (g :: * -> *) x.
(Traversable f, Traversable g) =>
ApT f g x -> ApIx f g
indices ApT f g a
u)) (forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
traverse a -> f b
f (forall (t :: * -> *) a. Foldable t => t a -> [a]
toList ApT f g a
u))

instance (Traversable f, Show (f Int), Traversable g, Show (g Int), Show a) => Show (ApT f g a) where
  showsPrec :: Int -> ApT f g a -> ShowS
showsPrec = forall (f :: * -> *) a. (Show1 f, Show a) => Int -> f a -> ShowS
showsPrec1

instance (Traversable f, Show (f Int), Traversable g, Show (g Int)) => Show1 (ApT f g) where
  liftShowsPrec :: forall a.
(Int -> a -> ShowS) -> ([a] -> ShowS) -> Int -> ApT f g a -> ShowS
liftShowsPrec Int -> a -> ShowS
showsPrecX [a] -> ShowS
showListX Int
p ApT f g a
u =
    Bool -> ShowS -> ShowS
showParen (Int
p forall a. Ord a => a -> a -> Bool
> Int
10) forall a b. (a -> b) -> a -> b
$ (String
"reconstruct " forall a. [a] -> [a] -> [a]
++) forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a. Show a => Int -> a -> ShowS
showsPrec Int
11 (forall (f :: * -> *) (g :: * -> *) x.
(Traversable f, Traversable g) =>
ApT f g x -> ApIx f g
indices ApT f g a
u) forall b c a. (b -> c) -> (a -> b) -> a -> c
. ShowS
space forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (f :: * -> *) a.
Show1 f =>
(Int -> a -> ShowS) -> ([a] -> ShowS) -> Int -> f a -> ShowS
liftShowsPrec Int -> a -> ShowS
showsPrecX [a] -> ShowS
showListX Int
11 (forall (t :: * -> *) a. Foldable t => t a -> [a]
toList ApT f g a
u)

instance (Eq1 f, Eq1 g, Eq a) => Eq (ApT f g a) where
  == :: ApT f g a -> ApT f g a -> Bool
(==) = forall (f :: * -> *) a. (Eq1 f, Eq a) => f a -> f a -> Bool
eq1

instance (Eq1 f, Eq1 g) => Eq1 (ApT f g) where
  liftEq :: forall a b. (a -> b -> Bool) -> ApT f g a -> ApT f g b -> Bool
liftEq a -> b -> Bool
eq (PureT g a
g1) (PureT g b
g2) = forall (f :: * -> *) a b.
Eq1 f =>
(a -> b -> Bool) -> f a -> f b -> Bool
liftEq a -> b -> Bool
eq g a
g1 g b
g2
  liftEq a -> b -> Bool
eq (ApT a -> b -> c -> a
x1 g a
g1 f b
f1 ApT f g c
r1) (ApT a -> b -> c -> b
x2 g a
g2 f b
f2 ApT f g c
r2)
    | forall (f :: * -> *) a b. Eq1 f => f a -> f b -> Bool
emptyEq g a
g1 g a
g2 = forall (f :: * -> *) a b. Eq1 f => f a -> f b -> Bool
boringEq f b
f1 f b
f2 Bool -> Bool -> Bool
&& forall (f :: * -> *) (g :: * -> *) a b.
(Eq1 f, Eq1 g) =>
ApT f g a -> ApT f g b -> Bool
boringEqApT ApT f g c
r1 ApT f g c
r2
    | forall (f :: * -> *) a b. Eq1 f => f a -> f b -> Bool
emptyEq f b
f1 f b
f2 = forall (f :: * -> *) a b. Eq1 f => f a -> f b -> Bool
boringEq g a
g1 g a
g2 Bool -> Bool -> Bool
&& forall (f :: * -> *) (g :: * -> *) a b.
(Eq1 f, Eq1 g) =>
ApT f g a -> ApT f g b -> Bool
boringEqApT ApT f g c
r1 ApT f g c
r2
    | Bool
otherwise =
        forall (f :: * -> *) a b.
Eq1 f =>
f a -> f b -> (a -> b -> Bool) -> Bool
liftEqFor g a
g1 g a
g2 forall a b. (a -> b) -> a -> b
$ \a
a1 a
a2 ->
          forall (f :: * -> *) a b.
Eq1 f =>
f a -> f b -> (a -> b -> Bool) -> Bool
liftEqFor f b
f1 f b
f2 forall a b. (a -> b) -> a -> b
$ \b
b1 b
b2 ->
            forall (f :: * -> *) a b.
Eq1 f =>
f a -> f b -> (a -> b -> Bool) -> Bool
liftEqFor ApT f g c
r1 ApT f g c
r2 forall a b. (a -> b) -> a -> b
$ \c
c1 c
c2 ->
              a -> b -> Bool
eq (a -> b -> c -> a
x1 a
a1 b
b1 c
c1) (a -> b -> c -> b
x2 a
a2 b
b2 c
c2)
  liftEq a -> b -> Bool
_ ApT f g a
_ ApT f g b
_ = Bool
False

instance (Ord1 f, Ord1 g, Ord a) => Ord (ApT f g a) where
  compare :: ApT f g a -> ApT f g a -> Ordering
compare = forall (f :: * -> *) a. (Ord1 f, Ord a) => f a -> f a -> Ordering
compare1

instance (Ord1 f, Ord1 g) => Ord1 (ApT f g) where
  liftCompare :: forall a b.
(a -> b -> Ordering) -> ApT f g a -> ApT f g b -> Ordering
liftCompare a -> b -> Ordering
cmp (PureT g a
g1) (PureT g b
g2) = forall (f :: * -> *) a b.
Ord1 f =>
(a -> b -> Ordering) -> f a -> f b -> Ordering
liftCompare a -> b -> Ordering
cmp g a
g1 g b
g2
  liftCompare a -> b -> Ordering
cmp (ApT a -> b -> c -> a
x1 g a
g1 f b
f1 ApT f g c
r1) (ApT a -> b -> c -> b
x2 g a
g2 f b
f2 ApT f g c
r2)
    | forall (f :: * -> *) a b. Eq1 f => f a -> f b -> Bool
emptyEq g a
g1 g a
g2 = forall (f :: * -> *) a b. Ord1 f => f a -> f b -> Ordering
boringCompare f b
f1 f b
f2 forall a. Semigroup a => a -> a -> a
<> forall (f :: * -> *) (g :: * -> *) a b.
(Ord1 f, Ord1 g) =>
ApT f g a -> ApT f g b -> Ordering
boringCompareApT ApT f g c
r1 ApT f g c
r2
    | forall (f :: * -> *) a b. Eq1 f => f a -> f b -> Bool
emptyEq f b
f1 f b
f2 = forall (f :: * -> *) a b. Ord1 f => f a -> f b -> Ordering
boringCompare g a
g1 g a
g2 forall a. Semigroup a => a -> a -> a
<> forall (f :: * -> *) (g :: * -> *) a b.
(Ord1 f, Ord1 g) =>
ApT f g a -> ApT f g b -> Ordering
boringCompareApT ApT f g c
r1 ApT f g c
r2
    | Bool
otherwise =
        forall (f :: * -> *) a b.
Ord1 f =>
f a -> f b -> (a -> b -> Ordering) -> Ordering
liftCompareFor g a
g1 g a
g2 forall a b. (a -> b) -> a -> b
$ \a
a1 a
a2 ->
          forall (f :: * -> *) a b.
Ord1 f =>
f a -> f b -> (a -> b -> Ordering) -> Ordering
liftCompareFor f b
f1 f b
f2 forall a b. (a -> b) -> a -> b
$ \b
b1 b
b2 ->
            forall (f :: * -> *) a b.
Ord1 f =>
f a -> f b -> (a -> b -> Ordering) -> Ordering
liftCompareFor ApT f g c
r1 ApT f g c
r2 forall a b. (a -> b) -> a -> b
$ \c
c1 c
c2 ->
              a -> b -> Ordering
cmp (a -> b -> c -> a
x1 a
a1 b
b1 c
c1) (a -> b -> c -> b
x2 a
a2 b
b2 c
c2)
  liftCompare a -> b -> Ordering
_ PureT{} ApT{} = Ordering
LT
  liftCompare a -> b -> Ordering
_ ApT{} PureT{} = Ordering
GT

-- | Order shuffled 'liftEq'
liftEqFor :: Eq1 f => f a -> f b -> (a -> b -> Bool) -> Bool
liftEqFor :: forall (f :: * -> *) a b.
Eq1 f =>
f a -> f b -> (a -> b -> Bool) -> Bool
liftEqFor f a
f1 f b
f2 a -> b -> Bool
eq = forall (f :: * -> *) a b.
Eq1 f =>
(a -> b -> Bool) -> f a -> f b -> Bool
liftEq a -> b -> Bool
eq f a
f1 f b
f2

-- | Order shuffled 'liftCompare'
liftCompareFor :: Ord1 f => f a -> f b -> (a -> b -> Ordering) -> Ordering
liftCompareFor :: forall (f :: * -> *) a b.
Ord1 f =>
f a -> f b -> (a -> b -> Ordering) -> Ordering
liftCompareFor f a
f1 f b
f2 a -> b -> Ordering
cmp = forall (f :: * -> *) a b.
Ord1 f =>
(a -> b -> Ordering) -> f a -> f b -> Ordering
liftCompare a -> b -> Ordering
cmp f a
f1 f b
f2

emptyEq, boringEq :: Eq1 f => f a -> f b -> Bool
emptyEq :: forall (f :: * -> *) a b. Eq1 f => f a -> f b -> Bool
emptyEq = forall (f :: * -> *) a b.
Eq1 f =>
(a -> b -> Bool) -> f a -> f b -> Bool
liftEq (\a
_ b
_ -> Bool
False)
boringEq :: forall (f :: * -> *) a b. Eq1 f => f a -> f b -> Bool
boringEq = forall (f :: * -> *) a b.
Eq1 f =>
(a -> b -> Bool) -> f a -> f b -> Bool
liftEq (\a
_ b
_ -> Bool
True)

boringCompare :: Ord1 f => f a -> f b -> Ordering
boringCompare :: forall (f :: * -> *) a b. Ord1 f => f a -> f b -> Ordering
boringCompare = forall (f :: * -> *) a b.
Ord1 f =>
(a -> b -> Ordering) -> f a -> f b -> Ordering
liftCompare (\a
_ b
_ -> Ordering
EQ)

boringEqApT :: (Eq1 f, Eq1 g) => ApT f g a -> ApT f g b -> Bool
boringEqApT :: forall (f :: * -> *) (g :: * -> *) a b.
(Eq1 f, Eq1 g) =>
ApT f g a -> ApT f g b -> Bool
boringEqApT (PureT g a
g1) (PureT g b
g2) = forall (f :: * -> *) a b. Eq1 f => f a -> f b -> Bool
boringEq g a
g1 g b
g2
boringEqApT (ApT a -> b -> c -> a
_ g a
g1 f b
f1 ApT f g c
r1) (ApT a -> b -> c -> b
_ g a
g2 f b
f2 ApT f g c
r2) = forall (f :: * -> *) a b. Eq1 f => f a -> f b -> Bool
boringEq g a
g1 g a
g2 Bool -> Bool -> Bool
&& forall (f :: * -> *) a b. Eq1 f => f a -> f b -> Bool
boringEq f b
f1 f b
f2 Bool -> Bool -> Bool
&& forall (f :: * -> *) (g :: * -> *) a b.
(Eq1 f, Eq1 g) =>
ApT f g a -> ApT f g b -> Bool
boringEqApT ApT f g c
r1 ApT f g c
r2
boringEqApT ApT f g a
_ ApT f g b
_ = Bool
False

boringCompareApT :: (Ord1 f, Ord1 g) => ApT f g a -> ApT f g b -> Ordering
boringCompareApT :: forall (f :: * -> *) (g :: * -> *) a b.
(Ord1 f, Ord1 g) =>
ApT f g a -> ApT f g b -> Ordering
boringCompareApT (PureT g a
g1) (PureT g b
g2) = forall (f :: * -> *) a b. Ord1 f => f a -> f b -> Ordering
boringCompare g a
g1 g b
g2
boringCompareApT (ApT a -> b -> c -> a
_ g a
g1 f b
f1 ApT f g c
r1) (ApT a -> b -> c -> b
_ g a
g2 f b
f2 ApT f g c
r2) = forall (f :: * -> *) a b. Ord1 f => f a -> f b -> Ordering
boringCompare g a
g1 g a
g2 forall a. Semigroup a => a -> a -> a
<> forall (f :: * -> *) a b. Ord1 f => f a -> f b -> Ordering
boringCompare f b
f1 f b
f2 forall a. Semigroup a => a -> a -> a
<> forall (f :: * -> *) (g :: * -> *) a b.
(Ord1 f, Ord1 g) =>
ApT f g a -> ApT f g b -> Ordering
boringCompareApT ApT f g c
r1 ApT f g c
r2
boringCompareApT PureT{} ApT{} = Ordering
LT
boringCompareApT ApT{} PureT{} = Ordering
GT