-- | A data structure for a static forest.

module Data.Forest.Static where

import           Control.DeepSeq (NFData(..))
import           Control.Applicative ((<$>),(<*>))
import           Control.Monad (replicateM)
import           Data.Foldable (toList)
import           Data.Graph.Inductive.Basic
import           Data.List (span,uncons,sort)
import           Data.Traversable (mapAccumL)
import           Data.Tree (Tree)
import           Debug.Trace
import qualified Data.List as L
import qualified Data.Map.Strict as S
import qualified Data.Set as Set
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
import           Test.QuickCheck
import           GHC.Generics(Generic)
import           Data.Aeson (ToJSON(..),FromJSON(..))



-- | Kind of possible @TreeOrder@s.
--
-- TODO @In@ for in-order traversal?
--
-- TODO @Unordered@ for trees that have no sorted order?

data TreeOrder = Pre | Post | Unordered



-- | 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. The labels of type @a@ (in @label@)
-- require a vector structure @v@ for @O(1)@ access.

data Forest (p ∷ TreeOrder) v a = Forest
  { forall (p :: TreeOrder) (v :: * -> *) a. Forest p v a -> v a
label     ∷ !(v a)
    -- ^ Each node @k@ in @[0..n-1]@ has a label at @label ! k@.
  , forall (p :: TreeOrder) (v :: * -> *) a. Forest p v a -> Vector Int
parent    ∷ !(VU.Vector Int)
    -- ^ Each node @k@ has a parent node, or @-1@ if there is no such
    -- parent.
  , forall (p :: TreeOrder) (v :: * -> *) a.
Forest p v a -> Vector (Vector Int)
children  ∷ !(V.Vector (VU.Vector Int))
    -- ^ Each node @k@ has a vector of indices for its children. For leaf
    -- nodes, the vector is empty.
  , forall (p :: TreeOrder) (v :: * -> *) a. Forest p v a -> Vector Int
lsib      ∷ !(VU.Vector Int)
    -- ^ The left sibling for a node @k@. Will *not* cross subtrees. I.e.
    -- if @k@ is @lsib@ of @l@, then @k@ and @l@ have the same parent.
  , forall (p :: TreeOrder) (v :: * -> *) a. Forest p v a -> Vector Int
rsib      ∷ !(VU.Vector Int)
    -- ^ The right sibling for a node @k@.
  , forall (p :: TreeOrder) (v :: * -> *) a. Forest p v a -> Vector Int
roots     ∷ !(VU.Vector Int)
    -- ^ The roots of the individual trees, the forest was constructed
    -- from.
  }
  deriving (Forest p v a -> Forest p v a -> Bool
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
forall (p :: TreeOrder) (v :: * -> *) a.
Eq (v a) =>
Forest p v a -> Forest p v a -> Bool
/= :: Forest p v a -> Forest p v a -> Bool
$c/= :: forall (p :: TreeOrder) (v :: * -> *) a.
Eq (v a) =>
Forest p v a -> Forest p v a -> Bool
== :: Forest p v a -> Forest p v a -> Bool
$c== :: forall (p :: TreeOrder) (v :: * -> *) a.
Eq (v a) =>
Forest p v a -> Forest p v a -> Bool
Eq,Forest p v a -> Forest p v a -> Bool
Forest p v a -> Forest p v a -> Ordering
forall a.
Eq a
-> (a -> a -> Ordering)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> a)
-> (a -> a -> a)
-> Ord a
forall {p :: TreeOrder} {v :: * -> *} {a}.
Ord (v a) =>
Eq (Forest p v a)
forall (p :: TreeOrder) (v :: * -> *) a.
Ord (v a) =>
Forest p v a -> Forest p v a -> Bool
forall (p :: TreeOrder) (v :: * -> *) a.
Ord (v a) =>
Forest p v a -> Forest p v a -> Ordering
forall (p :: TreeOrder) (v :: * -> *) a.
Ord (v a) =>
Forest p v a -> Forest p v a -> Forest p v a
min :: Forest p v a -> Forest p v a -> Forest p v a
$cmin :: forall (p :: TreeOrder) (v :: * -> *) a.
Ord (v a) =>
Forest p v a -> Forest p v a -> Forest p v a
max :: Forest p v a -> Forest p v a -> Forest p v a
$cmax :: forall (p :: TreeOrder) (v :: * -> *) a.
Ord (v a) =>
Forest p v a -> Forest p v a -> Forest p v a
>= :: Forest p v a -> Forest p v a -> Bool
$c>= :: forall (p :: TreeOrder) (v :: * -> *) a.
Ord (v a) =>
Forest p v a -> Forest p v a -> Bool
> :: Forest p v a -> Forest p v a -> Bool
$c> :: forall (p :: TreeOrder) (v :: * -> *) a.
Ord (v a) =>
Forest p v a -> Forest p v a -> Bool
<= :: Forest p v a -> Forest p v a -> Bool
$c<= :: forall (p :: TreeOrder) (v :: * -> *) a.
Ord (v a) =>
Forest p v a -> Forest p v a -> Bool
< :: Forest p v a -> Forest p v a -> Bool
$c< :: forall (p :: TreeOrder) (v :: * -> *) a.
Ord (v a) =>
Forest p v a -> Forest p v a -> Bool
compare :: Forest p v a -> Forest p v a -> Ordering
$ccompare :: forall (p :: TreeOrder) (v :: * -> *) a.
Ord (v a) =>
Forest p v a -> Forest p v a -> Ordering
Ord,ReadPrec [Forest p v a]
ReadPrec (Forest p v a)
ReadS [Forest p v a]
forall a.
(Int -> ReadS a)
-> ReadS [a] -> ReadPrec a -> ReadPrec [a] -> Read a
forall (p :: TreeOrder) (v :: * -> *) a.
Read (v a) =>
ReadPrec [Forest p v a]
forall (p :: TreeOrder) (v :: * -> *) a.
Read (v a) =>
ReadPrec (Forest p v a)
forall (p :: TreeOrder) (v :: * -> *) a.
Read (v a) =>
Int -> ReadS (Forest p v a)
forall (p :: TreeOrder) (v :: * -> *) a.
Read (v a) =>
ReadS [Forest p v a]
readListPrec :: ReadPrec [Forest p v a]
$creadListPrec :: forall (p :: TreeOrder) (v :: * -> *) a.
Read (v a) =>
ReadPrec [Forest p v a]
readPrec :: ReadPrec (Forest p v a)
$creadPrec :: forall (p :: TreeOrder) (v :: * -> *) a.
Read (v a) =>
ReadPrec (Forest p v a)
readList :: ReadS [Forest p v a]
$creadList :: forall (p :: TreeOrder) (v :: * -> *) a.
Read (v a) =>
ReadS [Forest p v a]
readsPrec :: Int -> ReadS (Forest p v a)
$creadsPrec :: forall (p :: TreeOrder) (v :: * -> *) a.
Read (v a) =>
Int -> ReadS (Forest p v a)
Read,Int -> Forest p v a -> ShowS
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
forall (p :: TreeOrder) (v :: * -> *) a.
Show (v a) =>
Int -> Forest p v a -> ShowS
forall (p :: TreeOrder) (v :: * -> *) a.
Show (v a) =>
[Forest p v a] -> ShowS
forall (p :: TreeOrder) (v :: * -> *) a.
Show (v a) =>
Forest p v a -> String
showList :: [Forest p v a] -> ShowS
$cshowList :: forall (p :: TreeOrder) (v :: * -> *) a.
Show (v a) =>
[Forest p v a] -> ShowS
show :: Forest p v a -> String
$cshow :: forall (p :: TreeOrder) (v :: * -> *) a.
Show (v a) =>
Forest p v a -> String
showsPrec :: Int -> Forest p v a -> ShowS
$cshowsPrec :: forall (p :: TreeOrder) (v :: * -> *) a.
Show (v a) =>
Int -> Forest p v a -> ShowS
Show,forall a.
(forall x. a -> Rep a x) -> (forall x. Rep a x -> a) -> Generic a
forall (p :: TreeOrder) (v :: * -> *) a x.
Rep (Forest p v a) x -> Forest p v a
forall (p :: TreeOrder) (v :: * -> *) a x.
Forest p v a -> Rep (Forest p v a) x
$cto :: forall (p :: TreeOrder) (v :: * -> *) a x.
Rep (Forest p v a) x -> Forest p v a
$cfrom :: forall (p :: TreeOrder) (v :: * -> *) a x.
Forest p v a -> Rep (Forest p v a) x
Generic)

