{-# LANGUAGE FlexibleContexts #-}

-----------------------------------------------------------------------------
-- |
-- Module      :  Generics.Regular.Rewriting.Strategies
-- Copyright   :  (c) 2008 Universiteit Utrecht
-- License     :  BSD3
--
-- Maintainer  :  generics@haskell.org
-- Stability   :  experimental
-- Portability :  non-portable
--
-- Summary: Generic functions for traversal strategies.
-----------------------------------------------------------------------------

module Generics.Regular.Rewriting.Strategies (

  -- * Apply a function to the children of a value.
  once,
  one,

  -- * Apply a (monadic) function exhaustively top-down.
  topdownM,
  topdown,

  -- * Apply a (monadic) function exhaustively bottom-up.
  bottomupM,
  bottomup,

  -- * Apply a (monadic) function to immediate children.
  composM,
  compos

) where

import Control.Monad

import Generics.Regular.Base (Regular(..), PF)
import Generics.Regular.Functions


-----------------------------------------------------------------------------
-- Functions to apply a function to the children of a value.
-----------------------------------------------------------------------------

{-# INLINE once #-}
-- | Applies a function to the first subtree (possibly the tree itself) on which
-- it succeeds, using a preorder traversal.
once :: (Regular a, GMap (PF a), Functor m, MonadPlus m) => (a -> m a) -> a -> m a
once f x = f x `mplus` one (once f) x

{-# INLINE one #-}
-- | Applies a function to the first immediate child of a value on which it succeeds.
one :: (Regular a, GMap (PF a), Functor m, MonadPlus m) => (a -> m a) -> a -> m a
one f x = fmap to rs
  where 
    S _ rs = fmapM try (from x)
    try x' = S x' (f x')

-- | Same monad to that in the SYB3 paper. It is used as follows: the first 
-- argument contains the original value, and the second arguments contain 
-- the transformed values.
data S m a = S a (m a)

instance MonadPlus m => Monad (S m) where
  return x = S x mzero
  (S x xs) >>= k = 
    S r (rs2 `mplus` rs1)
    where 
      S r rs1 = k x
      rs2     = 
        do x' <- xs
           let S r' _ = k x'
           return r'


-----------------------------------------------------------------------------
-- Apply a (monadic) function exhaustively top-down.
-----------------------------------------------------------------------------

{-# INLINE topdownM #-}
-- | Applies a monadic function exhaustively in a top-down fashion.
topdownM :: (Regular a, GMap (PF a), Functor m, Monad m) => (a -> m a) -> a -> m a
topdownM f x = f x >>= composM (topdownM f)

{-# INLINE topdown #-}
-- | Applies a function exhaustively in a top-down fashion
topdown :: (Regular a, Functor (PF a)) => (a -> a) -> a -> a
topdown f x = compos (topdown f) (f x)


-----------------------------------------------------------------------------
-- Apply a (monadic) function exhaustively bottom-up.
-----------------------------------------------------------------------------

{-# INLINE bottomupM #-}
-- | Applies a monadic function exhaustively in a bottom-up fashion.
bottomupM :: (Regular a, GMap (PF a), Functor m, Monad m) => (a -> m a) -> a -> m a
bottomupM f x = composM (bottomupM f) x >>= f

{-# INLINE bottomup #-}
-- | Applies a function exhaustively in a bottom-up fashion
bottomup :: (Regular a, Functor (PF a)) => (a -> a) -> a -> a
bottomup f x = f (compos (bottomup f) x)


-----------------------------------------------------------------------------
-- Apply a (monadic) function to immediate children.
-----------------------------------------------------------------------------

{-# INLINE composM #-}
-- | Applies a monadic function to all the immediate children of a value.
composM :: (Regular a, GMap (PF a), Functor m, Monad m) => (a -> m a) -> a -> m a
composM f = fmap to . fmapM f . from

{-# INLINE compos #-}
-- | Applies a function to all the immediate children of a value.
compos :: (Regular a, Functor (PF a)) => (a -> a) -> a -> a
compos f = to . fmap f . from