TraversableF.hs

------------------------------------------------------------------------
-- |
-- 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.Functor.Product (Product(Pair))
import Data.Functor.Sum (Sum(InL, InR))
import Data.Kind
import Data.Monoid hiding (Product, Sum)
import Data.Proxy (Proxy(Proxy))
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 _ = coerce

instance (FunctorF f, FunctorF g) => FunctorF (Product f g) where
  fmapF f (Pair x y) = Pair (fmapF f x) (fmapF f y)

instance FunctorF Proxy where
  fmapF _ = coerce
  {-# INLINE fmapF #-}

instance (FunctorF f, FunctorF g) => FunctorF (Sum f g) where
  fmapF f (InL x) = InL (fmapF f x)
  fmapF f (InR x) = InR (fmapF f x)

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

-- | This is a coercion used to avoid overhead associated
-- with function composition.
(#.) :: Coercible b c => (b -> c) -> (a -> b) -> (a -> c)
(#.) _f = 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 f = foldrF (mappend . f) mempty

  -- | Right-associative fold of a structure.
  foldrF :: (forall s . e s -> b -> b) -> b -> t e -> b
  foldrF f z t = appEndo (foldMapF (Endo #. f) t) z

  -- | Left-associative fold of a structure.
  foldlF :: (forall s . b -> e s -> b) -> b -> t e -> b
  foldlF f z t = appEndo (getDual (foldMapF (\e -> Dual (Endo (\r -> f r e))) t)) 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' f0 z0 xs = foldlF (f' f0) id xs z0
    where f' f k x z = k $! f x z

  -- | Left-associative fold of a parameterized structure
  -- with a strict accumulator.
  foldlF' :: (forall s . b -> e s -> b) -> b -> t e -> b
  foldlF' f0 z0 xs = foldrF (f' f0) id xs z0
    where f' f x k z = k $! f z x

  -- | Convert structure to list.
  toListF :: (forall tp . f tp -> a) -> t f -> [a]
  toListF f t = build (\c n -> foldrF (\e v -> c (f e) v) n 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 f z0 xs = foldrF f' return xs z0
  where f' x k z = f z x >>= 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' f z0 xs = seq z0 (foldrF f' return xs z0)
  where f' x k z = f z x >>= \r -> seq r (k 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 f z0 xs = foldlF f' return xs z0
  where f' k x z = f x z >>= 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' f z0 xs = seq z0 $ foldlF f' return xs z0
  where f' k x z = f x z >>= \r -> seq r (k r)

-- | Return 'True' if all values satisfy the predicate.
allF :: FoldableF t => (forall tp . f tp -> Bool) -> t f -> Bool
allF p = getAll #. foldMapF (All #. p)

-- | Return 'True' if any values satisfy the predicate.
anyF :: FoldableF t => (forall tp . f tp -> Bool) -> t f -> Bool
anyF p = getAny #. foldMapF (Any #. p)

-- | Return number of elements that we fold over.
lengthF :: FoldableF t => t f -> Int
lengthF = foldrF (const (+1)) 0

instance FoldableF (Const x) where
  foldMapF _ _ = mempty

instance (FoldableF f, FoldableF g) => FoldableF (Product f g) where
  foldMapF f (Pair x y) = foldMapF f x <> foldMapF f y

instance FoldableF Proxy where
  foldMapF _ _ = mempty
  {-# INLINE foldMapF #-}

instance (FoldableF f, FoldableF g) => FoldableF (Sum f g) where
  foldMapF f (InL x) = foldMapF f x
  foldMapF f (InR y) = foldMapF f y

------------------------------------------------------------------------
-- 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 _ (Const x) = pure (Const x)

instance (TraversableF f, TraversableF g) => TraversableF (Product f g) where
  traverseF f (Pair x y) = Pair <$> traverseF f x <*> traverseF f y

instance TraversableF Proxy where
  traverseF _ _ = pure Proxy
  {-# INLINE traverseF #-}

instance (TraversableF f, TraversableF g) => TraversableF (Sum f g) where
  traverseF f (InL x) = InL <$> traverseF f x
  traverseF f (InR y) = InR <$> traverseF f y

-- | Flipped 'traverseF'
forF :: (TraversableF t, Applicative m) => t e -> (forall s . e s -> m (f s)) -> m (t f)
forF f x = traverseF x 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 f = runIdentity #. traverseF (Identity #. 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 f = getConst #. traverseF (Const #. 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_ f = foldrF (\e r -> f e *> r) (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_ v f = traverseF_ 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 f (Compose v) = Compose $ fmapF (fmapFC f) v

instance ( TraversableF (s :: (k -> Type) -> Type)
         , TraversableFC (t :: (l -> Type) -> (k -> Type))
         ) =>
         FoldableF (Compose s t) where
  foldMapF = 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 f (Compose v) = Compose <$> traverseF (traverseFC f) v