-- Copyright 2013 Evan Laforge
-- This program is distributed under the terms of the GNU General Public
-- License 3.0, see COPYING or http://www.gnu.org/licenses/gpl-3.0.txt

-- | Some more utilities for "Data.Tree".
module Util.Trees (
    edges
    , paths, flatPaths
    , find
    , findAll
    , findWithParents
    , leaves
    , filter
) where
import           Prelude hiding (filter)
import qualified Data.Foldable as Foldable
import qualified Data.List as List
import           Data.Maybe (mapMaybe)
import qualified Data.Tree as Tree
import           Data.Tree (Tree(..))

import           Control.Monad


-- | The edges of a forest, as (parent, child).
edges :: [Tree a] -> [(a, a)]
edges :: forall a. [Tree a] -> [(a, a)]
edges = forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
concatMap forall a b. (a -> b) -> a -> b
$ \(Node a
val [Tree a]
subs) ->
    [(a
val, a
sub_val) | Node a
sub_val [Tree a]
_ <- [Tree a]
subs] forall a. [a] -> [a] -> [a]
++ forall a. [Tree a] -> [(a, a)]
edges [Tree a]
subs

-- | Get every element along with its parents.  The parents are closest first
-- root last.
paths :: [Tree a] -> [(Tree a, [Tree a])]
paths :: forall a. [Tree a] -> [(Tree a, [Tree a])]
paths = forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
concatMap (forall {a}. [Tree a] -> Tree a -> [(Tree a, [Tree a])]
go [])
    where
    go :: [Tree a] -> Tree a -> [(Tree a, [Tree a])]
go [Tree a]
parents tree :: Tree a
tree@(Node a
_ [Tree a]
subs) =
        (Tree a
tree, [Tree a]
parents) forall a. a -> [a] -> [a]
: forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
concatMap ([Tree a] -> Tree a -> [(Tree a, [Tree a])]
go (Tree a
treeforall a. a -> [a] -> [a]
:[Tree a]
parents)) [Tree a]
subs

-- | Like 'paths' but the parents and children have been flattened.
flatPaths :: [Tree a] -> [(a, [a], [a])]
    -- ^ (element, parents, children).  Parents are closest first root last,
    -- children are in depth first order.
flatPaths :: forall a. [Tree a] -> [(a, [a], [a])]
flatPaths = forall a b. (a -> b) -> [a] -> [b]
map forall {b} {b}. (Tree b, [Tree b]) -> (b, [b], [b])
flatten forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a. [Tree a] -> [(Tree a, [Tree a])]
paths
    where
    flatten :: (Tree b, [Tree b]) -> (b, [b], [b])
flatten (Node b
val [Tree b]
subs, [Tree b]
parents) =
        (b
val, forall a b. (a -> b) -> [a] -> [b]
map forall a. Tree a -> a
Tree.rootLabel [Tree b]
parents, forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
concatMap forall a. Tree a -> [a]
Tree.flatten [Tree b]
subs)

-- | Given a predicate, return the first depthwise matching element and
-- its children.
find :: (a -> Bool) -> [Tree a] -> Maybe (Tree a)
find :: forall a. (a -> Bool) -> [Tree a] -> Maybe (Tree a)
find a -> Bool
p [Tree a]
trees = case forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Maybe a
List.find (a -> Bool
p forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a. Tree a -> a
rootLabel) [Tree a]
trees of
    Maybe (Tree a)
Nothing -> forall (t :: * -> *) (m :: * -> *) a.
(Foldable t, MonadPlus m) =>
t (m a) -> m a
msum forall a b. (a -> b) -> a -> b
$ forall a b. (a -> b) -> [a] -> [b]
map (forall a. (a -> Bool) -> [Tree a] -> Maybe (Tree a)
find a -> Bool
p forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a. Tree a -> [Tree a]
subForest) [Tree a]
trees
    Just Tree a
tree -> forall a. a -> Maybe a
Just Tree a
tree

-- | Return the unmodified subtrees that match the predicate.  This is like
-- 'find', except that it returns all matching subtrees instead of the first
-- one.
findAll :: (a -> Bool) -> [Tree a] -> [Tree a]
findAll :: forall a. (a -> Bool) -> [Tree a] -> [Tree a]
findAll a -> Bool
f = forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
concatMap Tree a -> [Tree a]
node
    where
    node :: Tree a -> [Tree a]
node (Node a
a [Tree a]
trees)
        | a -> Bool
f a
a = [forall a. a -> [Tree a] -> Tree a
Node a
a [Tree a]
trees]
        | Bool
otherwise = forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
concatMap Tree a -> [Tree a]
node [Tree a]
trees

-- | Like 'paths', but return only the element that matches the predicate.
findWithParents :: (a -> Bool) -> [Tree a] -> Maybe (Tree a, [Tree a])
findWithParents :: forall a. (a -> Bool) -> [Tree a] -> Maybe (Tree a, [Tree a])
findWithParents a -> Bool
f = forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Maybe a
List.find (a -> Bool
f forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a. Tree a -> a
Tree.rootLabel forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a b. (a, b) -> a
fst) forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a. [Tree a] -> [(Tree a, [Tree a])]
paths

-- | Get all leaves.  The list will never be null.
leaves :: Tree a -> [a]
leaves :: forall a. Tree a -> [a]
leaves = forall (t :: * -> *) a. Foldable t => t a -> [a]
Foldable.toList

-- | Filter out nodes that don't match.
filter :: (a -> Bool) -> Tree a -> Maybe (Tree a)
filter :: forall a. (a -> Bool) -> Tree a -> Maybe (Tree a)
filter a -> Bool
f = Tree a -> Maybe (Tree a)
node
    where
    node :: Tree a -> Maybe (Tree a)
node (Node a
a [Tree a]
trees)
        | a -> Bool
f a
a = forall a. a -> Maybe a
Just forall a b. (a -> b) -> a -> b
$ forall a. a -> [Tree a] -> Tree a
Node a
a (forall a b. (a -> Maybe b) -> [a] -> [b]
mapMaybe Tree a -> Maybe (Tree a)
node [Tree a]
trees)
        | Bool
otherwise = forall a. Maybe a
Nothing