------------------------------------------------------------------------
-- |
-- 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 (x :: k). IndexF (t f z) x -> f x -> m -> m)
-> m -> t f z -> m
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 -> m -> m -> m
forall a. Monoid a => a -> a -> a
mappend (IndexF (t f z) x -> f x -> m
forall (x :: k). IndexF (t f z) x -> f x -> m
f IndexF (t f z) x
i f x
x)) m
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 = Endo b -> b -> b
forall a. Endo a -> a -> a
appEndo ((forall (x :: k). IndexF (t f z) x -> f x -> Endo b)
-> t f z -> Endo b
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 -> (b -> b) -> Endo b
forall a. (a -> a) -> Endo a
Endo (IndexF (t f z) x -> f x -> b -> b
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 =
    Endo b -> b -> b
forall a. Endo a -> a -> a
appEndo (Dual (Endo b) -> Endo b
forall a. Dual a -> a
getDual ((forall (x :: k). IndexF (t f z) x -> f x -> Dual (Endo b))
-> t f z -> Dual (Endo b)
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 -> Endo b -> Dual (Endo b)
forall a. a -> Dual a
Dual ((b -> b) -> Endo b
forall a. (a -> a) -> Endo a
Endo (\b
r -> IndexF (t f z) x -> b -> f x -> b
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 (x :: k). IndexF (t f z) x -> (b -> b) -> f x -> b -> b)
-> (b -> b) -> t f z -> b -> b
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 -> f x -> b -> b)
-> IndexF (t f z) x -> (b -> b) -> f x -> b -> b
forall t t t a b.
(t -> t -> t -> a) -> t -> (a -> b) -> t -> t -> b
f' IndexF (t f z) x -> f x -> b -> b
forall (x :: k). IndexF (t f z) x -> f x -> b -> b
f0) b -> b
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 (a -> b) -> a -> b
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 (x :: k). f x -> (b -> b) -> b -> b)
-> (b -> b) -> t f x -> b -> b
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 ((b -> f x -> b) -> f x -> (b -> b) -> b -> b
forall t t a b. (t -> t -> a) -> t -> (a -> b) -> t -> b
f' b -> f x -> b
forall (x :: k). b -> f x -> b
f0) b -> b
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 (a -> b) -> a -> b
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 b. (a -> b -> b) -> b -> b) -> [a]
forall a. (forall b. (a -> b -> b) -> b -> b) -> [a]
build (\a -> b -> b
c b
n -> (forall (x :: k). IndexF (t f z) x -> f x -> b -> b)
-> b -> t f z -> b
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 (IndexF (t f z) x -> f x -> a
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 (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 (x :: k).
 IndexF (t f z) x -> (b -> m b) -> f x -> b -> m b)
-> (b -> m b) -> t f z -> b -> m b
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 -> IndexF (t f z) x -> b -> f x -> m b
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 m b -> (b -> m b) -> m b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= b -> m b
k) b -> m b
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 (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 (x :: k).
 IndexF (t f z) x -> (b -> m b) -> f x -> b -> m b)
-> (b -> m b) -> t f z -> b -> m b
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 -> IndexF (t f z) x -> f x -> b -> m b
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 m b -> (b -> m b) -> m b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= b -> m b
k) b -> m b
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 (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 (All -> Bool) -> (t f z -> All) -> t f z -> Bool
forall (p :: * -> * -> *) a b c (q :: * -> * -> *).
(Profunctor p, Coercible c b) =>
q b c -> p a b -> p a c
#. (forall (x :: k). IndexF (t f z) x -> f x -> All) -> t f z -> All
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 (IndexF (t f z) x -> f x -> Bool
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 (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 (Any -> Bool) -> (t f z -> Any) -> t f z -> Bool
forall (p :: * -> * -> *) a b c (q :: * -> * -> *).
(Profunctor p, Coercible c b) =>
q b c -> p a b -> p a c
#. (forall (x :: k). IndexF (t f z) x -> f x -> Any) -> t f z -> Any
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 (IndexF (t f z) x -> f x -> Bool
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 (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 = Identity (t g z) -> t g z
forall a. Identity a -> a
runIdentity (Identity (t g z) -> t g z)
-> (t f z -> Identity (t g z)) -> t f z -> t g z
forall (p :: * -> * -> *) a b c (q :: * -> * -> *).
(Profunctor p, Coercible c b) =>
q b c -> p a b -> p a c
#. (forall (x :: k). IndexF (t f z) x -> f x -> Identity (g x))
-> t f z -> Identity (t g z)
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 -> g x -> Identity (g x)
forall a. a -> Identity a
Identity (IndexF (t f z) x -> f x -> g x
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 (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 = Const m (t Any z) -> m
forall a k (b :: k). Const a b -> a
getConst (Const m (t Any z) -> m)
-> (t f z -> Const m (t Any z)) -> t f z -> m
forall (p :: * -> * -> *) a b c (q :: * -> * -> *).
(Profunctor p, Coercible c b) =>
q b c -> p a b -> p a c
#. (forall (x :: k). IndexF (t f z) x -> f x -> Const m (Any x))
-> t f z -> Const m (t Any z)
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 -> m -> Const m (Any x)
forall k a (b :: k). a -> Const a b
Const (IndexF (t f z) x -> f x -> m
forall (x :: k). IndexF (t f z) x -> f x -> m
f IndexF (t f z) x
i f x
x))
{-# INLINEABLE ifoldMapFCDefault #-}