module Data.Named.Tree
(
addWords
, Span (..)
, leafSpan
, (<>)
, spanSet
, span
, spanTree
, spanForest
, unSpanTree
, unSpanForest
, sortTree
, sortForest
, mapTrees
) where
import Prelude hiding (span)
import Data.List (sortBy)
import Data.Ord (comparing)
import qualified Data.Tree as T
import qualified Data.IntSet as S
import qualified Data.Map as M
mapTrees :: (a -> b) -> T.Forest a -> T.Forest b
mapTrees f = map (fmap f)
data Span = Span
{ beg :: Int
, end :: Int }
deriving (Show, Eq, Ord)
leafSpan :: Int -> Span
leafSpan i = Span i i
(<>) :: Span -> Span -> Span
Span p q <> Span p' q' = Span (min p p') (max q q')
spanSet :: Span -> S.IntSet
spanSet s = S.fromList [beg s .. end s]
span :: T.Tree (a, Span) -> Span
span = snd . T.rootLabel
spanTree :: (k -> Int) -> T.Tree k -> T.Tree (k, Span)
spanTree f (T.Node k []) = T.Node (k, leafSpan (f k)) []
spanTree f (T.Node k ts) =
let us = spanForest f ts
s = foldl1 (<>) (map span us)
in T.Node (k, s) us
spanForest :: (k -> Int) -> T.Forest k -> T.Forest (k, Span)
spanForest f ts = map (spanTree f) ts
unSpanTree :: T.Tree (k, Span) -> T.Tree k
unSpanTree = fmap fst
unSpanForest :: T.Forest (k, Span) -> T.Forest k
unSpanForest = map unSpanTree
sortTree :: T.Tree (k, Span) -> T.Tree (k, Span)
sortTree (T.Node x ts) = T.Node x (sortForest ts)
sortForest :: T.Forest (k, Span) -> T.Forest (k, Span)
sortForest = sortBy (comparing span) . map sortTree
addWords :: Ord k => T.Forest k -> [k] -> T.Forest k
addWords [] xs = [T.Node x [] | x <- xs]
addWords ts xs
= unSpanForest . T.subForest
. sortTree . fillTree
. dummyRoot
. spanForest f $ ts
where
f = (M.!) $ M.fromList (zip xs [0..])
g = (M.!) $ M.fromList (zip [0..] xs)
dummyRoot = T.Node (undefined, bounds)
bounds = Span 0 (length xs 1)
fillForest = map fillTree
fillTree (T.Node n []) = T.Node n []
fillTree (T.Node (k, s) us) =
let m = spanSet s S.\\ S.unions (map (spanSet . span) us)
mkLeaf i = T.Node (g i, leafSpan i) []
in T.Node (k, s) (fillForest us ++ map mkLeaf (S.toList m))