{-# LANGUAGE FlexibleContexts    #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE Rank2Types          #-}

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

module Generics.MultiRec.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.MultiRec.Rewriting.Rewriting

import Generics.MultiRec.Base
import Generics.MultiRec.HFunctor
import Generics.MultiRec.Compos (composM, compos)


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

-- | Applies a function to the first subtree (possibly the tree itself)
-- on which it succeeds, using a preorder traversal.
once :: (Fam phi, HFunctor phi (PF phi), Functor m, MonadPlus m)
     => (forall ix'. phi ix' -> ix' -> m ix') -> phi ix -> ix -> m ix
once f p x = f p x `mplus` one (once f) p x

-- | Applies a function to the first immediate child of a value on which
-- it succeeds.
one :: (Fam phi, HFunctor phi (PF phi), Functor m, MonadPlus m)
    => (forall ix'. phi ix' -> ix' -> m ix') -> phi ix -> ix -> m ix
one f p x = fmap (to p) rs
  where 
    S _ rs = hmapM try p (from p x)
    try s (I0 x') = liftM I0 (S x' (f s 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.
-----------------------------------------------------------------------------

-- | Applies a monadic function exhaustively in a top-down fashion.
topdownM :: (Fam phi, HFunctor phi (PF phi), Functor m, Monad m)
         => (forall ix'. phi ix' -> ix' -> m ix') -> phi ix -> ix -> m ix
topdownM f p x = f p x >>= composM (topdownM f) p

-- | Applies a function exhaustively in a top-down fashion
topdown :: (Fam phi, HFunctor phi (PF phi))
        => (forall ix'. phi ix' -> ix' -> ix') -> phi ix -> ix -> ix
topdown f p x = compos (topdown f) p (f p x)


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

-- | Applies a monadic function exhaustively in a bottom-up fashion.
bottomupM :: (Fam phi, HFunctor phi (PF phi), Functor m, Monad m)
          => (forall ix'. phi ix' -> ix' -> m ix') -> phi ix -> ix -> m ix
bottomupM f p x = composM (bottomupM f) p x >>= f p

-- | Applies a function exhaustively in a bottom-up fashion
bottomup :: (Fam phi, HFunctor phi (PF phi))
          => (forall ix'. phi ix' -> ix' -> ix') -> phi ix -> ix -> ix
bottomup f p x = f p (compos (bottomup f) p x)