-- | Tree-related utilities.
module Data.Tree.Util where

import           Data.Bifoldable
import           Data.Bifunctor
import           Data.Bitraversable
import           Control.Lens
import           Control.Monad ((>=>))
import qualified Data.List as List
import qualified Data.List.NonEmpty as NonEmpty
import           Data.List.NonEmpty (NonEmpty(..))
import           Data.Maybe (listToMaybe,maybeToList)
import           Data.Tree

--------------------------------------------------------------------------------

-- $setup
-- >>> :{
-- let myTree = Node 0 [ Node 1 []
--                     , Node 2 []
--                     , Node 3 [ Node 4 [] ]
--                     ]
-- :}


--------------------------------------------------------------------------------

-- | Nodes in a tree are typically either an internal node or a leaf node
data TreeNode v a = InternalNode v | LeafNode a deriving (Int -> TreeNode v a -> ShowS
[TreeNode v a] -> ShowS
TreeNode v a -> String
(Int -> TreeNode v a -> ShowS)
-> (TreeNode v a -> String)
-> ([TreeNode v a] -> ShowS)
-> Show (TreeNode v a)
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
forall v a. (Show v, Show a) => Int -> TreeNode v a -> ShowS
forall v a. (Show v, Show a) => [TreeNode v a] -> ShowS
forall v a. (Show v, Show a) => TreeNode v a -> String
showList :: [TreeNode v a] -> ShowS
$cshowList :: forall v a. (Show v, Show a) => [TreeNode v a] -> ShowS
show :: TreeNode v a -> String
$cshow :: forall v a. (Show v, Show a) => TreeNode v a -> String
showsPrec :: Int -> TreeNode v a -> ShowS
$cshowsPrec :: forall v a. (Show v, Show a) => Int -> TreeNode v a -> ShowS
Show,TreeNode v a -> TreeNode v a -> Bool
(TreeNode v a -> TreeNode v a -> Bool)
-> (TreeNode v a -> TreeNode v a -> Bool) -> Eq (TreeNode v a)
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
forall v a. (Eq v, Eq a) => TreeNode v a -> TreeNode v a -> Bool
/= :: TreeNode v a -> TreeNode v a -> Bool
$c/= :: forall v a. (Eq v, Eq a) => TreeNode v a -> TreeNode v a -> Bool
== :: TreeNode v a -> TreeNode v a -> Bool
$c== :: forall v a. (Eq v, Eq a) => TreeNode v a -> TreeNode v a -> Bool
Eq)

instance Bifunctor TreeNode where
  bimap :: (a -> b) -> (c -> d) -> TreeNode a c -> TreeNode b d
bimap = (a -> b) -> (c -> d) -> TreeNode a c -> TreeNode b d
forall (t :: * -> * -> *) a b c d.
Bitraversable t =>
(a -> b) -> (c -> d) -> t a c -> t b d
bimapDefault
instance Bifoldable TreeNode where
  bifoldMap :: (a -> m) -> (b -> m) -> TreeNode a b -> m
bifoldMap = (a -> m) -> (b -> m) -> TreeNode a b -> m
forall (t :: * -> * -> *) m a b.
(Bitraversable t, Monoid m) =>
(a -> m) -> (b -> m) -> t a b -> m
bifoldMapDefault
instance Bitraversable TreeNode where
  bitraverse :: (a -> f c) -> (b -> f d) -> TreeNode a b -> f (TreeNode c d)
bitraverse a -> f c
f b -> f d
g = \case
    InternalNode a
v -> c -> TreeNode c d
forall v a. v -> TreeNode v a
InternalNode (c -> TreeNode c d) -> f c -> f (TreeNode c d)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> a -> f c
f a
v
    LeafNode b
l     -> d -> TreeNode c d
forall v a. a -> TreeNode v a
LeafNode     (d -> TreeNode c d) -> f d -> f (TreeNode c d)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> b -> f d
g b
l

