{-# LANGUAGE RankNTypes                 #-}
{-# LANGUAGE FlexibleContexts           #-}
{-# LANGUAGE GADTs                      #-}

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

import Control.Monad
import Control.Monad.State

import Generics.MultiRec
import Generics.MultiRec.Zipper
import Generics.MultiRec.Any

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

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

enterMonad :: (El phi a, Fam phi, Zipper phi (PF phi))
           => a -> ZipperMonad phi I0 a (Any phi)
enterMonad x = put ([], enter proof x) >> return (Any proof x)

moveMonad :: (EqS phi, El phi a)
          => (Loc phi I0 a -> Maybe (Loc phi I0 a))
          -> ZipperMonad phi I0 a (Any phi)
moveMonad d = StateT (\(s,l) -> do l' <- d l
                                   let a = on (\p (I0 x) -> Any p x) l'
                                   return (a, (s,l')))

upMonad, downMonad, leftMonad, rightMonad :: (EqS phi, El phi a)
                                          => ZipperMonad phi I0 a (Any phi)
upMonad    = moveMonad up
downMonad  = moveMonad down
leftMonad  = moveMonad left
rightMonad = moveMonad right

updateMonad :: (EqS phi, El phi a)
            => (forall xi. phi xi -> xi -> Maybe xi) 
               -> ZipperMonad phi I0 a (Any phi)
updateMonad f = do (s,l) <- get
                   let l' = update (\p -> maybe (error "updateMonad") id . f p) l
                       a  = on (\p (I0 x) -> Any p x) l'
                   put (s,l')
                   return a
saveMonad :: (EqS phi, El phi a) => ZipperMonad phi I0 a (Any phi)
saveMonad = do (s,l) <- get
               let a = on (\p (I0 x) -> Any p x) l
               put (s++[a],l)
               return a

loadMonad :: (EqS phi, El phi a) => ZipperMonad phi I0 a (Any phi)
loadMonad = do (s:ss,l) <- get
               let l' = update (\p x -> maybe x id (matchAny p s)) l
               put (ss,l')
               return s

topMonad :: (EqS phi, El phi a) => ZipperMonad phi I0 a (Any phi)
topMonad = moveMonad goUp where
  goUp l = maybe (Just l) goUp (up l)

leaveMonad :: (EqS phi, El phi a) 
              => Loc phi I0 a -> ZipperMonad phi I0 a b -> Maybe a
leaveMonad s m = maybe Nothing (matchAny proof) $ evalStateT (m >> topMonad) ([],s)

navigate :: (Fam phi, EqS phi, El phi a, Zipper phi (PF phi))
            => phi a -> a -> ZipperMonad phi I0 a b -> Maybe a
navigate p x = leaveMonad (enter p x)