Tree aに対する継続zipper

TreeZipper.hs

{-# LANGUAGE FlexibleContexts, FlexibleInstances, MultiParamTypeClasses, Rank2Types #-}
module TreeZipper
  (
  -- * Directions
    After(..), Before(..)
  , Exit(..), Enter(..)
  , BackForth(..)
  -- * Default directions
  , To(..), From(..), Next(..)
  -- * Zipper
  , Zipper(..)
  , zipper, start, continue
  -- * Walk combinators
  , stop, around, throughout, backForth, through
  , walk, walkTree
  ) where

import Control.Applicative
import Control.Monad.Cont (ContT(..), runContT)
import Control.Monad.Error (Error(..), runErrorT, throwError)
import Control.Monad.Trans (lift)
import Control.Monad.Writer (WriterT, runWriterT, tell)
import Data.Char (isSpace)
import Data.Maybe (fromMaybe, isJust)
import Data.Monoid (Any(..))
import Data.Tree (Tree(..))
import Text.Show.Functions ()

-- | A Walk is a monadic traversal
type Walk from to part whole
  = forall m. (Monad m, Applicative m)
  => (from -> part -> m (Maybe part, to)) -- A visitor is a way to enumerate parts of a value
  -> whole                                -- The value whose parts to enumerate
  -> m (Maybe whole)                      -- A monadic traversal

-- | The simplest case where only one incoming direction
data Before = Before deriving (Eq, Ord, Read, Show)

-- | The simplest case where only one outgoing direction
data After = After deriving (Eq, Ord, Read, Show)

-- | Two incoming directions
data Exit from = Exit Bool from deriving (Eq, Ord, Read, Show)

-- | Two outgoing directions
data Enter to = Enter | To to deriving (Eq, Ord, Read, Show)

-- | The extra outgoing direction
data BackForth to = Back | Forth to deriving (Eq, Ord, Read, Show)


class From from where
  before :: from -- ^ A default incoming direction

class Eq to => To to where
  after :: to -- ^ A default outgoing direction

class (From from, To to, Show from, Read to) => Next from to where
  next :: from -> to -- ^ 


instance To After where
  after = After

instance To to => To (Enter to) where
  after = To after

instance To to => To (BackForth to) where
  after = Forth after


instance From Before where
  before = Before

instance From from => From (Exit from) where
  before = Exit False before


instance Next Before After where
  next Before = After

instance Next Before to => Next (Exit Before) (Enter to) where
  next (Exit False Before) = Enter
  next (Exit True  Before) = To after

instance Next from to => Next from (BackForth to) where
  next = Forth . next

-- | The trivial one-stop walk from an incoming direction 'from'
stop :: from -> Walk from After a a
stop from visit a = fst <$> visit from a

-- | Walk around immediate subtrees
walk :: Walk from to (Tree a) (Tree a) -> Walk from to (Tree a) (Tree a)
walk walk' visit t@(Node a subtrees) = do
  (t', Any dirty) <- runWriterT $ Node a <$> sequence (map f subtrees)
  return $ scavenge dirty t'
    where
      f subtree = do
        msubtree <- lift $ walk' visit subtree
        tell . Any $ isJust msubtree
        return $ fromMaybe subtree msubtree

-- | Throws away a new whole if no part changed
scavenge :: Bool -> a -> Maybe a
scavenge True a  = Just a
scavenge False _ = Nothing

-- | A debugging visitor
debug :: (Next from to, Show a) => from -> Tree a -> IO (Maybe (Tree a), to)
debug from part@(Node a bs) = do
  putStrLn $ show from ++ ": " ++ show part
  return (Nothing, next from)

-- | A recursive tourism
around :: Walk from to part whole                -- Steps through parts of a whole
       -> Walk (Exit from) (Enter to) part part  -- Steps through subparts of a part
       -> Walk (Exit from) (Enter to) part whole -- Steps through recursively
around walkOuter walkInner visit = walkOuter (visit' False False)
    where
      visit' dirty around from part = do
        (part1', to) <- visit (Exit around from) part
        let (dirty1, part1) = pollute dirty part part1'
        case to of
          Enter -> do
            part2' <- walkInner visit part1
            let (dirty2, part2) = pollute dirty1 part1 part2'
            visit' dirty2 True from part2
          To to -> return (scavenge dirty1 part1, to)

pollute :: Bool -> a -> Maybe a -> (Bool, a)
pollute dirty a Nothing  = (dirty, a)
pollute _     _ (Just a) = (True, a)

newYork :: Walk from to part whole
newYork  _ _ = return Nothing

-- | Walks through all nodes of a tree
throughout :: Walk from to (Tree a) (Tree a)
           -> Walk (Exit from) (Enter to) (Tree a) (Tree a)
throughout level = level `around` walk (throughout level)

-- | An interactive visitor
keyboard :: (Next from to, Show a, Read a)
         => from -> a -> IO (Maybe a, to)
keyboard from x = do
  putStr $ show from ++ ": " ++ show x ++ "\n? "
  line <- getLine
  return $ if all isSpace line
             then (Nothing, next from)
             else read line

-- | A zipper is a suspended walk
data Zipper from to part a
  = Done a
  | Stop from part (Maybe part -> to -> Zipper from to part a)
  deriving Show

-- Zipper monad is designed to be the reification of a walk
instance Monad (Zipper from to part) where
  return = Done
  Done a           >>= k = k a
  Stop from part c >>= k = Stop from part c'
    where c' part' to = c part' to >>= k

-- | Converts a walk to a Zipper
zipper :: Walk from to part whole
       -> whole
       -> Zipper from to part (Maybe whole)
zipper walk whole = runContT (walk visit whole) return
  where visit from part = ContT (Stop from part . curry)

start :: Tree a
      -> Zipper (Exit Before) (Enter After) (Tree a) (Maybe (Tree a))
start = zipper (throughout (stop Before))

continue :: Zipper from to part whole
         -> Maybe part
         -> to
         -> Zipper from to part whole
continue (Done _)     = error "Zipper is done, not at a stop"
continue (Stop _ _ c) = c

-- | A serial record of a visitor's behavior over the course of walk
data Diff part to = Diff (Maybe part) to (Diff part to)

instance (Show part, Show to) => Show (Diff part to) where
  showsPrec = loop 3
    where loop 0 _ _ = showString "..."
          loop l d (Diff part to diff) = showParen (d > 0)
            $ showString "Diff " . showsPrec 11 part . showChar ' '
            . showsPrec 11 to . showString " $ " . loop (l - 1) 0 diff

-- | An empty 'Diff' record
same :: To to => Diff part to
same = Diff Nothing after same

-- | Replays a Diff against a zipper
replay :: Zipper from to part whole -> Diff part to -> whole
replay (Done whole) _                   = whole
replay (Stop _ _ c) (Diff part to diff) = replay (c part to) diff

-- | Converts any zipper into a walk over a 'Diff'
walkDiff :: To to
         => Zipper from to part whole
         -> Walk (Exit (Zipper from to part whole))
                 (Enter After)
                 (Diff part to)
                 (Diff part to)
walkDiff zipper = stop zipper `around`
  \visit ~(Diff part to diff) -> case zipper of
    Done _     -> return Nothing
    Stop _ _ c -> fmap (Diff part to) <$> walkDiff (c part to) visit diff

-- | Turns any walk with outgoing direction type 'to' into a walk with
--   outoging direction bype 'BackForth to'
backForth :: (To to, Error (Maybe whole))
          => Walk from to part whole
          -> Walk from (BackForth to) part whole
backForth walk visit whole =
  either id (>>= replay za) <$> runErrorT (walkDiff za visit' same)
  where
    za = zipper walk whole
    visit' (Exit _ (Done whole)) _ = throwError whole
    visit' (Exit _ (Stop from part _))
           (Diff partD toD diffD) = do
      (part, to) <- lift $ visit from $ fromMaybe part partD
      let diff' | isJust part = Just $ Diff part toD diffD
                | otherwise   = Nothing
      return $ case to of
        Back -> (diff', after)
        Forth to -> ( if to == toD
                        then diff'
                        else Just (Diff part to same)
                    , Enter)

instance Error (Maybe whole) where
  noMsg = Nothing

walkTree :: Walk Before After (Tree a) (Tree a)
walkTree = walk (stop Before)

through :: Walk (Exit Before) (Enter (BackForth After)) (Tree a) (Tree a)
through = stop' `around` throughout (backForth (walk stop'))
  where stop' visit a = fst <$> visit before a

-- tests
tree :: Tree Int
tree = Node 0 [ Node 1 [ Node 3 [], Node 4 [] ], Node 2 [ Node 5 [ Node 6 [] ] ]]

main :: IO ()
main = do
  putStrLn "-- Top level: Go right --"
  stop Before debug tree
  putStrLn "-- Immediate children: Go right --"
  walk (stop Before) debug tree
  putStrLn "-- Top level: Go down, go right and go up --"
  (stop Before `around` newYork) debug tree
  putStrLn "-- Second level: Go down, go right and go up --"
  (stop Before `around` walk (stop Before `around` newYork)) debug tree
  putStrLn "-- All levels: Go down, go right and go up --"
  throughout (stop Before) debug tree
  putStrLn "-- All levels: Interactive session --"
  throughout (stop Before) keyboard tree
  putStrLn "-- New zipper --"
  through keyboard tree
  return ()

{-
Prelude> walkTree keyboard tree
Before: Node {rootLabel = 1, subForest = [Node {rootLabel = 3, subForest = []},Node {rootLabel = 4, subForest = []}]}
? (Just (Node {rootLabel = 100, subForest = []}), After)
Before: Node {rootLabel = 2, subForest = [Node {rootLabel = 5, subForest = [Node {rootLabel = 6, subForest = []}]}]}
? 
Just (Node {rootLabel = 0, subForest = [Node {rootLabel = 100, subForest = []},Node {rootLabel = 2, subForest = [Node {rootLabel = 5, subForest = [Node {rootLabel = 6, subForest = []}]}]}]})
-}