-- | A TreeNode is isomorphic to Either
_TreeNodeEither :: Iso' (TreeNode v p) (Either v p)
_TreeNodeEither :: p (Either v p) (f (Either v p))
-> p (TreeNode v p) (f (TreeNode v p))
_TreeNodeEither = (TreeNode v p -> Either v p)
-> (Either v p -> TreeNode v p)
-> Iso (TreeNode v p) (TreeNode v p) (Either v p) (Either v p)
forall s a b t. (s -> a) -> (b -> t) -> Iso s t a b
iso TreeNode v p -> Either v p
forall a b. TreeNode a b -> Either a b
tne Either v p -> TreeNode v p
forall v a. Either v a -> TreeNode v a
etn
  where
    tne :: TreeNode a b -> Either a b
tne = \case
      InternalNode a
v -> a -> Either a b
forall a b. a -> Either a b
Left a
v
      LeafNode b
l     -> b -> Either a b
forall a b. b -> Either a b
Right b
l
    etn :: Either v a -> TreeNode v a
etn = (v -> TreeNode v a)
-> (a -> TreeNode v a) -> Either v a -> TreeNode v a
forall a c b. (a -> c) -> (b -> c) -> Either a b -> c
either v -> TreeNode v a
forall v a. v -> TreeNode v a
InternalNode a -> TreeNode v a
forall v a. a -> TreeNode v a
LeafNode

--------------------------------------------------------------------------------
-- * Zipper on rose trees

-- | Zipper for rose trees
data Zipper a = Zipper { Zipper a -> Tree a
focus      :: Tree a
                       , Zipper a -> [([Tree a], a, [Tree a])]
ancestors  :: [([Tree a], a, [Tree a])] -- left siblings in reverse order
                       }
              deriving (Int -> Zipper a -> ShowS
[Zipper a] -> ShowS
Zipper a -> String
(Int -> Zipper a -> ShowS)
-> (Zipper a -> String) -> ([Zipper a] -> ShowS) -> Show (Zipper a)
forall a. Show a => Int -> Zipper a -> ShowS
forall a. Show a => [Zipper a] -> ShowS
forall a. Show a => Zipper a -> String
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
showList :: [Zipper a] -> ShowS
$cshowList :: forall a. Show a => [Zipper a] -> ShowS
show :: Zipper a -> String
$cshow :: forall a. Show a => Zipper a -> String
showsPrec :: Int -> Zipper a -> ShowS
$cshowsPrec :: forall a. Show a => Int -> Zipper a -> ShowS
Show,Zipper a -> Zipper a -> Bool
(Zipper a -> Zipper a -> Bool)
-> (Zipper a -> Zipper a -> Bool) -> Eq (Zipper a)
forall a. Eq a => Zipper a -> Zipper a -> Bool
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
/= :: Zipper a -> Zipper a -> Bool
$c/= :: forall a. Eq a => Zipper a -> Zipper a -> Bool
== :: Zipper a -> Zipper a -> Bool
$c== :: forall a. Eq a => Zipper a -> Zipper a -> Bool
Eq)

-- | Create a new zipper focussiong on the root.
root :: Tree a -> Zipper a
root :: Tree a -> Zipper a
root = (Tree a -> [([Tree a], a, [Tree a])] -> Zipper a)
-> [([Tree a], a, [Tree a])] -> Tree a -> Zipper a
forall a b c. (a -> b -> c) -> b -> a -> c
flip Tree a -> [([Tree a], a, [Tree a])] -> Zipper a
forall a. Tree a -> [([Tree a], a, [Tree a])] -> Zipper a
Zipper []

-- | Move the focus to the parent of this node.
up               :: Zipper a -> Maybe (Zipper a)
up :: Zipper a -> Maybe (Zipper a)
up (Zipper Tree a
t [([Tree a], a, [Tree a])]
as) = case [([Tree a], a, [Tree a])]
as of
                     []              -> Maybe (Zipper a)
