------------------------------------------------------------------------
-- |
-- Module           : Data.Parameterized.TraversableF
-- Copyright        : (c) Galois, Inc 2014-2019
-- Maintainer       : Joe Hendrix <jhendrix@galois.com>
-- Description      : Traversing structures having a single parametric type
--
-- This module declares classes for working with structures that accept
-- a single parametric type parameter.
------------------------------------------------------------------------
{-# LANGUAGE InstanceSigs #-}
{-# LANGUAGE PolyKinds #-}
{-# LANGUAGE RankNTypes #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE Trustworthy #-}
module Data.Parameterized.TraversableF
  ( FunctorF(..)
  , FoldableF(..)
  , foldlMF
  , foldlMF'
  , foldrMF
  , foldrMF'
  , TraversableF(..)
  , traverseF_
  , forF_
  , forF
  , fmapFDefault
  , foldMapFDefault
  , allF
  , anyF
  , lengthF
  ) where

import Control.Applicative
import Control.Monad.Identity
import Data.Coerce
import Data.Functor.Compose (Compose(..))
import Data.Kind
import Data.Monoid
import GHC.Exts (build)

import Data.Parameterized.TraversableFC

-- | A parameterized type that is a functor on all instances.
class FunctorF m where
  fmapF :: (forall x . f x -> g x) -> m f -> m g

instance FunctorF (Const x) where
  fmapF :: forall (f :: k -> *) (g :: k -> *).
(forall (x :: k). f x -> g x) -> Const x f -> Const x g
fmapF forall (x :: k). f x -> g x
_ = coerce :: forall a b. Coercible a b => a -> b
coerce

------------------------------------------------------------------------
-- FoldableF

-- | This is a coercion used to avoid overhead associated
-- with function composition.
(#.) :: Coercible b c => (b -> c) -> (a -> b) -> (a -> c)
#. :: forall b c a. Coercible b c => (b -> c) -> (a -> b) -> a -> c
(#.) b -> c
_f = coerce :: forall a b. Coercible a b => a -> b
coerce

-- | This is a generalization of the 'Foldable' class to
-- structures over parameterized terms.
class FoldableF (t :: (k -> Type) -> Type) where
  {-# MINIMAL foldMapF | foldrF #-}

  -- | Map each element of the structure to a monoid,
  -- and combine the results.
  foldMapF :: Monoid m => (forall s . e s -> m) -> t e -> m
  foldMapF forall (s :: k). e s -> m
f = forall k (t :: (k -> *) -> *) (e :: k -> *) b.
FoldableF t =>
(forall (s :: k). e s -> b -> b) -> b -> t e -> b
foldrF (forall a. Monoid a => a -> a -> a
mappend forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (s :: k). e s -> m
f) forall a. Monoid a => a
mempty

  -- | Right-associative fold of a structure.
  foldrF :: (forall s . e s -> b -> b) -> b -> t e -> b
  foldrF forall (s :: k). e s -> b -> b
f b
z t e
t = forall a. Endo a -> a -> a
appEndo (forall k (t :: (k -> *) -> *) m (e :: k -> *).
(FoldableF t, Monoid m) =>
(forall (s :: k). e s -> m) -> t e -> m
foldMapF (forall a. (a -> a) -> Endo a
Endo forall b c a. Coercible b c => (b -> c) -> (a -> b) -> a -> c
#. forall (s :: k). e s -> b -> b
f) t e
t) b
z

  -- | Left-associative fold of a structure.
  foldlF :: (forall s . b -> e s -> b) -> b -> t e -> b
  foldlF forall (s :: k). b -> e s -> b
f b
z t e
t = forall a. Endo a -> a -> a
appEndo (forall a. Dual a -> a
getDual (forall k (t :: (k -> *) -> *) m (e :: k -> *).
(FoldableF t, Monoid m) =>
(forall (s :: k). e s -> m) -> t e -> m
foldMapF (\e s
e -> forall a. a -> Dual a
Dual (forall a. (a -> a) -> Endo a
Endo (\b
r -> forall (s :: k). b -> e s -> b
f b
r e s
e))) t e
t)) b
z

  -- | Right-associative fold of a structure,
  -- but with strict application of the operator.
  foldrF' :: (forall s . e s -> b -> b) -> b -> t e -> b
  foldrF' forall (s :: k). e s -> b -> b
f0 b
z0 t e
xs = forall k (t :: (k -> *) -> *) b (e :: k -> *).
FoldableF t =>
(forall (s :: k). b -> e s -> b) -> b -> t e -> b
foldlF (forall {t} {t} {a} {b}. (t -> t -> a) -> (a -> b) -> t -> t -> b
f' forall (s :: k). e s -> b -> b
f0) forall a. a -> a
id t e
xs b
z0
    where f' :: (t -> t -> a) -> (a -> b) -> t -> t -> b
f' t -> t -> a
f a -> b
k t
x t
z = a -> b
k forall a b. (a -> b) -> a -> b
$! t -> t -> a
f t
x t
z

  -- | Left-associative fold of a parameterized structure
  -- with a strict accumulator.
  foldlF' :: (forall s . b -> e s -> b) -> b -> t e -> b
  foldlF' forall (s :: k). b -> e s -> b
f0 b
z0 t e
xs = forall k (t :: (k -> *) -> *) (e :: k -> *) b.
FoldableF t =>
(forall (s :: k). e s -> b -> b) -> b -> t e -> b
foldrF (forall {t} {t} {a} {b}. (t -> t -> a) -> t -> (a -> b) -> t -> b
f' forall (s :: k). b -> e s -> b
f0) forall a. a -> a
id t e
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.
  toListF :: (forall tp . f tp -> a) -> t f -> [a]
  toListF forall (tp :: k). f tp -> a
f t f
t = forall a. (forall b. (a -> b -> b) -> b -> b) -> [a]
build (\a -> b -> b
c b
n -> forall k (t :: (k -> *) -> *) (e :: k -> *) b.
FoldableF t =>
(forall (s :: k). e s -> b -> b) -> b -> t e -> b
foldrF (\f s
e b
v -> a -> b -> b
c (forall (tp :: k). f tp -> a
f f s
e) b
v) b
n t f
t)

-- | Monadic fold over the elements of a structure from left to right.
foldlMF :: (FoldableF t, Monad m) => (forall x . b -> f x -> m b) -> b -> t f -> m b
foldlMF :: forall {k} (t :: (k -> *) -> *) (m :: * -> *) b (f :: k -> *).
(FoldableF t, Monad m) =>
(forall (x :: k). b -> f x -> m b) -> b -> t f -> m b
foldlMF forall (x :: k). b -> f x -> m b
f b
z0 t f
xs = forall k (t :: (k -> *) -> *) (e :: k -> *) b.
FoldableF t =>
(forall (s :: k). e s -> b -> b) -> b -> t e -> b
foldrF forall {x :: k} {b}. f x -> (b -> m b) -> b -> m b
f' forall (m :: * -> *) a. Monad m => a -> m a
return t f
xs b
z0
  where f' :: f x -> (b -> m b) -> b -> m b
f' f x
x b -> m b
k b
z = forall (x :: k). b -> f x -> m b
f b
z f x
x forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= b -> m b
k

-- | Monadic strict fold over the elements of a structure from left to right.
foldlMF' :: (FoldableF t, Monad m) => (forall x . b -> f x -> m b) -> b -> t f -> m b
foldlMF' :: forall {k} (t :: (k -> *) -> *) (m :: * -> *) b (f :: k -> *).
(FoldableF t, Monad m) =>
(forall (x :: k). b -> f x -> m b) -> b -> t f -> m b
foldlMF' forall (x :: k). b -> f x -> m b
f b
z0 t f
xs = seq :: forall a b. a -> b -> b
seq b
z0 (forall k (t :: (k -> *) -> *) (e :: k -> *) b.
FoldableF t =>
(forall (s :: k). e s -> b -> b) -> b -> t e -> b
foldrF forall {x :: k} {b}. f x -> (b -> m b) -> b -> m b
f' forall (m :: * -> *) a. Monad m => a -> m a
return t f
xs b
z0)
  where f' :: f x -> (b -> m b) -> b -> m b
f' f x
x b -> m b
k b
z = forall (x :: k). b -> f x -> m b
f b
z f x
x forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= \b
r -> seq :: forall a b. a -> b -> b
seq b
r (b -> m b
k b
r)

-- | Monadic fold over the elements of a structure from right to left.
foldrMF :: (FoldableF t, Monad m) => (forall x . f x -> b -> m b) -> b -> t f -> m b
foldrMF :: forall {k} (t :: (k -> *) -> *) (m :: * -> *) (f :: k -> *) b.
(FoldableF t, Monad m) =>
(forall (x :: k). f x -> b -> m b) -> b -> t f -> m b
foldrMF forall (x :: k). f x -> b -> m b
f b
z0 t f
xs = forall k (t :: (k -> *) -> *) b (e :: k -> *).
FoldableF t =>
(forall (s :: k). b -> e s -> b) -> b -> t e -> b
foldlF forall {b} {x :: k}. (b -> m b) -> f x -> b -> m b
f' forall (m :: * -> *) a. Monad m => a -> m a
return t f
xs b
z0
  where f' :: (b -> m b) -> f x -> b -> m b
f' b -> m b
k f x
x b
z = forall (x :: k). f x -> b -> m b
f f x
x b
z forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= b -> m b
k

-- | Monadic strict fold over the elements of a structure from right to left.
foldrMF' :: (FoldableF t, Monad m) => (forall x . f x -> b -> m b) -> b -> t f -> m b
foldrMF' :: forall {k} (t :: (k -> *) -> *) (m :: * -> *) (f :: k -> *) b.
(FoldableF t, Monad m) =>
(forall (x :: k). f x -> b -> m b) -> b -> t f -> m b
foldrMF' forall (x :: k). f x -> b -> m b
f b
z0 t f
xs = seq :: forall a b. a -> b -> b
seq b
z0 forall a b. (a -> b) -> a -> b
$ forall k (t :: (k -> *) -> *) b (e :: k -> *).
FoldableF t =>
(forall (s :: k). b -> e s -> b) -> b -> t e -> b
foldlF forall {b} {x :: k}. (b -> m b) -> f x -> b -> m b
f' forall (m :: * -> *) a. Monad m => a -> m a
return t f
xs b
z0
  where f' :: (b -> m b) -> f x -> b -> m b
f' b -> m b
k f x
x b
z = forall (x :: k). f x -> b -> m b
f f x
x b
z forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= \b
r -> seq :: forall a b. a -> b -> b
seq b
r (b -> m b
k b
r)

-- | Return 'True' if all values satisfy the predicate.
allF :: FoldableF t => (forall tp . f tp -> Bool) -> t f -> Bool
allF :: forall {k} (t :: (k -> *) -> *) (f :: k -> *).
FoldableF t =>
(forall (tp :: k). f tp -> Bool) -> t f -> Bool
allF forall (tp :: k). f tp -> Bool
p = All -> Bool
getAll forall b c a. Coercible b c => (b -> c) -> (a -> b) -> a -> c
#. forall k (t :: (k -> *) -> *) m (e :: k -> *).
(FoldableF t, Monoid m) =>
(forall (s :: k). e s -> m) -> t e -> m
foldMapF (Bool -> All
All forall b c a. Coercible b c => (b -> c) -> (a -> b) -> a -> c
#. forall (tp :: k). f tp -> Bool
p)

-- | Return 'True' if any values satisfy the predicate.
anyF :: FoldableF t => (forall tp . f tp -> Bool) -> t f -> Bool
anyF :: forall {k} (t :: (k -> *) -> *) (f :: k -> *).
FoldableF t =>
(forall (tp :: k). f tp -> Bool) -> t f -> Bool
anyF forall (tp :: k). f tp -> Bool
p = Any -> Bool
getAny forall b c a. Coercible b c => (b -> c) -> (a -> b) -> a -> c
#. forall k (t :: (k -> *) -> *) m (e :: k -> *).
(FoldableF t, Monoid m) =>
(forall (s :: k). e s -> m) -> t e -> m
foldMapF (Bool -> Any
Any forall b c a. Coercible b c => (b -> c) -> (a -> b) -> a -> c
#. forall (tp :: k). f tp -> Bool
p)

-- | Return number of elements that we fold over.
lengthF :: FoldableF t => t f -> Int
lengthF :: forall {k} (t :: (k -> *) -> *) (f :: k -> *).
FoldableF t =>
t f -> Int
lengthF = forall k (t :: (k -> *) -> *) (e :: k -> *) b.
FoldableF t =>
(forall (s :: k). e s -> b -> b) -> b -> t e -> b
foldrF (forall a b. a -> b -> a
const (forall a. Num a => a -> a -> a
+Int
1)) Int
0

instance FoldableF (Const x) where
  foldMapF :: forall m (e :: k -> *).
Monoid m =>
(forall (s :: k). e s -> m) -> Const x e -> m
foldMapF forall (s :: k). e s -> m
_ Const x e
_ = forall a. Monoid a => a
mempty

------------------------------------------------------------------------
-- TraversableF

class (FunctorF t, FoldableF t) => TraversableF t where
  traverseF :: Applicative m
            => (forall s . e s -> m (f s))
            -> t e
            -> m (t f)

instance TraversableF (Const x) where
  traverseF :: forall (m :: * -> *) (e :: k -> *) (f :: k -> *).
Applicative m =>
(forall (s :: k). e s -> m (f s)) -> Const x e -> m (Const x f)
traverseF forall (s :: k). e s -> m (f s)
_ (Const x
x) = forall (f :: * -> *) a. Applicative f => a -> f a
pure (forall {k} a (b :: k). a -> Const a b
Const x
x)

-- | Flipped 'traverseF'
forF :: (TraversableF t, Applicative m) => t e -> (forall s . e s -> m (f s)) -> m (t f)
forF :: forall {k} (t :: (k -> *) -> *) (m :: * -> *) (e :: k -> *)
       (f :: k -> *).
(TraversableF t, Applicative m) =>
t e -> (forall (s :: k). e s -> m (f s)) -> m (t f)
forF t e
f forall (s :: k). e s -> m (f s)
x = forall {k} (t :: (k -> *) -> *) (m :: * -> *) (e :: k -> *)
       (f :: k -> *).
(TraversableF t, Applicative m) =>
(forall (s :: k). e s -> m (f s)) -> t e -> m (t f)
traverseF forall (s :: k). e s -> m (f s)
x t e
f
{-# INLINE forF #-}

-- | This function may be used as a value for `fmapF` in a `FunctorF`
-- instance.
fmapFDefault :: TraversableF t => (forall s . e s -> f s) -> t e -> t f
fmapFDefault :: forall {k} (t :: (k -> *) -> *) (e :: k -> *) (f :: k -> *).
TraversableF t =>
(forall (s :: k). e s -> f s) -> t e -> t f
fmapFDefault forall (s :: k). e s -> f s
f = forall a. Identity a -> a
runIdentity forall b c a. Coercible b c => (b -> c) -> (a -> b) -> a -> c
#. forall {k} (t :: (k -> *) -> *) (m :: * -> *) (e :: k -> *)
       (f :: k -> *).
(TraversableF t, Applicative m) =>
(forall (s :: k). e s -> m (f s)) -> t e -> m (t f)
traverseF (forall a. a -> Identity a
Identity forall b c a. Coercible b c => (b -> c) -> (a -> b) -> a -> c
#. forall (s :: k). e s -> f s
f)
{-# INLINE fmapFDefault #-}

-- | This function may be used as a value for `Data.Foldable.foldMap`
-- in a `Foldable` instance.
foldMapFDefault :: (TraversableF t, Monoid m) => (forall s . e s -> m) -> t e -> m
foldMapFDefault :: forall {k} (t :: (k -> *) -> *) m (e :: k -> *).
(TraversableF t, Monoid m) =>
(forall (s :: k). e s -> m) -> t e -> m
foldMapFDefault forall (s :: k). e s -> m
f = forall {k} a (b :: k). Const a b -> a
getConst forall b c a. Coercible b c => (b -> c) -> (a -> b) -> a -> c
#. forall {k} (t :: (k -> *) -> *) (m :: * -> *) (e :: k -> *)
       (f :: k -> *).
(TraversableF t, Applicative m) =>
(forall (s :: k). e s -> m (f s)) -> t e -> m (t f)
traverseF (forall {k} a (b :: k). a -> Const a b
Const forall b c a. Coercible b c => (b -> c) -> (a -> b) -> a -> c
#. forall (s :: k). e s -> m
f)

-- | Map each element of a structure to an action, evaluate
-- these actions from left to right, and ignore the results.
traverseF_ :: (FoldableF t, Applicative f) => (forall s . e s  -> f a) -> t e -> f ()
traverseF_ :: forall {k} (t :: (k -> *) -> *) (f :: * -> *) (e :: k -> *) a.
(FoldableF t, Applicative f) =>
(forall (s :: k). e s -> f a) -> t e -> f ()
traverseF_ forall (s :: k). e s -> f a
f = forall k (t :: (k -> *) -> *) (e :: k -> *) b.
FoldableF t =>
(forall (s :: k). e s -> b -> b) -> b -> t e -> b
foldrF (\e s
e f ()
r -> forall (s :: k). e s -> f a
f e s
e forall (f :: * -> *) a b. Applicative f => f a -> f b -> f b
*> f ()
r) (forall (f :: * -> *) a. Applicative f => a -> f a
pure ())


-- | Map each element of a structure to an action, evaluate
-- these actions from left to right, and ignore the results.
forF_ :: (FoldableF t, Applicative m) => t f -> (forall x. f x -> m a) -> m ()
forF_ :: forall {k} (t :: (k -> *) -> *) (m :: * -> *) (f :: k -> *) a.
(FoldableF t, Applicative m) =>
t f -> (forall (x :: k). f x -> m a) -> m ()
forF_ t f
v forall (x :: k). f x -> m a
f = forall {k} (t :: (k -> *) -> *) (f :: * -> *) (e :: k -> *) a.
(FoldableF t, Applicative f) =>
(forall (s :: k). e s -> f a) -> t e -> f ()
traverseF_ forall (x :: k). f x -> m a
f t f
v
{-# INLINE forF_ #-}

------------------------------------------------------------------------
-- TraversableF (Compose s t)

instance ( FunctorF (s :: (k -> Type) -> Type)
         , FunctorFC (t :: (l -> Type) -> (k -> Type))
         ) =>
         FunctorF (Compose s t) where
  fmapF :: forall (f :: l -> *) (g :: l -> *).
(forall (x :: l). f x -> g x) -> Compose s t f -> Compose s t g
fmapF forall (x :: l). f x -> g x
f (Compose s (t f)
v) = forall {k} {k1} (f :: k -> *) (g :: k1 -> k) (a :: k1).
f (g a) -> Compose f g a
Compose forall a b. (a -> b) -> a -> b
$ forall {k} (m :: (k -> *) -> *) (f :: k -> *) (g :: k -> *).
FunctorF m =>
(forall (x :: k). f x -> g x) -> m f -> m g
fmapF (forall k l (t :: (k -> *) -> l -> *) (f :: k -> *) (g :: k -> *).
FunctorFC t =>
(forall (x :: k). f x -> g x) -> forall (x :: l). t f x -> t g x
fmapFC forall (x :: l). f x -> g x
f) s (t f)
v

instance ( TraversableF (s :: (k -> Type) -> Type)
         , TraversableFC (t :: (l -> Type) -> (k -> Type))
         ) =>
         FoldableF (Compose s t) where
  foldMapF :: forall m (e :: l -> *).
Monoid m =>
(forall (s :: l). e s -> m) -> Compose s t e -> m
foldMapF = forall {k} (t :: (k -> *) -> *) m (e :: k -> *).
(TraversableF t, Monoid m) =>
(forall (s :: k). e s -> m) -> t e -> m
foldMapFDefault

-- | Traverse twice over: go under the @t@, under the @s@ and lift @m@ out.
instance ( TraversableF (s :: (k -> Type) -> Type)
         , TraversableFC (t :: (l -> Type) -> (k -> Type))
         ) =>
         TraversableF (Compose s t) where
  traverseF :: forall (f :: l -> Type) (g :: l -> Type) m. (Applicative m) =>
               (forall (u :: l). f u -> m (g u))
            -> Compose s t f -> m (Compose s t g)
  traverseF :: forall (f :: l -> *) (g :: l -> *) (m :: * -> *).
Applicative m =>
(forall (u :: l). f u -> m (g u))
-> Compose s t f -> m (Compose s t g)
traverseF forall (u :: l). f u -> m (g u)
f (Compose s (t f)
v) = forall {k} {k1} (f :: k -> *) (g :: k1 -> k) (a :: k1).
f (g a) -> Compose f g a
Compose forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall {k} (t :: (k -> *) -> *) (m :: * -> *) (e :: k -> *)
       (f :: k -> *).
(TraversableF t, Applicative m) =>
(forall (s :: k). e s -> m (f s)) -> t e -> m (t f)
traverseF (forall k l (t :: (k -> *) -> l -> *) (f :: k -> *) (g :: k -> *)
       (m :: * -> *).
(TraversableFC t, Applicative m) =>
(forall (x :: k). f x -> m (g x))
-> forall (x :: l). t f x -> m (t g x)
traverseFC forall (u :: l). f u -> m (g u)
f) s (t f)
v