-- 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