instance (NFData (v a)) ⇒ NFData (Forest p v a)

instance ToJSON (v a) ⇒ ToJSON (Forest p v a)

instance FromJSON (v a) ⇒ FromJSON (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@.
--
-- Siblings span trees in the forest!

forestWith ∷ (VG.Vector v a) ⇒ (forall a . [T.Tree a] → [a]) → [T.Tree a] → Forest (p∷TreeOrder) v a
forestWith :: forall (v :: * -> *) a (p :: TreeOrder).
Vector v a =>
(forall a. [Tree a] -> [a]) -> [Tree a] -> Forest p v a
forestWith forall a. [Tree a] -> [a]
f [Tree a]
ts
  = Forest { label :: v a
label    = forall (v :: * -> *) a. Vector v a => [a] -> v a
VG.fromList forall a b. (a -> b) -> a -> b
$ forall a. [Tree a] -> [a]
f [Tree a]
ts
           , parent :: Vector Int
parent   = forall a. Unbox a => [a] -> Vector a
VU.fromList forall a b. (a -> b) -> a -> b
$ forall a b. (a -> b) -> [a] -> [b]
map (\(Int
_,Int
k,[Int]
_ ,a
_) -> Int
k             ) forall a b. (a -> b) -> a -> b
$ forall a. [Tree a] -> [a]
f [Tree (Int, Int, [Int], a)]
pcs
           , children :: Vector (Vector Int)
children = forall a. [a] -> Vector a
V.fromList  forall a b. (a -> b) -> a -> b
$ forall a b. (a -> b) -> [a] -> [b]
map (\(Int
_,Int
_,[Int]
cs,a
_) -> forall a. Unbox a => [a] -> Vector a
VU.fromList [Int]
cs) forall a b. (a -> b) -> a -> b
$ forall a. [Tree a] -> [a]
f [Tree (Int, Int, [Int], a)]
pcs
           , lsib :: Vector Int
lsib     = forall a. Unbox a => [a] -> Vector a
VU.fromList forall a b. (a -> b) -> a -> b
$ forall a b. (a -> b) -> [a] -> [b]
map forall a b. (a, b) -> a
fst forall a b. (a -> b) -> a -> b
$ forall k a. Map k a -> [a]
S.elems Map Int (Int, Int)
lr
           , rsib :: Vector Int
rsib     = forall a. Unbox a => [a] -> Vector a
VU.fromList forall a b. (a -> b) -> a -> b
$ forall a b. (a -> b) -> [a] -> [b]
map forall a b. (a, b) -> b
snd forall a b. (a -> b) -> a -> b
$ forall k a. Map k a -> [a]
S.elems Map Int (Int, Int)
lr
           , roots :: Vector Int
roots    = forall a. Unbox a => [a] -> Vector a
VU.fromList forall a b. (a -> b) -> a -> b
$ forall a b. (a -> b) -> [a] -> [b]
map (forall a b. (a, b) -> a
fst forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a. Tree a -> a
T.rootLabel) [Tree (Int, a)]
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 :: [Tree Int]
ps = forall a. Int -> [Tree a] -> [Tree Int]
addIndicesF' Int
0 [Tree a]
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 :: Vector Int
backp = forall a. Unbox a => [a] -> Vector a
VU.fromList forall a b. (a -> b) -> a -> b
$ forall a b. (a -> b) -> [a] -> [b]
map forall a b. (a, b) -> b
snd forall a b. (a -> b) -> a -> b
$ forall a. Ord a => [a] -> [a]
sort forall a b. (a -> b) -> a -> b
$ forall a b. [a] -> [b] -> [(a, b)]
zip (forall a. [Tree a] -> [a]
f [Tree Int]
ps) [Int
0..]
    -- Step 3: decorate the forest with indices in the correct order. Keep
    -- the label in @snd@.
    us :: [Tree (Int, a)]
us = forall a b. (a -> b) -> [a] -> [b]
map (forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (\(Int
k,a
l) -> (Vector Int
backp forall (v :: * -> *) a.
(HasCallStack, Vector v a) =>
v a -> Int -> a
VG.! Int
k,a
l))) forall a b. (a -> b) -> a -> b
$ forall a. Int -> [Tree a] -> [Tree (Int, a)]
addIndicesF Int
0 [Tree a]
ts
    -- Step 4: add the correct relations (children, lrSibling, parents)
    pcs :: [Tree (Int, Int, [Int], a)]
pcs = forall a. Int -> [Tree (Int, a)] -> [Tree (Int, Int, [Int], a)]
parentChildrenF (-Int
1) [Tree (Int, a)]
us
    -- A map with the left and right sibling
    lr :: Map Int (Int, Int)
lr  = forall a. [Tree (Int, a)] -> Map Int (Int, Int)
lrSiblingF [Tree (Int, a)]
us



-- | Construct a pre-ordered forest.

forestPre ∷ (VG.Vector v a) ⇒ [T.Tree a] → Forest Pre v a
forestPre :: forall (v :: * -> *) a. Vector v a => [Tree a] -> Forest 'Pre v a
forestPre = forall (v :: * -> *) a (p :: TreeOrder).
Vector v a =>
(forall a. [Tree a] -> [a]) -> [Tree a] -> Forest p v a
forestWith forall a. [Tree a] -> [a]
preorderF

-- | Construct a post-ordered forest.

forestPost ∷ (VG.Vector v a) ⇒ [T.Tree a] → Forest Post v a
forestPost :: forall (v :: * -> *) a. Vector v a => [Tree a] -> Forest 'Post v a
forestPost = forall (v :: * -> *) a (p :: TreeOrder).
Vector v a =>
(forall a. [Tree a] -> [a]) -> [Tree a] -> Forest p v a
forestWith forall a. [Tree a] -> [a]
postorderF

-- | Add @pre-ordered@ @(!)@ indices. First argument is the starting index.

addIndices ∷ Int → T.Tree a → T.Tree (Int,a)
addIndices :: forall a. Int -> Tree a -> Tree (Int, a)
addIndices Int
k = forall a b. (a, b) -> b
snd forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (t :: * -> *) s a b.
Traversable t =>
(s -> a -> (s, b)) -> s -> t a -> (s, t b)
mapAccumL (\Int
i a
e -> (Int
iforall a. Num a => a -> a -> a
+Int
1, (Int
i,a
e))) Int
k

-- | Add @pre-ordered@ @(!)@ indices, but to a forest.

addIndicesF ∷ Int → [T.Tree a] → [T.Tree (Int,a)]
addIndicesF :: forall a. Int -> [Tree a] -> [Tree (Int, a)]
addIndicesF Int
k = forall a b. (a, b) -> b
snd forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (t :: * -> *) s a b.
Traversable t =>
(s -> a -> (s, b)) -> s -> t a -> (s, t b)
mapAccumL forall {b}. Int -> Tree b -> (Int, Tree (Int, b))
go Int
k
  where go :: Int -> Tree b -> (Int, Tree (Int, b))
go = forall (t :: * -> *) s a b.
Traversable t =>
(s -> a -> (s, b)) -> s -> t a -> (s, t b)
mapAccumL (\Int
i b
e -> (Int
iforall a. Num a => a -> a -> a
+Int
1, (Int
i,b
e)))

-- | Add @pre-ordered@ @(!)@ indices to a forest, but throw the label away as
-- well.

addIndicesF' ∷ Int → [T.Tree a] → [T.Tree Int]
addIndicesF' :: forall a. Int -> [Tree a] -> [Tree Int]
addIndicesF' Int
k = forall a b. (a, b) -> b
snd forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (t :: * -> *) s a b.
Traversable t =>
(s -> a -> (s, b)) -> s -> t a -> (s, t b)
mapAccumL forall {a}. Int -> Tree a -> (Int, Tree Int)
go Int
k
  where go :: Int -> Tree a -> (Int, Tree Int)
go = forall (t :: * -> *) s a b.
Traversable t =>
(s -> a -> (s, b)) -> s -> t a -> (s, t b)
mapAccumL (\Int
i a
e -> (Int
iforall a. Num a => a -> a -> a
+Int
1, Int
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 :: forall a. Int -> [Tree (Int, a)] -> [Tree (Int, Int, [Int], a)]
parentChildrenF Int
k [Tree (Int, a)]
ts = [ forall a. a -> [Tree a] -> Tree a
T.Node (Int
i,Int
k,forall {b} {b}. [Tree (b, b)] -> [b]
children [Tree (Int, a)]
sf,a
l) (forall a. Int -> [Tree (Int, a)] -> [Tree (Int, Int, [Int], a)]
parentChildrenF Int
i [Tree (Int, a)]
sf)  | T.Node (Int
i,a
l) [Tree (Int, a)]
sf <- [Tree (Int, a)]
ts ]
  where children :: [Tree (b, b)] -> [b]
children [Tree (b, b)]
sf = forall a b. (a -> b) -> [a] -> [b]
map (forall a b. (a, b) -> a
fst forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a. Tree a -> a
T.rootLabel) [Tree (b, b)]
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 :: forall a. [Tree (Int, a)] -> Map Int (Int, Int)
lrSiblingF = forall k a. Ord k => k -> Map k a -> Map k a
S.delete (-Int
1) forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a. Tree (Int, a) -> Map Int (Int, Int)
lrSibling forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a. a -> [Tree a] -> Tree a
T.Node (-Int
1,forall a. HasCallStack => String -> a
error String
"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 :: forall a. Tree (Int, a) -> Map Int (Int, Int)
lrSibling = forall k a. Ord k => [(k, a)] -> Map k a
S.fromList forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a b. (a -> b) -> [a] -> [b]
map forall {b} {b}. (Num b, Eq b) => (b, b, [b]) -> (b, (b, b))
splt forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a. Tree a -> [a]
T.flatten forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall {a} {b}. [a] -> Tree (a, b) -> Tree (a, b, [a])
go ([]::[Int])
  where go :: [a] -> Tree (a, b) -> Tree (a, b, [a])
go [a]
sib (T.Node (a
k,b
lbl) [Tree (a, b)]
frst) = let cs :: [a]
cs = [a
l | T.Node (a
l,b
_) [Tree (a, b)]
_ <- [Tree (a, b)]
frst] in forall a. a -> [Tree a] -> Tree a
T.Node (a
k,b
lbl,[a]
sib) [ [a] -> Tree (a, b) -> Tree (a, b, [a])
go [a]
cs Tree (a, b)
t | Tree (a, b)
t <- [Tree (a, b)]
frst]
        splt :: (b, b, [b]) -> (b, (b, b))
splt (b
k,b
_,[])  = (b
k,(-b
1,-b
1))
        splt (b
k,b
_,[b]
sbl) = let ([b]
ls,[b]
rs) = forall a. (a -> Bool) -> [a] -> ([a], [a])
span (forall a. Eq a => a -> a -> Bool
/=b
k) [b]
sbl in (b
k,(forall a. [a] -> a
last forall a b. (a -> b) -> a -> b
$ (-b
1)forall a. a -> [a] -> [a]
:[b]
ls,forall a. [a] -> a
head forall a b. (a -> b) -> a -> b
$ forall a. [a] -> [a]
tail [b]
rs forall a. [a] -> [a] -> [a]
++ [-b
1]))

-- | Return the left-most leaf for each node.

leftMostLeaves ∷ Forest p v a → VU.Vector Int
leftMostLeaves :: forall (p :: TreeOrder) (v :: * -> *) a. Forest p v a -> Vector Int
leftMostLeaves Forest p v a
f = forall (v :: * -> *) a b.
(Vector v a, Vector v b) =>
(a -> b) -> v a -> v b
VG.map (forall (p :: TreeOrder) (v :: * -> *) a. Forest p v a -> Int -> Int
leftMostLeaf Forest p v a
f) forall a b. (a -> b) -> a -> b
$ forall (v :: * -> *) a. (Vector v a, Num a) => a -> Int -> v a
VG.enumFromN Int
0 forall a b. (a -> b) -> a -> b
$ forall (v :: * -> *) a. Vector v a => v a -> Int
VG.length forall a b. (a -> b) -> a -> b
$ forall (p :: TreeOrder) (v :: * -> *) a. Forest p v a -> Vector Int
parent Forest p v a
f

-- | Just the leaf-most leaf for a certain node.

leftMostLeaf ∷ Forest p v a → Int → Int
leftMostLeaf :: forall (p :: TreeOrder) (v :: * -> *) a. Forest p v a -> Int -> Int
leftMostLeaf Forest p v a
f = Int -> Int
go
  where go :: Int -> Int
go Int
k = let cs :: Vector Int
cs = forall (p :: TreeOrder) (v :: * -> *) a.
Forest p v a -> Vector (Vector Int)
children Forest p v a
f forall (v :: * -> *) a.
(HasCallStack, Vector v a) =>
v a -> Int -> a
VG.! Int
k
               in if forall (v :: * -> *) a. Vector v a => v a -> Bool
VG.null Vector Int
cs then Int
k else Int -> Int
go (forall (v :: * -> *) a. Vector v a => v a -> a
VG.head Vector Int
cs)

-- | Return the right-most leaf for each node.

rightMostLeaves ∷ Forest p v a → VU.Vector Int
rightMostLeaves :: forall (p :: TreeOrder) (v :: * -> *) a. Forest p v a -> Vector Int
rightMostLeaves Forest p v a
f = forall (v :: * -> *) a b.
(Vector v a, Vector v b) =>
(a -> b) -> v a -> v b
VG.map (forall (p :: TreeOrder) (v :: * -> *) a. Forest p v a -> Int -> Int
rightMostLeaf Forest p v a
f) forall a b. (a -> b) -> a -> b
$ forall (v :: * -> *) a. (Vector v a, Num a) => a -> Int -> v a
VG.enumFromN Int
0 forall a b. (a -> b) -> a -> b
$ forall (v :: * -> *) a. Vector v a => v a -> Int
VG.length forall a b. (a -> b) -> a -> b
$ forall (p :: TreeOrder) (v :: * -> *) a. Forest p v a -> Vector Int
parent Forest p v a
f

-- | Given a tree, and a node index, return the right-most leaf for the
-- node.

rightMostLeaf ∷ Forest p v a → Int → Int
rightMostLeaf :: forall (p :: TreeOrder) (v :: * -> *) a. Forest p v a -> Int -> Int
rightMostLeaf Forest p v a
f = Int -> Int
go
  where go :: Int -> Int
go Int
k = let cs :: Vector Int
cs = forall (p :: TreeOrder) (v :: * -> *) a.
Forest p v a -> Vector (Vector Int)
children Forest p v a
f forall (v :: * -> *) a.
(HasCallStack, Vector v a) =>
v a -> Int -> a
VG.! Int
k
               in  if forall (v :: * -> *) a. Vector v a => v a -> Bool
VG.null Vector Int
cs then Int
k else Int -> Int
go (forall (v :: * -> *) a. Vector v a => v a -> a
VG.last Vector Int
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 :: forall (v :: * -> *) a. Forest 'Post v a -> Vector Int
leftKeyRoots Forest 'Post v a
f = forall a. Unbox a => [a] -> Vector a
VU.fromList forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a. Ord a => [a] -> [a]
sort forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall k a. Map k a -> [a]
S.elems forall a b. (a -> b) -> a -> b
$ forall b a. Unbox b => (a -> b -> a) -> a -> Vector b -> a
VU.foldl' Map Int Int -> Int -> Map Int Int
go forall k a. Map k a
S.empty (forall a. (Unbox a, Num a) => a -> Int -> Vector a
VU.enumFromN (Int
0::Int) forall a b. (a -> b) -> a -> b
$ forall (v :: * -> *) a. Vector v a => v a -> Int
VG.length forall a b. (a -> b) -> a -> b
$ forall (p :: TreeOrder) (v :: * -> *) a. Forest p v a -> Vector Int
parent Forest 'Post v a
f)
        -- Build a map from left-most leaf to most root-near node.
  where go :: Map Int Int -> Int -> Map Int Int
go Map Int Int
s Int
k = forall k a. Ord k => (a -> a -> a) -> k -> a -> Map k a -> Map k a
S.insertWith forall a. Ord a => a -> a -> a
max (Vector Int
lml forall a. Unbox a => Vector a -> Int -> a
VU.! Int
k) Int
k Map Int Int
s
        lml :: Vector Int
lml  = forall (p :: TreeOrder) (v :: * -> *) a. Forest p v a -> Vector Int
leftMostLeaves Forest 'Post v a
f

-- | Returns the list of all sorted subsets of subforests in the forest.
-- If the forest is given in pre-order, then The subsets are returned in
-- reversed pre-order.
--
-- TODO turn this into @newtype vectors@ that enforce @size >= 1@.

sortedSubForests ∷ Forest p v a → [VU.Vector Int]
sortedSubForests :: forall (p :: TreeOrder) (v :: * -> *) a.
Forest p v a -> [Vector Int]
sortedSubForests Forest p v a
f =
  -- cleanup
  forall a b. (a -> b) -> [a] -> [b]
map forall a. Unbox a => [a] -> Vector a
VU.fromList
  forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a. Eq a => [a] -> [a]
L.nub     -- TODO revise later, is in @O(n^2)@
  forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (t :: * -> *) a. Foldable t => t [a] -> [a]
concat
  -- make sure that in our partial order we have smaller forests come
  -- first.
  forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a b. (a -> b) -> [a] -> [b]
map (forall a b. (a -> b) -> [a] -> [b]
map Srt -> [Int]
unSrt forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a. Set a -> [a]
Set.toList forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a. Ord a => [a] -> Set a
Set.fromList forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a b. (a -> b) -> [a] -> [b]
map [Int] -> Srt
Srt)
  -- get all nonempty ordered subforests
  forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a b. (a -> b) -> [a] -> [b]
map (forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
concatMap (forall a. [a] -> [a]
L.tail forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a. [a] -> [[a]]
L.subsequences))
  forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a b. (a -> b) -> [a] -> [b]
map (forall a. [a] -> [[a]]
L.permutations)
  forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a b. (a -> b) -> [a] -> [b]
map forall (v :: * -> *) a. Vector v a => v a -> [a]
VG.toList forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (v :: * -> *) a. Vector v a => v a -> [a]
VG.toList
  -- only nodes with children
  forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (v :: * -> *) a. Vector v a => (a -> Bool) -> v a -> v a
VG.filter (Bool -> Bool
not forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (v :: * -> *) a. Vector v a => v a -> Bool
VG.null)
  -- every node that has children in reverse order
  -- make sure that the roots are there, but come last
  forall a b. (a -> b) -> a -> b
$ forall (v :: * -> *) a. Vector v a => v a -> a -> v a
VG.snoc (forall (v :: * -> *) a. Vector v a => v a -> v a
VG.reverse (forall (p :: TreeOrder) (v :: * -> *) a.
Forest p v a -> Vector (Vector Int)
children Forest p v a
f)) (forall (p :: TreeOrder) (v :: * -> *) a. Forest p v a -> Vector Int
roots Forest p v a
f)

newtype Srt = Srt { Srt -> [Int]
unSrt ∷ [Int] }
  deriving (Srt -> Srt -> Bool
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
/= :: Srt -> Srt -> Bool
$c/= :: Srt -> Srt -> Bool
== :: Srt -> Srt -> Bool
$c== :: Srt -> Srt -> Bool
Eq,Int -> Srt -> ShowS
[Srt] -> ShowS
Srt -> String
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
showList :: [Srt] -> ShowS
$cshowList :: [Srt] -> ShowS
show :: Srt -> String
$cshow :: Srt -> String
showsPrec :: Int -> Srt -> ShowS
$cshowsPrec :: Int -> Srt -> ShowS
Show)

instance Ord Srt where
  Srt [Int]
xs <= :: Srt -> Srt -> Bool
<= Srt [Int]
ys = forall (t :: * -> *) a. Foldable t => t a -> Int
length [Int]
xs forall a. Ord a => a -> a -> Bool
<= forall (t :: * -> *) a. Foldable t => t a -> Int
length [Int]
ys

-- | Given a forest, return the list of trees that constitue the forest.

forestToTrees ∷ (VG.Vector v a) ⇒ Forest p v a → T.Forest a
forestToTrees :: forall (v :: * -> *) a (p :: TreeOrder).
Vector v a =>
Forest p v a -> Forest a
forestToTrees Forest{v a
Vector Int
Vector (Vector Int)
roots :: Vector Int
rsib :: Vector Int
lsib :: Vector Int
children :: Vector (Vector Int)
parent :: Vector Int
label :: v a
roots :: forall (p :: TreeOrder) (v :: * -> *) a. Forest p v a -> Vector Int
rsib :: forall (p :: TreeOrder) (v :: * -> *) a. Forest p v a -> Vector Int
lsib :: forall (p :: TreeOrder) (v :: * -> *) a. Forest p v a -> Vector Int
children :: forall (p :: TreeOrder) (v :: * -> *) a.
Forest p v a -> Vector (Vector Int)
parent :: forall (p :: TreeOrder) (v :: * -> *) a. Forest p v a -> Vector Int
label :: forall (p :: TreeOrder) (v :: * -> *) a. Forest p v a -> v a
..} = forall a b. (a -> b) -> [a] -> [b]
map Int -> Tree a
getTree forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (v :: * -> *) a. Vector v a => v a -> [a]
VG.toList forall a b. (a -> b) -> a -> b
$ Vector Int
roots
  where getTree :: Int -> Tree a
getTree Int
k = forall a. a -> [Tree a] -> Tree a
T.Node (v a
label forall (v :: * -> *) a.
(HasCallStack, Vector v a) =>
v a -> Int -> a
VG.! Int
k) (forall a b. (a -> b) -> [a] -> [b]
map Int -> Tree a
getTree forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (v :: * -> *) a. Vector v a => v a -> [a]
VG.toList forall a b. (a -> b) -> a -> b
$ Vector (Vector Int)
children forall (v :: * -> *) a.
(HasCallStack, Vector v a) =>
v a -> Int -> a
VG.! Int
k)



-- * QuickCheck

-- | Wrapped quickcheck instance for 'T.Tree'.

newtype QCTree a = QCTree { forall a. QCTree a -> Tree a
getTree ∷ T.Tree a }
  deriving (Int -> QCTree a -> ShowS
forall a. Show a => Int -> QCTree a -> ShowS
forall a. Show a => [QCTree a] -> ShowS
forall a. Show a => QCTree a -> String
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
showList :: [QCTree a] -> ShowS
$cshowList :: forall a. Show a => [QCTree a] -> ShowS
show :: QCTree a -> String
$cshow :: forall a. Show a => QCTree a -> String
showsPrec :: Int -> QCTree a -> ShowS
$cshowsPrec :: forall a. Show a => Int -> QCTree a -> ShowS
Show)

instance (Arbitrary a) ⇒ Arbitrary (QCTree a) where
  arbitrary :: Gen (QCTree a)
arbitrary =
    let go :: Gen (Tree a)
go = forall a. (Int -> Gen a) -> Gen a
sized forall a b. (a -> b) -> a -> b
$ \Int
n →
               do a
val ← forall a. Arbitrary a => Gen a
arbitrary
                  let n' :: Int
n' = Int
n forall a. Integral a => a -> a -> a
`div` Int
2
                  [Tree a]
nodes ← if Int
n' forall a. Ord a => a -> a -> Bool
> Int
0
                    then do Int
k ← forall a. Random a => (a, a) -> Gen a
choose (Int
0,Int
n')
                            forall a. Int -> Gen a -> Gen a
resize Int
n' forall a b. (a -> b) -> a -> b
$ forall (m :: * -> *) a. Applicative m => Int -> m a -> m [a]
replicateM Int
k (forall a. QCTree a -> Tree a
getTree forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall a. Arbitrary a => Gen a
arbitrary)
                    else forall (m :: * -> *) a. Monad m => a -> m a
return []
                  forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ forall a. a -> [Tree a] -> Tree a
T.Node a
val [Tree a]
nodes
    in  forall a. Tree a -> QCTree a
QCTree forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Gen (Tree a)
go
  shrink :: QCTree a -> [QCTree a]
shrink (QCTree (T.Node a
val [Tree a]
forest)) =
    [] -- [ QCTree $ T.Node v f | v <- shrink val, f <- map (map getTree) $ shrink $ map QCTree forest ]

--  -- * Test functions
--  
--  test1 :: [T.Tree Char]
--  test1 = [T.Node 'R' [T.Node 'a' [], T.Node 'b' []], T.Node 'S' [T.Node 'x' [], T.Node 'y' []]]
--  
--  test2 :: [T.Tree Char]
--  test2 = [T.Node 'R' [T.Node 'a' [], T.Node 'b' [], T.Node 'c' []]]
--  
--  runtest t = do
--    print (forestPre t :: Forest Pre V.Vector Char)
--    print (forestPost t :: 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)
--    print (sortedSubForests (forestPre t :: Forest Pre V.Vector Char))
--