-- 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 = (Tree a -> [(a, a)]) -> [Tree a] -> [(a, a)]
forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
concatMap ((Tree a -> [(a, a)]) -> [Tree a] -> [(a, a)])
-> (Tree a -> [(a, a)]) -> [Tree a] -> [(a, a)]
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] [(a, a)] -> [(a, a)] -> [(a, a)]
forall a. [a] -> [a] -> [a]
++ [Tree 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 = (Tree a -> [(Tree a, [Tree a])])
-> [Tree a] -> [(Tree a, [Tree a])]
forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
concatMap ([Tree a] -> Tree a -> [(Tree a, [Tree a])]
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) (Tree a, [Tree a]) -> [(Tree a, [Tree a])] -> [(Tree a, [Tree a])]
forall a. a -> [a] -> [a]
: (Tree a -> [(Tree a, [Tree a])])
-> [Tree a] -> [(Tree a, [Tree a])]
forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
concatMap ([Tree a] -> Tree a -> [(Tree a, [Tree a])]
go (Tree a
treeTree a -> [Tree a] -> [Tree a]
forall 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 = ((Tree a, [Tree a]) -> (a, [a], [a]))
-> [(Tree a, [Tree a])] -> [(a, [a], [a])]
forall a b. (a -> b) -> [a] -> [b]
map (Tree a, [Tree a]) -> (a, [a], [a])
forall {b} {b}. (Tree b, [Tree b]) -> (b, [b], [b])
flatten ([(Tree a, [Tree a])] -> [(a, [a], [a])])
-> ([Tree a] -> [(Tree a, [Tree a])])
-> [Tree a]
-> [(a, [a], [a])]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [Tree a] -> [(Tree a, [Tree a])]
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, (Tree b -> b) -> [Tree b] -> [b]
forall a b. (a -> b) -> [a] -> [b]
map Tree b -> b
forall a. Tree a -> a
Tree.rootLabel [Tree b]
parents, (Tree b -> [b]) -> [Tree b] -> [b]
forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
concatMap Tree b -> [b]
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 (Tree a -> Bool) -> [Tree a] -> Maybe (Tree a)
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Maybe a
List.find (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) [Tree a]
trees of
    Maybe (Tree a)
Nothing -> [Maybe (Tree a)] -> Maybe (Tree a)
forall (t :: * -> *) (m :: * -> *) a.
(Foldable t, MonadPlus m) =>
t (m a) -> m a
msum ([Maybe (Tree a)] -> Maybe (Tree a))
-> [Maybe (Tree a)] -> Maybe (Tree a)
forall a b. (a -> b) -> a -> b
$ (Tree a -> Maybe (Tree a)) -> [Tree a] -> [Maybe (Tree a)]
forall a b. (a -> b) -> [a] -> [b]
map ((a -> Bool) -> [Tree a] -> Maybe (Tree a)
forall a. (a -> Bool) -> [Tree a] -> Maybe (Tree a)
find a -> Bool
p ([Tree a] -> Maybe (Tree a))
-> (Tree a -> [Tree a]) -> Tree a -> Maybe (Tree a)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Tree a -> [Tree a]
forall a. Tree a -> [Tree a]
subForest) [Tree a]
trees
    Just Tree a
tree -> Tree a -> Maybe (Tree a)
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 = (Tree a -> [Tree a]) -> [Tree a] -> [Tree a]
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 = [a -> [Tree a] -> Tree a
forall a. a -> [Tree a] -> Tree a
Node a
a [Tree a]
trees]
        | Bool
otherwise = (Tree a -> [Tree a]) -> [Tree a] -> [Tree a]
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 = ((Tree a, [Tree a]) -> Bool)
-> [(Tree a, [Tree a])] -> Maybe (Tree a, [Tree a])
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Maybe a
List.find (a -> Bool
f (a -> Bool)
-> ((Tree a, [Tree a]) -> a) -> (Tree a, [Tree a]) -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Tree a -> a
forall a. Tree a -> a
Tree.rootLabel (Tree a -> a)
-> ((Tree a, [Tree a]) -> Tree a) -> (Tree a, [Tree a]) -> a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Tree a, [Tree a]) -> Tree a
forall a b. (a, b) -> a
fst) ([(Tree a, [Tree a])] -> Maybe (Tree a, [Tree a]))
-> ([Tree a] -> [(Tree a, [Tree a])])
-> [Tree a]
-> Maybe (Tree a, [Tree a])
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [Tree a] -> [(Tree a, [Tree a])]
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 = Tree a -> [a]
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 = Tree a -> Maybe (Tree a)
forall a. a -> Maybe a
Just (Tree a -> Maybe (Tree a)) -> Tree a -> Maybe (Tree a)
forall a b. (a -> b) -> a -> b
$ a -> [Tree a] -> Tree a
forall a. a -> [Tree a] -> Tree a
Node a
a ((Tree a -> Maybe (Tree a)) -> [Tree a] -> [Tree a]
forall a b. (a -> Maybe b) -> [a] -> [b]
mapMaybe Tree a -> Maybe (Tree a)
node [Tree a]
trees)
        | Bool
otherwise = Maybe (Tree a)
forall a. Maybe a
Nothing