module TreeZipper where

-- requires `cabal install pretty`
import Text.PrettyPrint


--
-- * Binary Tree
--

-- | Binary trees.
data Tree a
   = Node a (Tree a) (Tree a)
   | Leaf
  deriving Eq
   
-- Trunk a (Tree a)
-- Node3 a (Tree a) (Tree a) (Tree a)

instance Functor Tree where
  fmap _ Leaf         = Leaf
  fmap f (Node a l r) = Node (f a) (fmap f l) (fmap f r)


-- ** Basic editing operations on nodes

-- | Apply a function to the value of a node.
apply :: (a -> a) -> Tree a -> Tree a
apply f (Node a l r) = Node (f a) l r
apply _ Leaf         = Leaf

-- | Set the value of a node.
set :: a -> Tree a -> Tree a
set a (Node _ l r) = Node a l r
set a Leaf         = Node a Leaf Leaf

-- | Swap the branches.
swap :: Tree a -> Tree a
swap (Node a l r) = Node a r l
swap Leaf         = Leaf


--
-- * Tree Paths / Pointers
--

-- | A direction to take at an internal node.
data Dir = L | R
  deriving (Eq,Show)

-- | A path through the tree.
type Path = [Dir]

-- | Get the node at a position.
get :: Path -> Tree a -> Tree a
get []    t            = t
get (L:p) (Node _ l _) = get p l
get (R:p) (Node _ _ r) = get p r
get _     _            = Leaf

-- | Update the node at a position.
update :: (Tree a -> Tree a) -> Path -> Tree a -> Tree a
update f []    t            = f t
update f (L:p) (Node a l r) = Node a (update f p l) r
update f (R:p) (Node a l r) = Node a l (update f p r)
update f _     _            = f Leaf


--
-- * Tree Zippers
--

-- | A context is a binary tree with a hole in it.
data Context a
   = InL a (Context a) (Tree a)
   | InR a (Tree a) (Context a)
   | Root
  deriving Eq
   
-- InTrunk a (Context a)
-- In13 a (Context a) (Tree a) (Tree a)
-- In23 a (Tree a) (Context a) (Tree a)
-- In33 a (Tree a) (Tree a) (Context a)
   
-- | A tree zipper is like a binary tree with a movable pointer.
--   It's just a context and the tree that goes in the hole!
data Zipper a = Z (Tree a) (Context a)
  deriving Eq


-- ** Navigating the zipper

-- | Flipped function composition to make zipper sequences more readable.
(.>) :: (a -> b) -> (b -> c) -> a -> c
(.>) = flip (.)

-- | Point to the root of a tree.
enter :: Tree a -> Zipper a
enter t = Z t Root

-- | Exit out of a zipper, returning the resulting tree.
exit :: Zipper a -> Tree a
exit (Z t Root)        = t
exit (Z t (InL a c r)) = exit (Z (Node a t r) c)
exit (Z t (InR a l c)) = exit (Z (Node a l t) c)

-- | Run a zipper function on a tree.
runZipper :: (Zipper a -> Zipper a) -> Tree a -> Tree a
runZipper f = enter .> f .> exit

-- | Move down the left branch of the current node.
left :: Zipper a -> Zipper a
left (Z (Node a l r) c) = Z l (InL a c r)
left z                  = z

-- | Move down the right branch of the current node.
right :: Zipper a -> Zipper a
right (Z (Node a l r) c) = Z r (InR a l c)
right z                  = z

-- | Move up to the parent of the current node.
up :: Zipper a -> Zipper a
up (Z t (InL a c r)) = Z (Node a t r) c
up (Z t (InR a l c)) = Z (Node a l t) c
up z                 = z


-- ** Editing operations

-- | Modify the current node.
modify :: (Tree a -> Tree a) -> Zipper a -> Zipper a
modify f (Z t c) = Z (f t) c

-- | Replace the current node.
replace :: Tree a -> Zipper a -> Zipper a
replace t (Z _ c) = Z t c


-- ** Example

-- | Generate a tree of the given height containing all zeros.
zeros :: Int -> Tree Int
zeros 0 = Leaf
zeros n = let t = zeros (n-1) in Node 0 t t

-- | An edit script that adds 5 to the current focus, 10 to its left child
--   and 20 to its right child.
s1 :: Num a => Zipper a -> Zipper a
s1 = modify (apply (+5))
  .> left
  .> modify (apply (+10))
  .> up .> right
  .> modify (apply (+20))
  .> up


--
-- * Pretty printing
--

instance Show a => Show (Tree a) where
  show = render . prettyTree
instance Show a => Show (Context a) where
  show = render . prettyContext
instance Show a => Show (Zipper a) where
  show = render . prettyZipper

prettyZipper (Z t c) = text "Tree:"    $$ nest 2 (prettyTree t)
                    $$ text "Context:" $$ nest 2 (prettyContext c)

prettyNode name end fl fr a l r =
    parens (text name <+> text val
      $$ nest indL (fl l)
      $$ nest indR (fr r))
  where 
    val  = show a
    indL = 6 + length val
    indR = indL + if end l then 5 else 0

-- | Pretty print a tree.
prettyTree :: Show a => Tree a -> Doc
prettyTree Leaf         = text "Leaf"
prettyTree (Node a l r) = prettyNode "Node" isLeaf prettyTree prettyTree a l r

-- | Pretty print a context.
prettyContext :: Show a => Context a -> Doc
prettyContext Root        = text "Root"
prettyContext (InL a c r) = prettyNode "InL " isRoot prettyContext prettyTree a c r
prettyContext (InR a l c) = prettyNode "InR " isLeaf prettyTree prettyContext a l c

isLeaf Leaf = True
isLeaf _    = False

isRoot Root = True
isRoot _    = False