tree zipper

TreeZipper.hs

-- little helper function
-- f3 (f2 (f1 x)) == x -: f1 -: f2 -: f3
(-:) x f = f x

-- our data structure
data Tree a = Empty
            | Node (Tree a) a (Tree a)
            deriving Show

-- a step on the way through our tree.
-- either going to the left or to the right.
data Direction a = DLeft a (Tree a)
                 | DRight a (Tree a)
                 deriving Show

-- the way from the root node to the current location.
-- with this data the whole tree except the current node
-- can be reconstructed.
type Directions a = [Direction a]

-- the zipper represents all data under the current
-- location (Tree a) and the way to the current
-- location (Directions a).
type Zipper a = (Tree a, Directions a)

-- make a zipper from a tree
makeZipper :: Tree a -> Zipper a
makeZipper tree = (tree, [])

-- descend into the left node
left :: Zipper a -> Zipper a
left ((Node l v r), ds) = (l, DLeft v r : ds)

-- descend into the right node
right :: Zipper a -> Zipper a
right ((Node l v r), ds) = (r, DRight v l : ds)

-- go to the parent of the current zipper position
up :: Zipper a -> Zipper a
up (l, (DLeft v r):ds) = (Node l v r, ds)
up (r, (DRight v l):ds) = (Node l v r, ds)

-- go to the top of the tree
top :: Zipper a -> Zipper a
top (z, []) = (z, [])
top z = top $ up z

-- apply function f to the value at the current zipper position
modify :: (a -> a) -> Zipper a -> Zipper a
modify _ (Empty, ds) = (Empty, ds)
modify f (Node l v r, ds) = (Node l (f v) r, ds)

-- attach a new tree element at the current zipper position
attach :: Tree a -> Zipper a -> Zipper a
attach t (_, ds) = (t, ds)


-- pretty print a tree
ppTree :: (Show a) => Int -> Tree a -> String
ppTree i t = pp t 0
    where pp Empty level = (s level) ++ "(E)\n"
          pp (Node l v r) level =
            pp l (level + i) ++
            (s level) ++ show v ++ "\n" ++
            pp r (level + i)
          s l = take l $ repeat ' '

sampleTree = Node 
                (Node Empty 'B' (Node Empty 'D' Empty))
                'A'
                (Node Empty 'C' Empty)

main = do
    let t = makeZipper sampleTree -: left -: right -: (modify $ const 'X') -: top -: right -: right -: (attach $ Node Empty 'Y' Empty) -: top
    mapM_ putStrLn [ ppTree 4 sampleTree
                   , "-----------------------------------------------------------------"
                   , "change 'D' to 'X'   and   add Node 'Y' as right child of Node 'C'"
                   , "-----------------------------------------------------------------"
                   , ppTree 4 $ fst t
                   ]