------------------------------------------------------------------------
-- |
-- Module           : Data.Parameterized.TraversableFC.WithIndex
-- Copyright        : (c) Galois, Inc 2021
-- Maintainer       : Langston Barrett
-- Description      : 'TraversableFC' classes, but with indices.
--
-- As in the package indexed-traversable.
------------------------------------------------------------------------
{-# LANGUAGE PolyKinds #-}
{-# LANGUAGE RankNTypes #-}
{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE TypeOperators #-}

module Data.Parameterized.TraversableFC.WithIndex
  ( FunctorFCWithIndex(..)
  , FoldableFCWithIndex(..)
  , ifoldlMFC
  , ifoldrMFC
  , iallFC
  , ianyFC
  , TraversableFCWithIndex(..)
  , imapFCDefault
  , ifoldMapFCDefault
  ) where

import Data.Functor.Const (Const(Const, getConst))
import Data.Functor.Identity (Identity(Identity, runIdentity))
import Data.Kind
import Data.Monoid (All(..), Any(..), Endo(Endo), appEndo, Dual(Dual, getDual))
import Data.Profunctor.Unsafe ((#.))
import GHC.Exts (build)

import Data.Parameterized.Classes
import Data.Parameterized.TraversableFC

class FunctorFC t => FunctorFCWithIndex (t :: (k -> Type) -> l -> Type) where
  -- | Like 'fmapFC', but with an index.
  --
  -- @
  -- 'fmapFC' f ≡ 'imapFC' ('const' f)
  -- @
  imapFC ::
    forall f g z.
    (forall x. IndexF (t f z) x -> f x -> g x)
    -> t f z
    -> t g z

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

class (FoldableFC t, FunctorFCWithIndex t) => FoldableFCWithIndex (t :: (k -> Type) -> l -> Type) where

  -- | Like 'foldMapFC', but with an index.
  --
  -- @
  -- 'foldMapFC' f ≡ 'ifoldMapFC' ('const' f)
  -- @
  ifoldMapFC ::
    forall f m z.
    Monoid m =>
    (forall x. IndexF (t f z) x -> f x -> m) ->
    t f z ->
    m
  ifoldMapFC forall (x :: k). IndexF (t f z) x -> f x -> m
f = forall k l (t :: (k -> *) -> l -> *) (z :: l) (f :: k -> *) b.
FoldableFCWithIndex t =>
(forall (x :: k). IndexF (t f z) x -> f x -> b -> b)
-> b -> t f z -> b
ifoldrFC (\IndexF (t f z) x
i f x
x -> forall a. Monoid a => a -> a -> a
mappend (forall (x :: k). IndexF (t f z) x -> f x -> m
f IndexF (t f z) x
i f x
x)) forall a. Monoid a => a
mempty

  -- | Like 'foldrFC', but with an index.
  ifoldrFC ::
    forall z f b.
    (forall x. IndexF (t f z) x -> f x -> b -> b) ->
    b ->
    t f z ->
    b
  ifoldrFC forall (x :: k). IndexF (t f z) x -> f x -> b -> b
f b
z t f z
t = forall a. Endo a -> a -> a
appEndo (forall k l (t :: (k -> *) -> l -> *) (f :: k -> *) m (z :: l).
(FoldableFCWithIndex t, Monoid m) =>
(forall (x :: k). IndexF (t f z) x -> f x -> m) -> t f z -> m
ifoldMapFC (\IndexF (t f z) x
i f x
x -> forall a. (a -> a) -> Endo a
Endo (forall (x :: k). IndexF (t f z) x -> f x -> b -> b
f IndexF (t f z) x
i f x
x)) t f z
t) b
z

  -- | Like 'foldlFC', but with an index.
  ifoldlFC ::
    forall f b z.
    (forall x. IndexF (t f z) x -> b -> f x -> b) ->
    b ->
    t f z ->
    b
  ifoldlFC forall (x :: k). IndexF (t f z) x -> b -> f x -> b
f b
z t f z
t =
    forall a. Endo a -> a -> a
appEndo (forall a. Dual a -> a
getDual (forall k l (t :: (k -> *) -> l -> *) (f :: k -> *) m (z :: l).
(FoldableFCWithIndex t, Monoid m) =>
(forall (x :: k). IndexF (t f z) x -> f x -> m) -> t f z -> m
ifoldMapFC (\IndexF (t f z) x
i f x
e -> forall a. a -> Dual a
Dual (forall a. (a -> a) -> Endo a
Endo (\b
r -> forall (x :: k). IndexF (t f z) x -> b -> f x -> b
f IndexF (t f z) x
i b
r f x
e))) t f z
t)) b
z

  -- | Like 'ifoldrFC', but with an index.
  ifoldrFC' ::
    forall f b z.
    (forall x. IndexF (t f z) x -> f x -> b -> b) ->
    b ->
    t f z ->
    b
  ifoldrFC' forall (x :: k). IndexF (t f z) x -> f x -> b -> b
f0 b
z0 t f z
xs = forall k l (t :: (k -> *) -> l -> *) (f :: k -> *) b (z :: l).
FoldableFCWithIndex t =>
(forall (x :: k). IndexF (t f z) x -> b -> f x -> b)
-> b -> t f z -> b
ifoldlFC (forall {t} {t} {t} {a} {b}.
(t -> t -> t -> a) -> t -> (a -> b) -> t -> t -> b
f' forall (x :: k). IndexF (t f z) x -> f x -> b -> b
f0) forall a. a -> a
id t f z
xs b
z0
    where f' :: (t -> t -> t -> a) -> t -> (a -> b) -> t -> t -> b
f' t -> t -> t -> a
f t
i a -> b
k t
x t
z = a -> b
k forall a b. (a -> b) -> a -> b
$! t -> t -> t -> a
f t
i t
x t
z

  -- | Like 'ifoldlFC', but with an index.
  ifoldlFC' :: forall f b. (forall x. b -> f x -> b) -> (forall x. b -> t f x -> b)
  ifoldlFC' forall (x :: k). b -> f x -> b
f0 b
z0 t f x
xs = forall k l (t :: (k -> *) -> l -> *) (f :: k -> *) b.
FoldableFC t =>
(forall (x :: k). f x -> b -> b)
-> forall (x :: l). b -> t f x -> b
foldrFC (forall {t} {t} {a} {b}. (t -> t -> a) -> t -> (a -> b) -> t -> b
f' forall (x :: k). b -> f x -> b
f0) forall a. a -> a
id t f x
xs b
z0
    where f' :: (t -> t -> a) -> t -> (a -> b) -> t -> b
f' t -> t -> a
f t
x a -> b
k t
z = a -> b
k forall a b. (a -> b) -> a -> b
$! t -> t -> a
f t
z t
x

  -- | Convert structure to list.
  itoListFC ::
    forall f a z.
    (forall x. IndexF (t f z) x -> f x -> a) ->
    t f z ->
    [a]
  itoListFC forall (x :: k). IndexF (t f z) x -> f x -> a
f t f z
t = forall a. (forall b. (a -> b -> b) -> b -> b) -> [a]
build (\a -> b -> b
c b
n -> forall k l (t :: (k -> *) -> l -> *) (z :: l) (f :: k -> *) b.
FoldableFCWithIndex t =>
(forall (x :: k). IndexF (t f z) x -> f x -> b -> b)
-> b -> t f z -> b
ifoldrFC (\IndexF (t f z) x
i f x
e b
v -> a -> b -> b
c (forall (x :: k). IndexF (t f z) x -> f x -> a
f IndexF (t f z) x
i f x
e) b
v) b
n t f z
t)

-- | Like 'foldlMFC', but with an index.
ifoldlMFC ::
  FoldableFCWithIndex t =>
  Monad m =>
  (forall x. IndexF (t f z) x -> b -> f x -> m b) ->
  b ->
  t f z ->
  m b
ifoldlMFC :: forall {k} {l} (t :: (k -> *) -> l -> *) (m :: * -> *)
       (f :: k -> *) (z :: l) b.
(FoldableFCWithIndex t, Monad m) =>
(forall (x :: k). IndexF (t f z) x -> b -> f x -> m b)
-> b -> t f z -> m b
ifoldlMFC forall (x :: k). IndexF (t f z) x -> b -> f x -> m b
f b
z0 t f z
xs = forall k l (t :: (k -> *) -> l -> *) (f :: k -> *) b (z :: l).
FoldableFCWithIndex t =>
(forall (x :: k). IndexF (t f z) x -> b -> f x -> b)
-> b -> t f z -> b
ifoldlFC (\IndexF (t f z) x
i b -> m b
k f x
x b
z -> forall (x :: k). IndexF (t f z) x -> b -> f x -> m b
f IndexF (t f z) x
i b
z f x
x forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= b -> m b
k) forall (m :: * -> *) a. Monad m => a -> m a
return t f z
xs b
z0

-- | Like 'foldrMFC', but with an index.
ifoldrMFC ::
  FoldableFCWithIndex t =>
  Monad m =>
  (forall x. IndexF (t f z) x -> f x -> b -> m b) ->
  b ->
  t f z ->
  m b
ifoldrMFC :: forall {k} {l} (t :: (k -> *) -> l -> *) (m :: * -> *)
       (f :: k -> *) (z :: l) b.
(FoldableFCWithIndex t, Monad m) =>
(forall (x :: k). IndexF (t f z) x -> f x -> b -> m b)
-> b -> t f z -> m b
ifoldrMFC forall (x :: k). IndexF (t f z) x -> f x -> b -> m b
f b
z0 t f z
xs = forall k l (t :: (k -> *) -> l -> *) (f :: k -> *) b (z :: l).
FoldableFCWithIndex t =>
(forall (x :: k). IndexF (t f z) x -> b -> f x -> b)
-> b -> t f z -> b
ifoldlFC (\IndexF (t f z) x
i b -> m b
k f x
x b
z -> forall (x :: k). IndexF (t f z) x -> f x -> b -> m b
f IndexF (t f z) x
i f x
x b
z forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= b -> m b
k) forall (m :: * -> *) a. Monad m => a -> m a
return t f z
xs b
z0

-- | Like 'allFC', but with an index.
iallFC ::
  FoldableFCWithIndex t =>
  (forall x. IndexF (t f z) x -> f x -> Bool) ->
  t f z ->
  Bool
iallFC :: forall {k} {l} (t :: (k -> *) -> l -> *) (f :: k -> *) (z :: l).
FoldableFCWithIndex t =>
(forall (x :: k). IndexF (t f z) x -> f x -> Bool) -> t f z -> Bool
iallFC forall (x :: k). IndexF (t f z) x -> f x -> Bool
p = All -> Bool
getAll forall (p :: * -> * -> *) a b c (q :: * -> * -> *).
(Profunctor p, Coercible c b) =>
q b c -> p a b -> p a c
#. forall k l (t :: (k -> *) -> l -> *) (f :: k -> *) m (z :: l).
(FoldableFCWithIndex t, Monoid m) =>
(forall (x :: k). IndexF (t f z) x -> f x -> m) -> t f z -> m
ifoldMapFC (\IndexF (t f z) x
i f x
x -> Bool -> All
All (forall (x :: k). IndexF (t f z) x -> f x -> Bool
p IndexF (t f z) x
i f x
x))

-- | Like 'anyFC', but with an index.
ianyFC ::
  FoldableFCWithIndex t =>
  (forall x. IndexF (t f z) x -> f x -> Bool) ->
  t f z -> Bool
ianyFC :: forall {k} {l} (t :: (k -> *) -> l -> *) (f :: k -> *) (z :: l).
FoldableFCWithIndex t =>
(forall (x :: k). IndexF (t f z) x -> f x -> Bool) -> t f z -> Bool
ianyFC forall (x :: k). IndexF (t f z) x -> f x -> Bool
p = Any -> Bool
getAny forall (p :: * -> * -> *) a b c (q :: * -> * -> *).
(Profunctor p, Coercible c b) =>
q b c -> p a b -> p a c
#. forall k l (t :: (k -> *) -> l -> *) (f :: k -> *) m (z :: l).
(FoldableFCWithIndex t, Monoid m) =>
(forall (x :: k). IndexF (t f z) x -> f x -> m) -> t f z -> m
ifoldMapFC (\IndexF (t f z) x
i f x
x -> Bool -> Any
Any (forall (x :: k). IndexF (t f z) x -> f x -> Bool
p IndexF (t f z) x
i f x
x))

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

class (TraversableFC t, FoldableFCWithIndex t) => TraversableFCWithIndex (t :: (k -> Type) -> l -> Type) where
  -- | Like 'traverseFC', but with an index.
  --
  -- @
  -- 'traverseFC' f ≡ 'itraverseFC' ('const' f)
  -- @
  itraverseFC ::
    forall m z f g.
    Applicative m =>
    (forall x. IndexF (t f z) x -> f x -> m (g x)) ->
    t f z ->
    m (t g z)

imapFCDefault ::
  forall t f g z.
  TraversableFCWithIndex t =>
  (forall x. IndexF (t f z) x -> f x -> g x)
  -> t f z
  -> t g z
imapFCDefault :: forall {k} {l} (t :: (k -> *) -> l -> *) (f :: k -> *)
       (g :: k -> *) (z :: l).
TraversableFCWithIndex t =>
(forall (x :: k). IndexF (t f z) x -> f x -> g x) -> t f z -> t g z
imapFCDefault forall (x :: k). IndexF (t f z) x -> f x -> g x
f = forall a. Identity a -> a
runIdentity forall (p :: * -> * -> *) a b c (q :: * -> * -> *).
(Profunctor p, Coercible c b) =>
q b c -> p a b -> p a c
#. forall k l (t :: (k -> *) -> l -> *) (m :: * -> *) (z :: l)
       (f :: k -> *) (g :: k -> *).
(TraversableFCWithIndex t, Applicative m) =>
(forall (x :: k). IndexF (t f z) x -> f x -> m (g x))
-> t f z -> m (t g z)
itraverseFC (\IndexF (t f z) x
i f x
x -> forall a. a -> Identity a
Identity (forall (x :: k). IndexF (t f z) x -> f x -> g x
f IndexF (t f z) x
i f x
x))
{-# INLINEABLE imapFCDefault #-}

ifoldMapFCDefault ::
  forall t m z f.
  TraversableFCWithIndex t =>
  Monoid m =>
  (forall x. IndexF (t f z) x -> f x -> m) ->
  t f z ->
  m
ifoldMapFCDefault :: forall {k} {l} (t :: (k -> *) -> l -> *) m (z :: l) (f :: k -> *).
(TraversableFCWithIndex t, Monoid m) =>
(forall (x :: k). IndexF (t f z) x -> f x -> m) -> t f z -> m
ifoldMapFCDefault forall (x :: k). IndexF (t f z) x -> f x -> m
f = forall {k} a (b :: k). Const a b -> a
getConst forall (p :: * -> * -> *) a b c (q :: * -> * -> *).
(Profunctor p, Coercible c b) =>
q b c -> p a b -> p a c
#. forall k l (t :: (k -> *) -> l -> *) (m :: * -> *) (z :: l)
       (f :: k -> *) (g :: k -> *).
(TraversableFCWithIndex t, Applicative m) =>
(forall (x :: k). IndexF (t f z) x -> f x -> m (g x))
-> t f z -> m (t g z)
itraverseFC (\IndexF (t f z) x
i f x
x -> forall {k} a (b :: k). a -> Const a b
Const (forall (x :: k). IndexF (t f z) x -> f x -> m
f IndexF (t f z) x
i f x
x))
{-# INLINEABLE ifoldMapFCDefault #-}