module Data.Named.Tree
(
NeTree
, NeForest
, Span (..)
, leafSpan
, (<>)
, spanSet
, span
, spanTree
, spanForest
, unSpanTree
, unSpanForest
, sortTree
, sortForest
, mapForest
, mapTree
, onLeaf
, onNode
, onEither
, onBoth
, groupForestLeaves
, groupTreeLeaves
, concatForestLeaves
, concatTreeLeaves
, module Data.Tree
) where
import Prelude hiding (span, (<>))
import Data.List (sortBy, groupBy)
import Data.Either (rights)
import Data.Ord (comparing)
import Data.Ix (Ix, range)
import Data.Tree
import qualified Data.Set as S
type NeTree a b = Tree (Either a b)
type NeForest a b = Forest (Either a b)
onLeaf :: (a -> b) -> Either c a -> Either c b
onLeaf _ (Left x) = Left x
onLeaf f (Right x) = Right (f x)
{-# INLINE onLeaf #-}
onNode :: (a -> b) -> Either a c -> Either b c
onNode f (Left x) = Left (f x)
onNode _ (Right x) = Right x
{-# INLINE onNode #-}
onEither :: (a -> c) -> (b -> d) -> Either a b -> Either c d
onEither f _ (Left x) = Left (f x)
onEither _ g (Right x) = Right (g x)
{-# INLINE onEither #-}
onBoth :: (a -> b) -> Either a a -> Either b b
onBoth f (Left x) = Left (f x)
onBoth f (Right x) = Right (f x)
{-# INLINE onBoth #-}
mapForest :: (a -> b) -> Forest a -> Forest b
mapForest = map . mapTree
{-# INLINE mapForest #-}
mapTree :: (a -> b) -> Tree a -> Tree b
mapTree = fmap
{-# INLINE mapTree #-}
groupForestLeaves :: (b -> b -> Bool) -> NeForest a b -> NeForest a [b]
groupForestLeaves f
= concatMap joinLeaves
. groupBy (both isLeaf)
. map (groupTreeLeaves f)
where
joinLeaves [x] = [x]
joinLeaves xs =
let ys = (concat . rights) (map rootLabel xs)
in [Node (Right ys') [] | ys' <- groupBy f ys]
both g x y = g x && g y
isLeaf (Node (Right _) []) = True
isLeaf _ = False
groupTreeLeaves :: (b -> b -> Bool) -> NeTree a b -> NeTree a [b]
groupTreeLeaves f (Node v xs) = Node (fmap (:[]) v) (groupForestLeaves f xs)
concatForestLeaves :: NeForest a [b] -> NeForest a b
concatForestLeaves = concatMap concatTreeLeaves
concatTreeLeaves :: NeTree a [b] -> NeForest a b
concatTreeLeaves (Node (Left x) xs) = [Node (Left x) (concatForestLeaves xs)]
concatTreeLeaves (Node (Right xs) _) = [Node (Right x) [] | x <- xs]
data Span w = Span
{ beg :: w
, end :: w }
deriving (Show, Eq, Ord)
leafSpan :: w -> Span w
leafSpan i = Span i i
(<>) :: Ord w => Span w -> Span w -> Span w
Span p q <> Span p' q' = Span (min p p') (max q q')
{-# INLINE (<>) #-}
spanSet :: Ix w => Span w -> S.Set w
spanSet s = S.fromList $ range (beg s, end s)
span :: Tree (a, Span w) -> Span w
span = snd . rootLabel
spanTree :: Ord w => Tree (Either n w) -> Tree (Either n w, Span w)
spanTree (Node (Right k) []) = Node (Right k, leafSpan k) []
spanTree (Node k ts) =
let us = spanForest ts
s = foldl1 (<>) (map span us)
in Node (k, s) us
spanForest :: Ord w => Forest (Either n w) -> Forest (Either n w, Span w)
spanForest = map spanTree
unSpanTree :: Tree (k, Span w) -> Tree k
unSpanTree = fmap fst
unSpanForest :: Forest (k, Span w) -> Forest k
unSpanForest = map unSpanTree
sortTree :: Ord w => Tree (k, Span w) -> Tree (k, Span w)
sortTree (Node x ts) = Node x (sortForest ts)
sortForest :: Ord w => Forest (k, Span w) -> Forest (k, Span w)
sortForest = sortBy (comparing span) . map sortTree