------------------------------------------------------------------------
-- |
-- 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.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 (x :: k). f x -> g x) -> Const x f -> Const x g
fmapF forall (x :: k). f x -> g x
_ = Const x f -> Const x g
coerce

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

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

-- | This is a generalization of the 'Foldable' class to
-- structures over parameterized terms.
class FoldableF (t :: (k -> *) -> *) 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 (s :: k). e s -> m -> m) -> m -> t e -> m
forall k (t :: (k -> *) -> *) (e :: k -> *) b.
FoldableF t =>
(forall (s :: k). e s -> b -> b) -> b -> t e -> b
foldrF (m -> m -> m
forall a. Monoid a => a -> a -> a
mappend (m -> m -> m) -> (e s -> m) -> e s -> m -> m
forall b c a. (b -> c) -> (a -> b) -> a -> c
. e s -> m
forall (s :: k). e s -> m
f) m
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 = Endo b -> b -> b
forall a. Endo a -> a -> a
appEndo ((forall (s :: k). e s -> Endo b) -> t e -> Endo b
forall k (t :: (k -> *) -> *) m (e :: k -> *).
(FoldableF t, Monoid m) =>
(forall (s :: k). e s -> m) -> t e -> m
foldMapF ((b -> b) -> Endo b
forall a. (a -> a) -> Endo a
Endo ((b -> b) -> Endo b) -> (e s -> b -> b) -> e s -> Endo b
forall b c a. Coercible b c => (b -> c) -> (a -> b) -> a -> c
#. e s -> b -> b
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 = Endo b -> b -> b
forall a. Endo a -> a -> a
appEndo (Dual (Endo b) -> Endo b
forall a. Dual a -> a
getDual ((forall (s :: k). e s -> Dual (Endo b)) -> t e -> Dual (Endo b)
forall k (t :: (k -> *) -> *) m (e :: k -> *).
(FoldableF t, Monoid m) =>
(forall (s :: k). e s -> m) -> t e -> m
foldMapF (\e s
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 -> b -> e s -> b
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 (s :: k). (b -> b) -> e s -> b -> b)
-> (b -> b) -> t e -> b -> b
forall k (t :: (k -> *) -> *) b (e :: k -> *).
FoldableF t =>
(forall (s :: k). b -> e s -> b) -> b -> t e -> b
foldlF ((e s -> b -> b) -> (b -> b) -> e s -> b -> b
forall t t a b. (t -> t -> a) -> (a -> b) -> t -> t -> b
f' e s -> b -> b
forall (s :: k). e s -> b -> b
f0) b -> b
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 (a -> b) -> a -> b
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 (s :: k). e s -> (b -> b) -> b -> b)
-> (b -> b) -> t e -> b -> b
forall k (t :: (k -> *) -> *) (e :: k -> *) b.
FoldableF t =>
(forall (s :: k). e s -> b -> b) -> b -> t e -> b
foldrF ((b -> e s -> b) -> e s -> (b -> b) -> b -> b
forall t t a b. (t -> t -> a) -> t -> (a -> b) -> t -> b
f' b -> e s -> b
forall (s :: k). b -> e s -> b
f0) b -> b
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 (a -> b) -> a -> b
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 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 (s :: k). f s -> b -> b) -> b -> t f -> b
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 (f s -> a
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 (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 (s :: k). f s -> (b -> m b) -> b -> m b)
-> (b -> m b) -> t f -> b -> m b
forall k (t :: (k -> *) -> *) (e :: k -> *) b.
FoldableF t =>
(forall (s :: k). e s -> b -> b) -> b -> t e -> b
foldrF forall (s :: k). f s -> (b -> m b) -> b -> m b
forall (x :: k) b. f x -> (b -> m b) -> b -> m b
f' b -> m b
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 = b -> f x -> m b
forall (x :: k). b -> f x -> m b
f 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

-- | 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 (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 = b -> m b -> m b
seq b
z0 ((forall (s :: k). f s -> (b -> m b) -> b -> m b)
-> (b -> m b) -> t f -> b -> m b
forall k (t :: (k -> *) -> *) (e :: k -> *) b.
FoldableF t =>
(forall (s :: k). e s -> b -> b) -> b -> t e -> b
foldrF forall (s :: k). f s -> (b -> m b) -> b -> m b
forall (x :: k) b. f x -> (b -> m b) -> b -> m b
f' b -> m b
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 = b -> f x -> m b
forall (x :: k). b -> f x -> m b
f 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
r -> b -> m b -> m 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 (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 (s :: k). (b -> m b) -> f s -> b -> m b)
-> (b -> m b) -> t f -> b -> m b
forall k (t :: (k -> *) -> *) b (e :: k -> *).
FoldableF t =>
(forall (s :: k). b -> e s -> b) -> b -> t e -> b
foldlF forall (s :: k). (b -> m b) -> f s -> b -> m b
forall b (x :: k). (b -> m b) -> f x -> b -> m b
f' b -> m b
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 = f x -> b -> m b
forall (x :: k). f x -> b -> m b
f 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

-- | 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 (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 = b -> m b -> m b
seq b
z0 (m b -> m b) -> m b -> m b
forall a b. (a -> b) -> a -> b
$ (forall (s :: k). (b -> m b) -> f s -> b -> m b)
-> (b -> m b) -> t f -> b -> m b
forall k (t :: (k -> *) -> *) b (e :: k -> *).
FoldableF t =>
(forall (s :: k). b -> e s -> b) -> b -> t e -> b
foldlF forall (s :: k). (b -> m b) -> f s -> b -> m b
forall b (x :: k). (b -> m b) -> f x -> b -> m b
f' b -> m b
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 = f x -> b -> m b
forall (x :: k). f x -> b -> m b
f 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
r -> b -> m b -> m 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 (tp :: k). f tp -> Bool) -> t f -> Bool
allF forall (tp :: k). f tp -> Bool
p = All -> Bool
getAll (All -> Bool) -> (t f -> All) -> t f -> Bool
forall b c a. Coercible b c => (b -> c) -> (a -> b) -> a -> c
#. (forall (s :: k). f s -> All) -> t f -> All
forall k (t :: (k -> *) -> *) m (e :: k -> *).
(FoldableF t, Monoid m) =>
(forall (s :: k). e s -> m) -> t e -> m
foldMapF (Bool -> All
All (Bool -> All) -> (f s -> Bool) -> f s -> All
forall b c a. Coercible b c => (b -> c) -> (a -> b) -> a -> c
#. f s -> Bool
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 (tp :: k). f tp -> Bool) -> t f -> Bool
anyF forall (tp :: k). f tp -> Bool
p = Any -> Bool
getAny (Any -> Bool) -> (t f -> Any) -> t f -> Bool
forall b c a. Coercible b c => (b -> c) -> (a -> b) -> a -> c
#. (forall (s :: k). f s -> Any) -> t f -> Any
forall k (t :: (k -> *) -> *) m (e :: k -> *).
(FoldableF t, Monoid m) =>
(forall (s :: k). e s -> m) -> t e -> m
foldMapF (Bool -> Any
Any (Bool -> Any) -> (f s -> Bool) -> f s -> Any
forall b c a. Coercible b c => (b -> c) -> (a -> b) -> a -> c
#. f s -> Bool
forall (tp :: k). f tp -> Bool
p)

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

instance FoldableF (Const x) where
  foldMapF :: (forall (s :: k). e s -> m) -> Const x e -> m
foldMapF forall (s :: k). e s -> m
_ Const x e
_ = m
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 (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) = Const x f -> m (Const x f)
forall (f :: * -> *) a. Applicative f => a -> f a
pure (x -> Const x f
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 :: 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 (s :: k). e s -> m (f s)) -> t e -> m (t f)
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 (s :: k). e s -> f s) -> t e -> t f
fmapFDefault forall (s :: k). e s -> f s
f = Identity (t f) -> t f
forall a. Identity a -> a
runIdentity (Identity (t f) -> t f) -> (t e -> Identity (t f)) -> t e -> t f
forall b c a. Coercible b c => (b -> c) -> (a -> b) -> a -> c
#. (forall (s :: k). e s -> Identity (f s)) -> t e -> Identity (t f)
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 (f s -> Identity (f s)
forall a. a -> Identity a
Identity (f s -> Identity (f s)) -> (e s -> f s) -> e s -> Identity (f s)
forall b c a. Coercible b c => (b -> c) -> (a -> b) -> a -> c
#. e s -> f s
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 (s :: k). e s -> m) -> t e -> m
foldMapFDefault forall (s :: k). e s -> m
f = Const m (t Any) -> m
forall a k (b :: k). Const a b -> a
getConst (Const m (t Any) -> m) -> (t e -> Const m (t Any)) -> t e -> m
forall b c a. Coercible b c => (b -> c) -> (a -> b) -> a -> c
#. (forall (s :: k). e s -> Const m (Any s)) -> t e -> Const m (t Any)
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 (m -> Const m (Any s)
forall k a (b :: k). a -> Const a b
Const (m -> Const m (Any s)) -> (e s -> m) -> e s -> Const m (Any s)
forall b c a. Coercible b c => (b -> c) -> (a -> b) -> a -> c
#. e s -> m
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 (s :: k). e s -> f a) -> t e -> f ()
traverseF_ forall (s :: k). e s -> f a
f = (forall (s :: k). e s -> f () -> f ()) -> f () -> t e -> 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 -> e s -> f a
forall (s :: k). e s -> f a
f e s
e f a -> f () -> f ()
forall (f :: * -> *) a b. Applicative f => f a -> f b -> f b
*> f ()
r) (() -> f ()
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_ :: t f -> (forall (x :: k). f x -> m a) -> m ()
forF_ t f
v forall (x :: k). f x -> m a
f = (forall (x :: k). f x -> m a) -> t f -> m ()
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 -> *) -> *)
         , FunctorFC (t :: (l -> *) -> (k -> *))
         ) =>
         FunctorF (Compose s t) where
  fmapF :: (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) = s (t g) -> Compose s t g
forall k k1 (f :: k -> *) (g :: k1 -> k) (a :: k1).
f (g a) -> Compose f g a
Compose (s (t g) -> Compose s t g) -> s (t g) -> Compose s t g
forall a b. (a -> b) -> a -> b
$ (forall (x :: k). t f x -> t g x) -> s (t f) -> s (t g)
forall k (m :: (k -> *) -> *) (f :: k -> *) (g :: k -> *).
FunctorF m =>
(forall (x :: k). f x -> g x) -> m f -> m g
fmapF ((forall (x :: l). f x -> g x) -> forall (x :: k). t f x -> t g x
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 -> *) -> *)
         , TraversableFC (t :: (l -> *) -> (k -> *))
         ) =>
         FoldableF (Compose s t) where
  foldMapF :: (forall (s :: l). e s -> m) -> Compose s t e -> m
foldMapF = (forall (s :: l). e s -> m) -> Compose s t e -> m
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 -> *) -> *)
         , TraversableFC (t :: (l -> *) -> (k -> *))
         ) =>
         TraversableF (Compose s t) where
  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))
-> Compose s t f -> m (Compose s t g)
traverseF forall (u :: l). f u -> m (g u)
f (Compose s (t f)
v) = s (t g) -> Compose s t g
forall k k1 (f :: k -> *) (g :: k1 -> k) (a :: k1).
f (g a) -> Compose f g a
Compose (s (t g) -> Compose s t g) -> m (s (t g)) -> m (Compose s t g)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (forall (s :: k). t f s -> m (t g s)) -> s (t f) -> m (s (t g))
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 (u :: l). f u -> m (g u))
-> forall (s :: k). t f s -> m (t g s)
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