module Data.Tree.BranchLeafLabel where
import Control.Monad.HT ((<=<))
import Data.Traversable as Traversable(Traversable(traverse))
import Control.Applicative (Applicative, (<*>), )
import qualified Control.Applicative as App
import qualified Control.Monad as Monad
import qualified Data.List as List
import Prelude hiding (map, mapM, )
type T i branch leaf = (i, Elem i branch leaf)
data Elem i branch leaf =
Branch branch [T i branch leaf]
| Leaf leaf
deriving (Show)
map :: (branch0 -> branch1)
-> (leaf0 -> leaf1)
-> (T i branch0 leaf0 -> T i branch1 leaf1)
map branchF leafF = fold (,) (Branch . branchF) (Leaf . leafF)
mapLabel :: (i -> j)
-> (T i branch leaf -> T j branch leaf)
mapLabel f = fold ((,) . f) Branch Leaf
mapCond ::
(branch -> Bool)
-> (branch -> branch)
-> (leaf -> leaf)
-> (T i branch leaf -> T i branch leaf)
mapCond descend branchF leafF =
let recourse =
switch (,)
(\branch -> Branch (branchF branch) .
if descend branch
then List.map recourse
else id)
(Leaf . leafF)
in recourse
mapSubTrees ::
(branch -> Bool)
-> ((branch, [T i branch leaf]) -> (branch, [T i branch leaf]))
-> (T i branch leaf -> T i branch leaf)
mapSubTrees p f =
let recourse =
switch (,)
(\branch subTrees -> uncurry Branch
(if p branch
then f (branch, subTrees)
else (branch, List.map recourse subTrees)))
Leaf
in recourse
filterBranch ::
(branch -> Bool)
-> (T i branch leaf -> [T i branch leaf])
filterBranch p =
foldLabel
(\i branch subTrees ->
let jointSubTrees = concat subTrees
in if p branch
then [(i, Branch branch jointSubTrees)]
else jointSubTrees)
(\i leaf -> [(i, Leaf leaf)])
fold :: (i -> a -> b)
-> (branch -> [b] -> a)
-> (leaf -> a)
-> (T i branch leaf -> b)
fold iF branchF leafF =
let recourse =
switch iF (\x -> branchF x . List.map recourse) leafF
in recourse
switch ::
(i -> a -> b)
-> (branch -> [T i branch leaf] -> a)
-> (leaf -> a)
-> (T i branch leaf -> b)
switch iF branchF leafF (i,n) =
iF i (switchElem branchF leafF n)
foldLabel ::
(i -> branch -> [b] -> b)
-> (i -> leaf -> b)
-> (T i branch leaf -> b)
foldLabel branchF leafF =
fold
(flip ($))
(\branch subTrees i -> branchF i branch subTrees)
(\leaf i -> leafF i leaf)
foldLabelAlt ::
(i -> branch -> [b] -> b)
-> (i -> leaf -> b)
-> (T i branch leaf -> b)
foldLabelAlt branchF leafF =
let recourse =
switchLabel (\i x -> branchF i x . List.map recourse) leafF
in recourse
switchLabel ::
(i -> branch -> [T i branch leaf] -> b)
-> (i -> leaf -> b)
-> (T i branch leaf -> b)
switchLabel branchF leafF (i,n) =
switchElem (branchF i) (leafF i) n
switchElem ::
(branch -> [T i branch leaf] -> a)
-> (leaf -> a)
-> (Elem i branch leaf -> a)
switchElem branchF _ (Branch x subTrees) =
branchF x subTrees
switchElem _ leafF (Leaf x) = leafF x
allSubTrees :: T i branch leaf -> [T i branch leaf]
allSubTrees tree =
tree :
switch (flip const) (const (concatMap allSubTrees)) (const []) tree
mapA :: Applicative m =>
(branch0 -> m branch1)
-> (leaf0 -> m leaf1)
-> (T i branch0 leaf0 -> m (T i branch1 leaf1))
mapA branchF leafF =
foldA
(App.pure . (,))
(App.liftA Branch . branchF)
(App.liftA Leaf . leafF)
mapCondA :: Applicative m =>
(branch -> Bool)
-> (branch -> m branch)
-> (leaf -> m leaf)
-> (T i branch leaf -> m (T i branch leaf))
mapCondA descend branchF leafF =
let recourse =
switch
(App.liftA . (,))
(\branch -> App.liftA2 Branch (branchF branch) .
if descend branch
then traverse recourse
else App.pure)
(App.liftA Leaf . leafF)
in recourse
foldA :: Applicative m =>
(i -> m (a -> b))
-> (branch -> m ([b] -> a))
-> (leaf -> m a)
-> (T i branch leaf -> m b)
foldA iF branchF leafF =
let recourse =
switch (\i x -> iF i <*> x)
(\x subTrees -> branchF x <*> traverse recourse subTrees)
leafF
in recourse
foldM :: Monad m =>
(i -> a -> m b)
-> (branch -> [b] -> m a)
-> (leaf -> m a)
-> (T i branch leaf -> m b)
foldM iF branchF leafF =
let recourse =
switch ((=<<) . iF)
(\x -> branchF x <=< Monad.mapM recourse)
leafF
in recourse