forall a. Maybe a
Nothing
                     (([Tree a]
ls,a
p,[Tree a]
rs):[([Tree a], a, [Tree a])]
as') -> Zipper a -> Maybe (Zipper a)
forall a. a -> Maybe a
Just (Zipper a -> Maybe (Zipper a)) -> Zipper a -> Maybe (Zipper a)
forall a b. (a -> b) -> a -> b
$ Tree a -> [([Tree a], a, [Tree a])] -> Zipper a
forall a. Tree a -> [([Tree a], a, [Tree a])] -> Zipper a
Zipper (a -> [Tree a] -> Tree a
forall a. a -> Forest a -> Tree a
Node a
p ([Tree a] -> [Tree a]
forall a. [a] -> [a]
reverse [Tree a]
ls [Tree a] -> [Tree a] -> [Tree a]
forall a. Semigroup a => a -> a -> a
<> [Tree a
t] [Tree a] -> [Tree a] -> [Tree a]
forall a. Semigroup a => a -> a -> a
<> [Tree a]
rs)) [([Tree a], a, [Tree a])]
as'

-- | Move the focus to the first child of this node.
--
-- >>> firstChild $ root myTree
-- Just (Zipper {focus = Node {rootLabel = 1, subForest = []}, ancestors = [([],0,[Node {rootLabel = 2, subForest = []},Node {rootLabel = 3, subForest = [Node {rootLabel = 4, subForest = []}]}])]})
firstChild                          :: Zipper a -> Maybe (Zipper a)
firstChild :: Zipper a -> Maybe (Zipper a)
firstChild (Zipper (Node a
x Forest a
chs) [(Forest a, a, Forest a)]
as) = case Forest a
chs of
                                        []       -> Maybe (Zipper a)
forall a. Maybe a
Nothing
                                        (Tree a
c:Forest a
chs') -> Zipper a -> Maybe (Zipper a)
forall a. a -> Maybe a
Just (Zipper a -> Maybe (Zipper a)) -> Zipper a -> Maybe (Zipper a)
forall a b. (a -> b) -> a -> b
$ Tree a -> [(Forest a, a, Forest a)] -> Zipper a
forall a. Tree a -> [([Tree a], a, [Tree a])] -> Zipper a
Zipper Tree a
c (([],a
x,Forest a
chs')(Forest a, a, Forest a)
-> [(Forest a, a, Forest a)] -> [(Forest a, a, Forest a)]
forall a. a -> [a] -> [a]
:[(Forest a, a, Forest a)]
as)

-- | Move the focus to the next sibling of this node
--
-- >>> (firstChild $ root myTree) >>= nextSibling
-- Just (Zipper {focus = Node {rootLabel = 2, subForest = []}, ancestors = [([Node {rootLabel = 1, subForest = []}],0,[Node {rootLabel = 3, subForest = [Node {rootLabel = 4, subForest = []}]}])]})
nextSibling               :: Zipper a -> Maybe (Zipper a)
nextSibling :: Zipper a -> Maybe (Zipper a)
nextSibling (Zipper Tree a
t [([Tree a], a, [Tree a])]
as) = case [([Tree a], a, [Tree a])]
as of
                              []                -> Maybe (Zipper a)
forall a. Maybe a
Nothing -- no parent
                              (([Tree a]
_,a
_,[]):[([Tree a], a, [Tree a])]
_)      -> Maybe (Zipper a)
forall a. Maybe a
Nothing -- no next sibling
                              (([Tree a]
ls,a
p,Tree a
r:[Tree a]
rs):[([Tree a], a, [Tree a])]
as') -> Zipper a -> Maybe (Zipper a)
forall a. a -> Maybe a
Just (Zipper a -> Maybe (Zipper a)) -> Zipper a -> Maybe (Zipper a)
forall a b. (a -> b) -> a -> b
$ Tree a -> [([Tree a], a, [Tree a])] -> Zipper a
forall a. Tree a -> [([Tree a], a, [Tree a])] -> Zipper a
Zipper Tree a
r ((Tree a
tTree a -> [Tree a] -> [Tree a]
forall a. a -> [a] -> [a]
:[Tree a]
ls,a
p,[Tree a]
rs)([Tree a], a, [Tree a])
-> [([Tree a], a, [Tree a])] -> [([Tree a], a, [Tree a])]
forall a. a -> [a] -> [a]
:[([Tree a], a, [Tree a])]
as')

-- | Move the focus to the next sibling of this node
prevSibling               :: Zipper a -> Maybe (Zipper a)
prevSibling :: Zipper a -> Maybe (Zipper a)
prevSibling (Zipper Tree a
t [([Tree a], a, [Tree a])]
as) = case [([Tree a], a, [Tree a])]
as of
                              []                -> Maybe (Zipper a)
forall a. Maybe a
Nothing -- no parent
                              (([],a
_,[Tree a]
_):[([Tree a], a, [Tree a])]
_)      -> Maybe (Zipper a)
forall a. Maybe a
Nothing -- no prev sibling
                              ((Tree a
l:[Tree a]
ls,a
p,[Tree a]
rs):[([Tree a], a, [Tree a])]
as') -> Zipper a -> Maybe (Zipper a)
forall a. a -> Maybe a
Just (Zipper a -> Maybe (Zipper a)) -> Zipper a -> Maybe (Zipper a)
forall a b. (a -> b) -> a -> b
$ Tree a -> [([Tree a], a, [Tree a])] -> Zipper a
forall a. Tree a -> [([Tree a], a, [Tree a])] -> Zipper a
Zipper Tree a
l (([Tree a]
ls,a
p,Tree a
tTree a -> [Tree a] -> [Tree a]
forall a. a -> [a] -> [a]
:[Tree a]
rs)([Tree a], a, [Tree a])
-> [([Tree a], a, [Tree a])] -> [([Tree a], a, [Tree a])]
forall a. a -> [a] -> [a]
:[([Tree a], a, [Tree a])]
as')

-- | Given a zipper that focussses on some subtree t, construct a list with
-- zippers that focus on each child.
allChildren :: Zipper a -> [Zipper a]
allChildren :: Zipper a -> [Zipper a]
allChildren = (Maybe (Zipper a) -> Maybe (Zipper a, Maybe (Zipper a)))
-> Maybe (Zipper a) -> [Zipper a]
forall b a. (b -> Maybe (a, b)) -> b -> [a]
List.unfoldr ((\Zipper a
ch -> (Zipper a
ch, Zipper a -> Maybe (Zipper a)
forall a. Zipper a -> Maybe (Zipper a)
nextSibling Zipper a
ch)) (Zipper a -> (Zipper a, Maybe (Zipper a)))
-> Maybe (Zipper a) -> Maybe (Zipper a, Maybe (Zipper a))
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$>) (Maybe (Zipper a) -> [Zipper a])
-> (Zipper a -> Maybe (Zipper a)) -> Zipper a -> [Zipper a]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Zipper a -> Maybe (Zipper a)
forall a. Zipper a -> Maybe (Zipper a)
firstChild

-- | Given a zipper that focussses on some subtree t, construct a list with
-- zippers that focus on each of the nodes in the subtree of t.
allTrees   :: Zipper a -> [Zipper a]
allTrees :: Zipper a -> [Zipper a]
allTrees Zipper a
r = Zipper a
r Zipper a -> [Zipper a] -> [Zipper a]
forall a. a -> [a] -> [a]
: (Zipper a -> [Zipper a]) -> [Zipper a] -> [Zipper a]
forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
concatMap Zipper a -> [Zipper a]
forall a. Zipper a -> [Zipper a]
allTrees (Zipper a -> [Zipper a]
forall a. Zipper a -> [Zipper a]
allChildren Zipper a
r)

-- | Creates a new tree from the zipper that thas the current node as root. The
-- ancestorTree (if there is any) forms the first child in this new root.
unZipperLocal                          :: Zipper a -> Tree a
unZipperLocal :: Zipper a -> Tree a
unZipperLocal (Zipper (Node a
x Forest a
chs) [(Forest a, a, Forest a)]
as) = a -> Forest a -> Tree a
forall a. a -> Forest a -> Tree a
Node a
x (Maybe (Tree a) -> Forest a
forall a. Maybe a -> [a]
maybeToList ([(Forest a, a, Forest a)] -> Maybe (Tree a)
forall a. [([Tree a], a, [Tree a])] -> Maybe (Tree a)
constructTree [(Forest a, a, Forest a)]
as) Forest a -> Forest a -> Forest a
forall a. Semigroup a => a -> a -> a
<> Forest a
chs)

-- | Constructs a tree from the list of ancestors (if there are any)
constructTree :: [([Tree a],a,[Tree a])] -> Maybe (Tree a)
constructTree :: [([Tree a], a, [Tree a])] -> Maybe (Tree a)
constructTree = [Tree a] -> Maybe (Tree a)
forall a. [a] -> Maybe a
listToMaybe
              ([Tree a] -> Maybe (Tree a))
-> ([([Tree a], a, [Tree a])] -> [Tree a])
-> [([Tree a], a, [Tree a])]
-> Maybe (Tree a)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (([Tree a], a, [Tree a]) -> [Tree a] -> [Tree a])
-> [Tree a] -> [([Tree a], a, [Tree a])] -> [Tree a]
forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr (\([Tree a]
ls,a
p,[Tree a]
rs) [Tree a]
tas -> [a -> [Tree a] -> Tree a
forall a. a -> Forest a -> Tree a
Node a
p ([Tree a]
tas [Tree a] -> [Tree a] -> [Tree a]
forall a. Semigroup a => a -> a -> a
<> [Tree a] -> [Tree a]
forall a. [a] -> [a]
reverse [Tree a]
ls [Tree a] -> [Tree a] -> [Tree a]
forall a. Semigroup a => a -> a -> a
<> [Tree a]
rs)]) []


--------------------------------------------------------------------------------

-- | Given a predicate on an element, find a node that matches the predicate, and turn that
-- node into the root of the tree.
--
-- running time: \(O(nT)\) where \(n\) is the size of the tree, and \(T\) is
-- the time to evaluate a predicate.
--
-- >>> findEvert (== 4) myTree
-- Just (Node {rootLabel = 4, subForest = [Node {rootLabel = 3, subForest = [Node {rootLabel = 0, subForest = [Node {rootLabel = 1, subForest = []},Node {rootLabel = 2, subForest = []}]}]}]})
-- >>> findEvert (== 5) myTree
-- Nothing
findEvert   :: (a -> Bool) -> Tree a -> Maybe (Tree a)
findEvert :: (a -> Bool) -> Tree a -> Maybe (Tree a)
findEvert a -> Bool
p = (Tree a -> Bool) -> Tree a -> Maybe (Tree a)
forall a. (Tree a -> Bool) -> Tree a -> Maybe (Tree a)
findEvert' (a -> Bool
p (a -> Bool) -> (Tree a -> a) -> Tree a -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Tree a -> a
forall a. Tree a -> a
rootLabel)

-- | Given a predicate matching on a subtree, find a node that matches the predicate, and turn that
-- node into the root of the tree.
--
-- running time: \(O(nT(n))\) where \(n\) is the size of the tree, and \(T(m)\) is
-- the time to evaluate a predicate on a subtree of size \(m\).
findEvert'   :: (Tree a -> Bool) -> Tree a -> Maybe (Tree a)
findEvert' :: (Tree a -> Bool) -> Tree a -> Maybe (Tree a)
findEvert' Tree a -> Bool
p = (Zipper a -> Tree a) -> Maybe (Zipper a) -> Maybe (Tree a)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap Zipper a -> Tree a
forall a. Zipper a -> Tree a
unZipperLocal (Maybe (Zipper a) -> Maybe (Tree a))
-> (Tree a -> Maybe (Zipper a)) -> Tree a -> Maybe (Tree a)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Zipper a -> Bool) -> [Zipper a] -> Maybe (Zipper a)
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Maybe a
List.find (Tree a -> Bool
p (Tree a -> Bool) -> (Zipper a -> Tree a) -> Zipper a -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Zipper a -> Tree a
forall a. Zipper a -> Tree a
focus) ([Zipper a] -> Maybe (Zipper a))
-> (Tree a -> [Zipper a]) -> Tree a -> Maybe (Zipper a)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Zipper a -> [Zipper a]
forall a. Zipper a -> [Zipper a]
allTrees (Zipper a -> [Zipper a])
-> (Tree a -> Zipper a) -> Tree a -> [Zipper a]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Tree a -> Zipper a
forall a. Tree a -> Zipper a
root

-- | Function to extract a path between a start node and an end node (if such a
--path exists). If there are multiple paths, no guarantees are given about
--which one is returned.
--
-- running time: \(O(n(T_p+T_s)\), where \(n\) is the size of the tree, and
-- \(T_p\) and \(T_s\) are the times it takes to evaluate the @isStartingNode@
-- and @isEndingNode@ predicates.
--
--
-- >>> findPath (== 1) (==4) myTree
-- Just [1,0,3,4]
-- >>>  findPath (== 1) (==2) myTree
-- Just [1,0,2]
-- >>>  findPath (== 1) (==1) myTree
-- Just [1]
-- >>>  findPath (== 1) (==2) myTree
-- Just [1,0,2]
-- >>>  findPath (== 4) (==2) myTree
-- Just [4,3,0,2]
findPath               :: (a -> Bool) -- ^ is this node a starting node
                          -> (a -> Bool) -- ^ is this node an ending node
                          -> Tree a -> Maybe [a]
findPath :: (a -> Bool) -> (a -> Bool) -> Tree a -> Maybe [a]
findPath a -> Bool
isStart a -> Bool
isEnd = (a -> Bool) -> Tree a -> Maybe (Tree a)
forall a. (a -> Bool) -> Tree a -> Maybe (Tree a)
findEvert a -> Bool
isStart (Tree a -> Maybe (Tree a))
-> (Tree a -> Maybe [a]) -> Tree a -> Maybe [a]
forall (m :: * -> *) a b c.
Monad m =>
(a -> m b) -> (b -> m c) -> a -> m c
>=> (a -> Bool) -> Tree a -> Maybe [a]
forall a. (a -> Bool) -> Tree a -> Maybe [a]
findNode a -> Bool
isEnd

-- | Given a predicate on a, find (the path to) a node that satisfies the predicate.
--
-- >>> findNode (== 4) myTree
-- Just [0,3,4]
findNode   :: (a -> Bool) -> Tree a -> Maybe [a]
findNode :: (a -> Bool) -> Tree a -> Maybe [a]
findNode a -> Bool
p = [[a]] -> Maybe [a]
forall a. [a] -> Maybe a
listToMaybe ([[a]] -> Maybe [a]) -> (Tree a -> [[a]]) -> Tree a -> Maybe [a]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Tree a -> Bool) -> Tree a -> [[a]]
forall a. (Tree a -> Bool) -> Tree a -> [[a]]
findNodes (a -> Bool
p (a -> Bool) -> (Tree a -> a) -> Tree a -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Tree a -> a
forall a. Tree a -> a
rootLabel)

-- | Find all paths to nodes that satisfy the predicate
--
-- running time: \(O(nT(n))\) where \(n\) is the size of the tree, and \(T(m)\) is
-- the time to evaluate a predicate on a subtree of size \(m\).
--
-- >>> findNodes ((< 4) . rootLabel) myTree
-- [[0],[0,1],[0,2],[0,3]]
-- >>> findNodes (even . rootLabel) myTree
-- [[0],[0,2],[0,3,4]]
-- >>> let size = length in findNodes ((> 1) . size) myTree
-- [[0],[0,3]]
findNodes   :: (Tree a -> Bool) -> Tree a -> [[a]]
findNodes :: (Tree a -> Bool) -> Tree a -> [[a]]
findNodes Tree a -> Bool
p = Tree a -> [[a]]
go
  where
    go :: Tree a -> [[a]]
go Tree a
t = let mh :: [[a]]
mh = [ [] | Tree a -> Bool
p Tree a
t ] -- [[]] iff 'p t'
           in ([a] -> [a]) -> [[a]] -> [[a]]
forall a b. (a -> b) -> [a] -> [b]
map (Tree a -> a
forall a. Tree a -> a
rootLabel Tree a
ta -> [a] -> [a]
forall a. a -> [a] -> [a]
:) ([[a]] -> [[a]]) -> [[a]] -> [[a]]
forall a b. (a -> b) -> a -> b
$ [[a]]
mh [[a]] -> [[a]] -> [[a]]
forall a. Semigroup a => a -> a -> a
<> (Tree a -> [[a]]) -> [Tree a] -> [[a]]
forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
concatMap Tree a -> [[a]]
go (Tree a -> [Tree a]
forall a. Plated a => a -> [a]
children Tree a
t)


-- | BFS Traversal of the rose tree that decomposes it into levels.
--
-- running time: \(O(n)\)
levels :: Tree a -> NonEmpty (NonEmpty a)
levels :: Tree a -> NonEmpty (NonEmpty a)
levels = NonEmpty (Tree a) -> NonEmpty (NonEmpty a)
forall a. NonEmpty (Tree a) -> NonEmpty (NonEmpty a)
go1 (NonEmpty (Tree a) -> NonEmpty (NonEmpty a))
-> (Tree a -> NonEmpty (Tree a)) -> Tree a -> NonEmpty (NonEmpty a)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Tree a -> [Tree a] -> NonEmpty (Tree a)
forall a. a -> [a] -> NonEmpty a
:| [])
  where
    go0   :: [Tree a] -> [NonEmpty a]
    go0 :: [Tree a] -> [NonEmpty a]
go0 [Tree a]
q = case [Tree a] -> Maybe (NonEmpty (Tree a))
forall a. [a] -> Maybe (NonEmpty a)
NonEmpty.nonEmpty [Tree a]
q of
              Maybe (NonEmpty (Tree a))
Nothing -> []
              Just NonEmpty (Tree a)
q1 -> NonEmpty (NonEmpty a) -> [NonEmpty a]
forall a. NonEmpty a -> [a]
NonEmpty.toList (NonEmpty (NonEmpty a) -> [NonEmpty a])
-> NonEmpty (NonEmpty a) -> [NonEmpty a]
forall a b. (a -> b) -> a -> b
$ NonEmpty (Tree a) -> NonEmpty (NonEmpty a)
forall a. NonEmpty (Tree a) -> NonEmpty (NonEmpty a)
go1 NonEmpty (Tree a)
q1
    {-# INLINE go0 #-}

    -- all work essentially happens here: given a bunch of trees whose
    -- root elements all have the same level, extract the values
    -- stored at these root nodes, collect all children in a big list,
    -- and explore those recursively.
    go1    :: NonEmpty (Tree a) -> NonEmpty (NonEmpty a)
    go1 :: NonEmpty (Tree a) -> NonEmpty (NonEmpty a)
go1 NonEmpty (Tree a)
qs = (Tree a -> a) -> NonEmpty (Tree a) -> NonEmpty a
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap Tree a -> a
forall a. Tree a -> a
root' NonEmpty (Tree a)
qs NonEmpty a -> [NonEmpty a] -> NonEmpty (NonEmpty a)
forall a. a -> [a] -> NonEmpty a
:| [Tree a] -> [NonEmpty a]
forall a. [Tree a] -> [NonEmpty a]
go0 ((Tree a -> [Tree a]) -> NonEmpty (Tree a) -> [Tree a]
forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
concatMap Tree a -> [Tree a]
forall a. Tree a -> Forest a
children' NonEmpty (Tree a)
qs)
    {-# INLINE go1 #-}

    root' :: Tree a -> a
root'     (Node a
x Forest a
_)   = a
x
    children' :: Tree a -> Forest a
children' (Node a
_ Forest a
chs) = Forest a
chs