{-# LANGUAGE FlexibleContexts           #-}

module Generics.Regular.Transformations.ZipperState (
  ZipperMonad, ZipperState, upMonad, downMonad, leftMonad, rightMonad, 
  navigate, saveMonad, loadMonad, topMonad, updateMonad
  ) where

import Control.Monad.State (StateT (..), evalStateT, get, put)

import Generics.Regular.Zipper
import Generics.Regular ( Regular, PF )

--------------------------------------------------------------------------------
-- A zipper with state
--------------------------------------------------------------------------------

type ZipperState a = ([a], Loc a)
type ZipperMonad a b = StateT (ZipperState a) Maybe b

moveMonad :: (Loc a -> Maybe (Loc a)) -> ZipperMonad a a
moveMonad m = StateT (\(s,l) -> m l >>= (\l' -> return (on l', (s,l'))))

upMonad, downMonad, leftMonad, rightMonad :: ZipperMonad a a
upMonad    = moveMonad up
downMonad  = moveMonad down
leftMonad  = moveMonad left
rightMonad = moveMonad right

updateMonad :: (a -> a) -> ZipperMonad a a
updateMonad f = do (s,l) <- get
                   let l' = update f l
                   put (s,l')
                   return (on l')

saveMonad :: ZipperMonad a a
saveMonad = do (s,l) <- get
               let a = on l
               put (s++[a],l)
               return a

loadMonad :: ZipperMonad a a
loadMonad = do (s:ss,l) <- get
               let l' = update (const s) l
               put (ss,l')
               return (on l')

topMonad :: ZipperMonad a a
topMonad = do (_, Loc x l) <- get
              case l of
                [] -> return x
                _  -> upMonad >> topMonad

leaveMonad :: Loc a -> ZipperMonad a b -> Maybe a
leaveMonad s m = evalStateT (m >> topMonad) ([],s)

navigate :: (Regular a, Zipper (PF a)) => a -> ZipperMonad a b -> Maybe a
navigate x m = leaveMonad (enter x) m