-- | A data structure for a static forest.
module Data.Forest.Static where
import Data.Foldable (toList)
import Data.Graph.Inductive.Basic
import Data.List (span,uncons,sort)
import Data.Traversable (mapAccumL)
import Debug.Trace
import qualified Data.Map.Strict as S
import qualified Data.Tree as T
import qualified Data.Vector as V
import qualified Data.Vector.Generic as VG
import qualified Data.Vector.Unboxed as VU
-- | Kind of possible @TreeOrder@s.
--
-- TODO @In@ for in-order traversal?
--
-- TODO @Unordered@ for trees that have no sorted order?
data TreeOrder = Pre | Post
-- | A static forest structure. While traversals are always explicitly
-- possible by following the indices, the nodes themselves shall always be
-- ordered by the type @p :: TreeOrder@. This is not completely enforced,
-- given that @Forest@ is exporting the constructor, but encouraged via
-- construction with helper functions.
data Forest (p :: TreeOrder) v a where
Forest :: (VG.Vector v a) =>
{ label :: v a
, parent :: VU.Vector Int
, children :: V.Vector (VU.Vector Int)
, lsib :: VU.Vector Int
, rsib :: VU.Vector Int
, roots :: VU.Vector Int
} -> Forest p v a
deriving instance (Show a, Show (v a)) => Show (Forest p v a)
-- | Construct a static 'Forest' with a tree traversal function. I.e.
-- @forestWith preorderF trees@ will construct a pre-order forest from the
-- list of @trees@.
forestWith :: (VG.Vector v a) => (forall a . [T.Tree a] -> [a]) -> [T.Tree a] -> Forest (p::TreeOrder) v a
forestWith f ts
= Forest { label = VG.fromList $ f ts
, parent = VU.fromList $ map (\(_,k,_ ,_) -> k ) $ f pcs
, children = V.fromList $ map (\(_,_,cs,_) -> VU.fromList cs) $ f pcs
, lsib = VU.fromList $ map fst $ S.elems lr
, rsib = VU.fromList $ map snd $ S.elems lr
, roots = VU.fromList $ map (fst . T.rootLabel) us
}
where
-- Step 1: construct a forest isomorphic to @ts@ but labelled with
-- a total order of unique identifiers. (That is: label with @Int@s).
-- The unique identifiers are in pre-order.
ps = addIndicesF' 0 ts
-- Step 2: use @f@ to produce a permutation map and apply this
-- permutation to turn the pre-order @ps@ into the required order.
backp = VU.fromList $ map snd $ sort $ zip (f ps) [0..]
-- Step 3: decorate the forest with indices in the correct order. Keep
-- the label in @snd@.
us = map (fmap (\(k,l) -> (backp VG.! k,l))) $ addIndicesF 0 ts
-- Step 4: add the correct relations (children, lrSibling, parents)
pcs = parentChildrenF (-1) us
-- A map with the left and right sibling
lr = lrSiblingF us
-- | Construct a pre-ordered forest.
forestPre :: (VG.Vector v a) => [T.Tree a] -> Forest Pre v a
forestPre = forestWith preorderF
-- | Construct a post-ordered forest.
forestPost :: (VG.Vector v a) => [T.Tree a] -> Forest Post v a
forestPost = forestWith postorderF
-- | Add @pre-ordered@ !!! indices. First argument is the starting index.
addIndices :: Int -> T.Tree a -> T.Tree (Int,a)
addIndices k = snd . mapAccumL (\i e -> (i+1, (i,e))) k
-- | Add @pre-ordered@ !!! indices, but to a forest.
addIndicesF :: Int -> [T.Tree a] -> [T.Tree (Int,a)]
addIndicesF k = snd . mapAccumL go k
where go = mapAccumL (\i e -> (i+1, (i,e)))
-- | Add @pre-ordered@ !!! indices to a forest, but throw the label away as
-- well.
addIndicesF' :: Int -> [T.Tree a] -> [T.Tree Int]
addIndicesF' k = snd . mapAccumL go k
where go = mapAccumL (\i e -> (i+1, i))
-- | Add parent + children information. Yields
-- @(Index,Parent,[Child],Label)@. Parent is @-1@ if root node.
parentChildrenF :: Int -> [T.Tree (Int,a)] -> [T.Tree (Int,Int,[Int],a)]
parentChildrenF k ts = [ T.Node (i,k,children sf,l) (parentChildrenF i sf) | T.Node (i,l) sf <- ts ]
where children sf = map (fst . T.rootLabel) sf
-- | Return a map with all the nearest siblings for each node, for a forest.
lrSiblingF :: [T.Tree (Int,a)] -> S.Map Int (Int,Int)
lrSiblingF = S.delete (-1) . lrSibling . T.Node (-1,error "laziness in lrSiblingF broken")
-- | Return a map with all the nearest siblings for each node, for a tree.
lrSibling :: T.Tree (Int,a) -> S.Map Int (Int,Int)
lrSibling = S.fromList . map splt . T.flatten . go ([]::[Int])
where go sib (T.Node (k,lbl) frst) = let cs = [l | T.Node (l,_) _ <- frst] in T.Node (k,lbl,sib) [ go cs t | t <- frst]
splt (k,_,[]) = (k,(-1,-1))
splt (k,_,sbl) = let (ls,rs) = span (/=k) sbl in (k,(last $ (-1):ls,head $ tail rs ++ [-1]))
-- | Return the left-most leaf for each node.
leftMostLeaves :: Forest p v a -> VU.Vector Int
leftMostLeaves f = VG.map go $ VG.enumFromN 0 $ VG.length $ parent f
where go k = let cs = children f VG.! k
in if VG.null cs then k else go (VG.head cs)
-- | Return all left key roots. These are the nodes that have no (super-)
-- parent with the same left-most leaf.
--
-- This function is somewhat specialized for tree editing.
--
-- TODO group by
leftKeyRoots :: Forest Post v a -> VU.Vector Int
leftKeyRoots f = VU.fromList . sort . S.elems $ VU.foldl' go S.empty (VU.enumFromN (0::Int) $ VG.length $ parent f)
-- Build a map from left-most leaf to most root-near node.
where go s k = S.insertWith max (lml VU.! k) k s
lml = leftMostLeaves f
{-
test :: [T.Tree Char]
test = [T.Node 'R' [T.Node 'a' [], T.Node 'b' []], T.Node 'S' [T.Node 'x' [], T.Node 'y' []]]
runtest = do
print (forestPre test :: Forest Pre V.Vector Char)
print (forestPost test :: Forest Post V.Vector Char)
print (forestPost [T.Node 'R' [T.Node 'a' []]] :: Forest Post V.Vector Char)
print (forestPost [T.Node 'R' [T.Node 'a' [], T.Node 'b' []]] :: Forest Post V.Vector Char)
-}