{-# LANGUAGE CPP, FlexibleInstances #-}

-- | The Zipper is a data structure which maintains a location in 
-- a tree, and allows O(1) movement and local changes
-- (to be more precise, in our case it is O(k) where k is the number
-- of children of the node at question; typically this is a very small number).
--
module Data.Generics.Fixplate.Zipper where

--------------------------------------------------------------------------------

import Prelude hiding (foldl,foldr,mapM,mapM_,concat,concatMap)
import Data.Foldable
import Data.Traversable ()
import Data.Maybe

import Text.Show ()
import Text.Read 

import Data.Generics.Fixplate.Base
import Data.Generics.Fixplate.Open
import Data.Generics.Fixplate.Misc

--------------------------------------------------------------------------------
-- * Types

-- | A context node. 
type Node f  =  Either (Mu f) (Path f)

-- | The context or path type. The invariant we must respect is that there is exactly
-- one child with the 'Right' constructor.
data Path f  =  Top
             |  Path { unPath :: f (Node f) } 
               
-- | The zipper type itself, which encodes a locations in thre tree @Mu f@.
data Loc f   =  Loc { focus :: Mu f , path :: Path f } 

--------------------------------------------------------------------------------

instance EqF f => Eq (Path f) where               
  Top     == Top      = True
  Path p1 == Path p2  = equalF p1 p2
  _       == _        = False
  
instance EqF f => Eq (Loc f) where               
  Loc f1 p1 == Loc f2 p2  = f1 == f2 && p1 == p2
  
instance ShowF f => Show (Path f) where               
  showsPrec d Top = showString "Top"
  showsPrec d (Path xs) = showParen (d>10) 
    $ showString "Path "
    . showsPrecF 11 xs

instance ShowF f => Show (Loc f) where
  showsPrec d (Loc foc path) = showParen (d>10) 
    $ showString "Loc "
    . showsPrec 11 foc
    . showChar ' '
    . showsPrec 11 path

instance ReadF f => Read (Path f) where 
#ifdef __GLASGOW_HASKELL__
  readPrec = parens $ 
    (do
      { Ident "Top" <- lexP
      ; return Top            
      })
    +++
    (prec app_prec $ do
      { Ident "Path" <- lexP
      ; p <- step readPrecF
      ; return (Path p)            
      })    
#else                                  
  readsPrec d r = readParen (d > app_prec)
     (\r -> [ (Top, s) 
            | ("Top", s) <- lex r]) r    
     ++
     (\r -> [ (Path p, t) 
            | ("Path", s) <- lex r
            , (f,t) <- readsPrecF (app_prec+1) s]) r    
            
#endif
    
instance ReadF f => Read (Loc f) where 
#ifdef __GLASGOW_HASKELL__
  readPrec = parens $ 
    (prec app_prec $ do
      { Ident "Loc" <- lexP
      ; f <- step readPrec
      ; p <- step readPrec
      ; return (Loc f p)            
      })
#else                                  
  readsPrec d r = readParen (d > app_prec)
     (\r -> [ (Loc f p, u) 
            | ("Loc", s) <- lex r
            , (f,t) <- readsPrec (app_prec+1) s
            , (p,u) <- readsPrec (app_prec+1) t]) r    
#endif

--------------------------------------------------------------------------------
-- * Converting to and from zippers

-- | Creates a zipper from a tree, with the focus at the root.
root :: Mu f -> Loc f
root t = Loc t Top

-- | Restores a tree from a zipper.
defocus :: Traversable f => Loc f -> Mu f
defocus (Loc foc path) = go foc path where
  go t Top = t
  go t (Path xs)  = go (Fix s) path' where
    (Just path', s) = mapAccumL h Nothing xs
    h  old  (Left   y)  =  (old     , y)
    h  _    (Right  p)  =  (Just p  , t)

-- | We attribute all nodes with a zipper focused at that location.
locations :: Traversable f => Mu f -> Attr f (Loc f)
locations tree = go (root tree) tree where
  go loc (Fix t) = Fix (Ann loc t') where
    t' = enumerateWith_ (\j x -> go (unsafeMoveDown j loc) x) t

-- | The list of all locations.
locationsList :: Traversable f => Mu f -> [Loc f]
locationsList = toList . Attrib . locations
    
-- | The zipper version of 'forget'.
locForget :: Functor f => Loc (Ann f a) -> Loc f    
locForget (Loc foc path) = Loc (forget foc) (go path) where
  go :: Functor f => Path (Ann f a) -> Path f    
  go Top = Top
  go (Path (Ann _ nodes)) = Path (fmap h nodes)
  
  h :: Functor f => Node (Ann f a) -> Node f    
  h (Left  t) = Left  (forget t)
  h (Right p) = Right (go p)
    
--------------------------------------------------------------------------------
-- * Manipulating the subtree at focus

-- | Extracts the subtree at focus. Synonym of 'focus'.
extract :: Loc f -> Mu f
extract = focus

-- | Replaces the subtree at focus. 
replace :: Mu f -> Loc f -> Loc f
replace new loc = loc { focus = new }

-- | Modifies the subtree at focus. 
modify :: (Mu f -> Mu f) -> Loc f -> Loc f
modify h loc = replace (h (focus loc)) loc

--------------------------------------------------------------------------------
-- * Safe movements

-- | Moves down to the child with the given index.
-- The leftmost children has index @0@.
moveDown :: Traversable f => Int -> Loc f -> Maybe (Loc f)
moveDown pos (Loc foc path) = new where
  new = case mfoc' of  
     Nothing    ->  Nothing
     Just foc'  ->  Just $ Loc foc' (Path nodes')
  ((mfoc',_),nodes')  =  mapAccumL g (Nothing,0) (unFix foc)    
  g (old,j) x  =  if j==pos 
    then  ((Just x  , j+1),  Right  path  ) 
    else  ((old     , j+1),  Left   x     )      

-- | Moves down to the leftmost child.
moveDownL :: Traversable f => Loc f -> Maybe (Loc f)
moveDownL (Loc foc path) = new where
  new = case mfoc' of  
     Nothing    ->  Nothing
     Just foc'  ->  Just $ Loc foc' (Path nodes')
  (mfoc',nodes')  =  mapAccumL g Nothing (unFix foc)    
  g old x  = case old of
    Nothing  -> (Just x  ,  Right  path  ) 
    _        -> (old     ,  Left   x     )      

-- | Moves down to the rightmost child.
moveDownR :: Traversable f => Loc f -> Maybe (Loc f)
moveDownR (Loc foc path) = new where
  new = case mfoc' of  
     Nothing    ->  Nothing
     Just foc'  ->  Just $ Loc foc' (Path nodes')
  (mfoc',nodes')  =  mapAccumR g Nothing (unFix foc)    
  g old x  = case old of
    Nothing  -> (Just x  ,  Right  path  ) 
    _        -> (old     ,  Left   x     )      
    
--------------------------------------------------------------------------------

-- | Moves up.
moveUp :: Traversable f => Loc f -> Maybe (Loc f)
moveUp (Loc foc path) = case path of
  Top         -> Nothing
  Path nodes  -> 
    case mpath of
      Nothing      -> error "moveUp: shouldn't happen"
      Just path'   -> Just $ case path' of
        Path nodes'    -> Loc (Fix foc') (Path nodes')
        Top            -> Loc (Fix foc') Top
    where      
      (mpath,foc') = mapAccumL g Nothing nodes 
      g old ei = case ei of
        Right  p  -> (Just p  , foc)
        Left   x  -> (old     , x  )

--------------------------------------------------------------------------------

moveRight :: Traversable f => Loc f -> Maybe (Loc f)
moveRight (Loc foc path) = case path of
  Top         -> Nothing
  Path nodes  -> 
    case two of
      Two foc' -> Just $ Loc foc' (Path nodes')
      _        -> Nothing
    where      
      (two,nodes') = mapAccumL g Empty nodes 
      g old ei = case ei of
        Right  p  ->  (One p  , Left foc  )
        Left   x  -> case old of
          One p ->    (Two x  , Right p   )
          _     ->    (old    , ei        )

moveLeft :: Traversable f => Loc f -> Maybe (Loc f)
moveLeft (Loc foc path) = case path of
  Top         -> Nothing
  Path nodes  -> 
    case two of
      Two foc' -> Just $ Loc foc' (Path nodes')
      _        -> Nothing
    where      
      (two,nodes') = mapAccumR g Empty nodes 
      g old ei = case ei of
        Right  p  ->  (One p  , Left foc  )
        Left   x  -> case old of
          One p ->    (Two x  , Right p   )
          _     ->    (old    , ei        )

--------------------------------------------------------------------------------
-- * Testing for borders

-- | Checks whether we are at the top (root).
isTop :: Loc f -> Bool
isTop (Loc _ p) = case p of { Top -> True ; _ -> False }

-- | Checks whether we cannot move down.
isBottom :: Traversable f => Loc f -> Bool
isBottom = isNothing . moveDownL

isLeftmost :: Traversable f => Loc f -> Bool
isLeftmost = isNothing . moveLeft

isRightmost :: Traversable f => Loc f -> Bool
isRightmost = isNothing . moveRight

--------------------------------------------------------------------------------
-- * Location queries

-- | Gives back the index of the given location among the children of its parent.
-- Indexing starts from zero. In case of root node (no parent), we also return zero.
horizontalPos :: Foldable f => Loc f -> Int
horizontalPos (Loc _ path) = case path of
  Top        -> 0
  Path nodes -> 
    case mpos of
      Right pos -> pos
      Left _    -> error "horizontalPos: shouldn't happen"
    where
      mpos = foldl g (Left 0) nodes 
      g old ei = case old of
        Right _ -> old
        Left  j -> case ei of
          Left _  -> Left (j+1)
          Right _ -> Right j

-- | We return the full path from the root as a sequence of child indices.
-- This means that
-- 
-- > loc == foldl (flip unsafeMoveDown) (moveTop loc) (fullPathDown loc)
--
fullPathDown :: Foldable f => Loc f -> [Int]
fullPathDown = reverse . fullPathUp

-- | The following equations hold for 'fullPathUp' and 'fullPathDown':
-- 
-- > fullPathUp == reverse . fullPathDown
-- > loc == foldr unsafeMoveDown (moveTop loc) (fullPathUp loc)
--
fullPathUp :: Foldable f => Loc f -> [Int]
fullPathUp (Loc _ pth) = go pth where    
  go path = case path of
    Top        -> []
    Path nodes -> 
      case mpos of
        Right (pos,parent) -> pos : go parent
        Left _             -> error "fullPathUp: shouldn't happen"
      where
        mpos = foldl g (Left 0) nodes 
        g old ei = case old of
          Right _ -> old
          Left  j -> case ei of
            Left _  -> Left (j+1)
            Right p -> Right (j,p)    

--------------------------------------------------------------------------------
-- * Compound movements
  
-- | Moves to the top, by repeatedly moving up.
moveTop :: Traversable f => Loc f -> Loc f
moveTop = tillNothing moveUp

-- | Moves left until it can.
-- It should be faster than repeated left steps.
leftmost :: Traversable f => Loc f -> Loc f
leftmost orig@(Loc foc path) = case path of
  Top         -> orig
  Path nodes  -> 
    case both of
      Both {}  -> Loc foc' (Path nodes')
      _        -> error "leftmost: shouldn't happen"
    where  
      -- this tricky implementation uses lazyness 
      -- so that we only need a single traversal
      (foc',pnew) = case both of { Both f p -> (f,p) ; _ -> error "leftmost: shouldn't happen" }    
      (both,nodes') = mapAccumL g None nodes 
      g old ei = case old of
        None -> case ei of
          Left  x  -> (First x    , Right pnew)
          Right p  -> (Both foc p , ei        )   -- we are already at the leftmost position
        First f -> case ei of
          Left  x  -> (old        , ei        )
          Right p  -> (Both f p   , Left  foc )
        Both {} -> (old, ei)

-- | Moves right until it can.
-- It should be faster than repeated right steps.
rightmost :: Traversable f => Loc f -> Loc f
rightmost orig@(Loc foc path) = case path of
  Top         -> orig
  Path nodes  -> 
    case both of
      Both {}  -> Loc foc' (Path nodes')
      _        -> error "rightmost: shouldn't happen"
    where  
      -- this tricky implementation uses lazyness 
      -- so that we only need a single traversal
      (foc',pnew) = case both of { Both f p -> (f,p) ; _ -> error "rightmost: shouldn't happen" }    
      (both,nodes') = mapAccumR g None nodes 
      g old ei = case old of
        None -> case ei of
          Left  x  -> (First x    , Right pnew)
          Right p  -> (Both foc p , ei        )   -- we are already at the rightmost position
        First f -> case ei of
          Left  x  -> (old        , ei        )
          Right p  -> (Both f p   , Left  foc )
        Both {} -> (old, ei)
          
--------------------------------------------------------------------------------
-- * Unsafe movements

unsafeMoveDown :: Traversable f => Int -> Loc f -> Loc f
unsafeMoveDown i = unsafe (moveDown i) "unsafeMoveDown: cannot move down"
  
unsafeMoveDownL :: Traversable f => Loc f -> Loc f
unsafeMoveDownR :: Traversable f => Loc f -> Loc f
unsafeMoveUp    :: Traversable f => Loc f -> Loc f

unsafeMoveDownL = unsafe moveDownL "unsafeMoveDownL: cannot move down"
unsafeMoveDownR = unsafe moveDownR "unsafeMoveDownR: cannot move down"  
unsafeMoveUp    = unsafe moveUp    "unsafeMoveUp: cannot move up"  

unsafeMoveLeft, unsafeMoveRight :: Traversable f => Loc f -> Loc f
unsafeMoveLeft  = unsafe moveLeft  "unsafeMoveLeft: cannot move left"  
unsafeMoveRight = unsafe moveRight "unsafeMoveRight: cannot move right"    

--------------------------------------------------------------------------------