module Debug.Trace.Tree.Rose (
Depth
, Offset
, Coords(..)
, Metadata(..)
, isFirstChild
, annotate
) where
import Control.Monad.State
import Data.Tree
type Depth = Int
type Offset = Int
data Coords = Coords {
depth :: Depth
, offset :: Offset
}
deriving (Show, Eq, Ord)
data Metadata = Metadata {
isSpine :: Bool
, nthChild :: Int
, coords :: Coords
}
deriving (Show, Eq, Ord)
isFirstChild :: Metadata -> Bool
isFirstChild Metadata{..} = nthChild == 0
markSpine :: Tree a -> Tree (a, Bool)
markSpine = go True
where
go :: Bool -> Tree a -> Tree (a, Bool)
go isSpine (Node a ts) = Node (a, isSpine)
$ map (uncurry go) $ zip (isSpine : repeat False) ts
markNthChild :: Tree a -> Tree (a, Int)
markNthChild = go 0
where
go :: Int -> Tree a -> Tree (a, Int)
go nth (Node a ts) = Node (a, nth) $ zipWith go [0..] ts
markDepth :: Tree a -> Tree (a, Depth)
markDepth = go 0
where
go :: Depth -> Tree a -> Tree (a, Depth)
go d (Node a ts) = Node (a, d) $ map (go (d + 1)) ts
markCoords :: forall a. Tree a -> Tree (a, Coords)
markCoords t = evalState (unfoldTreeM_BF go (markDepth t)) (Coords 0 0)
where
go :: MonadState Coords m
=> Tree (a, Depth) -> m ((a, Coords), [Tree (a, Depth)])
go (Node (a, depth) ts) = do
cs <- state $ \(Coords curDepth curOffset) ->
if depth > curDepth
then (Coords depth 0 , Coords depth 1)
else (Coords depth curOffset , Coords depth (curOffset + 1))
return ((a, cs), ts)
annotate :: Tree a -> Tree (a, Metadata)
annotate = fmap aux
. markCoords
. markNthChild
. markSpine
where
aux :: (((a, Bool), Int), Coords) -> (a, Metadata)
aux (((a, isSpine), nthChild), coords) = (a, Metadata{..})
_testTree :: Tree String
_testTree =
Node "a" [ Node "b" [ Node "d" []
, Node "e" []
]
, Node "c" [ Node "f" []
]
]
_testTree2 :: Tree String
_testTree2 =
Node "a" [ Node "b" []
, Node "c" [ Node "d" []
, Node "e" []
]
]