module TestBench.LabelTree where
import Data.Maybe (mapMaybe)
type Depth = Int
data LabelTree a = Leaf !Depth a
| Branch !Depth String [LabelTree a]
deriving (Eq, Ord, Show, Read, Functor)
foldLTree :: (Depth -> String -> [a] -> a) -> (Depth -> b -> a) -> LabelTree b -> a
foldLTree br lf = go
where
go tr = case tr of
Leaf d b -> lf d b
Branch d str trs -> br d str (map go trs)
mapMaybeTree :: (a -> Maybe b) -> LabelTree a -> Maybe (LabelTree b)
mapMaybeTree f = go
where
go tr = case tr of
Leaf d a -> Leaf d <$> f a
Branch d l trs -> case mapMaybe go trs of
[] -> Nothing
trs' -> Just (Branch d l trs')
mapMaybeForest :: (a -> Maybe b) -> (Depth -> String -> [b] -> b) -> [LabelTree a] -> [b]
mapMaybeForest f br = mapMaybe (fmap (foldLTree br (flip const)) . mapMaybeTree f)
leaves :: LabelTree a -> [a]
leaves = foldLTree (\_ _ lss -> concat lss) (\_ l -> [l])