module Data.Named.Tree
(
Span (..)
, leafSpan
, (<>)
, spanSet
, span
, spanTree
, spanForest
, unSpanTree
, unSpanForest
, sortTree
, sortForest
, mapLeaves
, mapNodes
, mapTrees
) where
import Prelude hiding (span)
import Data.List (sortBy)
import Data.Ord (comparing)
import Data.Ix (Ix, range)
import qualified Data.Tree as T
import qualified Data.Set as S
mapLeaves :: (a -> b) -> T.Tree (Either c a) -> T.Tree (Either c b)
mapLeaves f (T.Node (Left x) ts) = T.Node (Left x) (map (mapLeaves f) ts)
mapLeaves f (T.Node (Right x) _) = T.Node (Right $ f x) []
mapNodes :: (a -> b) -> T.Tree (Either a c) -> T.Tree (Either b c)
mapNodes f (T.Node (Left x) ts) = T.Node (Left $ f x) (map (mapNodes f) ts)
mapNodes _ (T.Node (Right x) _) = T.Node (Right x) []
mapTrees :: (a -> b) -> T.Forest a -> T.Forest b
mapTrees f = map (fmap f)
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')
spanSet :: Ix w => Span w -> S.Set w
spanSet s = S.fromList $ range (beg s, end s)
span :: T.Tree (a, Span w) -> Span w
span = snd . T.rootLabel
spanTree :: Ord w => T.Tree (Either n w) -> T.Tree (Either n w, Span w)
spanTree (T.Node (Right k) []) = T.Node (Right k, leafSpan k) []
spanTree (T.Node k ts) =
let us = spanForest ts
s = foldl1 (<>) (map span us)
in T.Node (k, s) us
spanForest :: Ord w => T.Forest (Either n w) -> T.Forest (Either n w, Span w)
spanForest = map spanTree
unSpanTree :: T.Tree (k, Span w) -> T.Tree k
unSpanTree = fmap fst
unSpanForest :: T.Forest (k, Span w) -> T.Forest k
unSpanForest = map unSpanTree
sortTree :: Ord w => T.Tree (k, Span w) -> T.Tree (k, Span w)
sortTree (T.Node x ts) = T.Node x (sortForest ts)
sortForest :: Ord w => T.Forest (k, Span w) -> T.Forest (k, Span w)
sortForest = sortBy (comparing span) . map sortTree