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 (Int -> Elem i branch leaf -> ShowS
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
forall i branch leaf.
(Show branch, Show i, Show leaf) =>
Int -> Elem i branch leaf -> ShowS
forall i branch leaf.
(Show branch, Show i, Show leaf) =>
[Elem i branch leaf] -> ShowS
forall i branch leaf.
(Show branch, Show i, Show leaf) =>
Elem i branch leaf -> String
showList :: [Elem i branch leaf] -> ShowS
$cshowList :: forall i branch leaf.
(Show branch, Show i, Show leaf) =>
[Elem i branch leaf] -> ShowS
show :: Elem i branch leaf -> String
$cshow :: forall i branch leaf.
(Show branch, Show i, Show leaf) =>
Elem i branch leaf -> String
showsPrec :: Int -> Elem i branch leaf -> ShowS
$cshowsPrec :: forall i branch leaf.
(Show branch, Show i, Show leaf) =>
Int -> Elem i branch leaf -> ShowS
Show)

map :: (branch0 -> branch1)
    -> (leaf0 -> leaf1)
    -> (T i branch0 leaf0 -> T i branch1 leaf1)
map :: forall branch0 branch1 leaf0 leaf1 i.
(branch0 -> branch1)
-> (leaf0 -> leaf1) -> T i branch0 leaf0 -> T i branch1 leaf1
map branch0 -> branch1
branchF leaf0 -> leaf1
leafF = forall i a b branch leaf.
(i -> a -> b)
-> (branch -> [b] -> a) -> (leaf -> a) -> T i branch leaf -> b
fold (,) (forall i branch leaf.
branch -> [T i branch leaf] -> Elem i branch leaf
Branch forall b c a. (b -> c) -> (a -> b) -> a -> c
. branch0 -> branch1
branchF) (forall i branch leaf. leaf -> Elem i branch leaf
Leaf forall b c a. (b -> c) -> (a -> b) -> a -> c
. leaf0 -> leaf1
leafF)

mapLabel :: (i -> j)
    -> (T i branch leaf -> T j branch leaf)
mapLabel :: forall i j branch leaf.
(i -> j) -> T i branch leaf -> T j branch leaf
mapLabel i -> j
f = forall i a b branch leaf.
(i -> a -> b)
-> (branch -> [b] -> a) -> (leaf -> a) -> T i branch leaf -> b
fold ((,) forall b c a. (b -> c) -> (a -> b) -> a -> c
. i -> j
f) forall i branch leaf.
branch -> [T i branch leaf] -> Elem i branch leaf
Branch forall i branch leaf. leaf -> Elem i branch leaf
Leaf

mapCond ::
       (branch -> Bool)
    -> (branch -> branch)
    -> (leaf -> leaf)
    -> (T i branch leaf -> T i branch leaf)
mapCond :: forall branch leaf i.
(branch -> Bool)
-> (branch -> branch)
-> (leaf -> leaf)
-> T i branch leaf
-> T i branch leaf
mapCond branch -> Bool
descend branch -> branch
branchF leaf -> leaf
leafF =
   let recourse :: T i branch leaf -> T i branch leaf
recourse =
          forall i a b branch leaf.
(i -> a -> b)
-> (branch -> [T i branch leaf] -> a)
-> (leaf -> a)
-> T i branch leaf
-> b
switch (,)
             (\branch
branch -> forall i branch leaf.
branch -> [T i branch leaf] -> Elem i branch leaf
Branch (branch -> branch
branchF branch
branch) forall b c a. (b -> c) -> (a -> b) -> a -> c
.
                if branch -> Bool
descend branch
branch
                   then forall a b. (a -> b) -> [a] -> [b]
List.map T i branch leaf -> T i branch leaf
recourse
                   else forall a. a -> a
id)
             (forall i branch leaf. leaf -> Elem i branch leaf
Leaf forall b c a. (b -> c) -> (a -> b) -> a -> c
. leaf -> leaf
leafF)
   in  forall {i}. T i branch leaf -> T i branch leaf
recourse

{- |
Process the subtrees for which the predicate holds.
If the predicate matches subtrees of a matching subtree,
the sub-subtrees are not mapped.
-}
mapSubTrees ::
       (branch -> Bool)
    -> ((branch, [T i branch leaf]) -> (branch, [T i branch leaf]))
    -> (T i branch leaf -> T i branch leaf)
mapSubTrees :: forall branch i leaf.
(branch -> Bool)
-> ((branch, [T i branch leaf]) -> (branch, [T i branch leaf]))
-> T i branch leaf
-> T i branch leaf
mapSubTrees branch -> Bool
p (branch, [T i branch leaf]) -> (branch, [T i branch leaf])
f =
   let recourse :: T i branch leaf -> T i branch leaf
recourse =
          forall i a b branch leaf.
(i -> a -> b)
-> (branch -> [T i branch leaf] -> a)
-> (leaf -> a)
-> T i branch leaf
-> b
switch (,)
             (\branch
branch [T i branch leaf]
subTrees -> forall a b c. (a -> b -> c) -> (a, b) -> c
uncurry forall i branch leaf.
branch -> [T i branch leaf] -> Elem i branch leaf
Branch
                (if branch -> Bool
p branch
branch
                   then (branch, [T i branch leaf]) -> (branch, [T i branch leaf])
f (branch
branch, [T i branch leaf]
subTrees)
                   else (branch
branch, forall a b. (a -> b) -> [a] -> [b]
List.map T i branch leaf -> T i branch leaf
recourse [T i branch leaf]
subTrees)))
             forall i branch leaf. leaf -> Elem i branch leaf
Leaf
   in  T i branch leaf -> T i branch leaf
recourse

filterBranch ::
       (branch -> Bool)
    -> (T i branch leaf -> [T i branch leaf])
filterBranch :: forall branch i leaf.
(branch -> Bool) -> T i branch leaf -> [T i branch leaf]
filterBranch branch -> Bool
p =
   forall i branch b leaf.
(i -> branch -> [b] -> b)
-> (i -> leaf -> b) -> T i branch leaf -> b
foldLabel
      (\i
i branch
branch [[T i branch leaf]]
subTrees ->
         let jointSubTrees :: [T i branch leaf]
jointSubTrees = forall (t :: * -> *) a. Foldable t => t [a] -> [a]
concat [[T i branch leaf]]
subTrees
         in  if branch -> Bool
p branch
branch
               then [(i
i, forall i branch leaf.
branch -> [T i branch leaf] -> Elem i branch leaf
Branch branch
branch [T i branch leaf]
jointSubTrees)]
               else [T i branch leaf]
jointSubTrees)
      (\i
i leaf
leaf -> [(i
i, forall i branch leaf. leaf -> Elem i branch leaf
Leaf leaf
leaf)])

fold :: (i -> a -> b)
     -> (branch -> [b] -> a)
     -> (leaf -> a)
     -> (T i branch leaf -> b)
fold :: forall i a b branch leaf.
(i -> a -> b)
-> (branch -> [b] -> a) -> (leaf -> a) -> T i branch leaf -> b
fold i -> a -> b
iF branch -> [b] -> a
branchF leaf -> a
leafF =
   let recourse :: T i branch leaf -> b
recourse =
          forall i a b branch leaf.
(i -> a -> b)
-> (branch -> [T i branch leaf] -> a)
-> (leaf -> a)
-> T i branch leaf
-> b
switch i -> a -> b
iF (\branch
x -> branch -> [b] -> a
branchF branch
x forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a b. (a -> b) -> [a] -> [b]
List.map T i branch leaf -> b
recourse) leaf -> a
leafF
   in  T i branch leaf -> b
recourse

switch ::
        (i -> a -> b)
     -> (branch -> [T i branch leaf] -> a)
     -> (leaf -> a)
     -> (T i branch leaf -> b)
switch :: forall i a b branch leaf.
(i -> a -> b)
-> (branch -> [T i branch leaf] -> a)
-> (leaf -> a)
-> T i branch leaf
-> b
switch i -> a -> b
iF branch -> [T i branch leaf] -> a
branchF leaf -> a
leafF (i
i,Elem i branch leaf
n) =
   i -> a -> b
iF i
i (forall branch i leaf a.
(branch -> [T i branch leaf] -> a)
-> (leaf -> a) -> Elem i branch leaf -> a
switchElem branch -> [T i branch leaf] -> a
branchF leaf -> a
leafF Elem i branch leaf
n)

foldLabel ::
        (i -> branch -> [b] -> b)
     -> (i -> leaf -> b)
     -> (T i branch leaf -> b)
foldLabel :: forall i branch b leaf.
(i -> branch -> [b] -> b)
-> (i -> leaf -> b) -> T i branch leaf -> b
foldLabel i -> branch -> [b] -> b
branchF i -> leaf -> b
leafF =
   forall i a b branch leaf.
(i -> a -> b)
-> (branch -> [b] -> a) -> (leaf -> a) -> T i branch leaf -> b
fold
      (forall a b c. (a -> b -> c) -> b -> a -> c
flip forall a b. (a -> b) -> a -> b
($))
      (\branch
branch [b]
subTrees i
i -> i -> branch -> [b] -> b
branchF i
i branch
branch [b]
subTrees)
      (\leaf
leaf i
i -> i -> leaf -> b
leafF i
i leaf
leaf)

foldLabelAlt ::
        (i -> branch -> [b] -> b)
     -> (i -> leaf -> b)
     -> (T i branch leaf -> b)
foldLabelAlt :: forall i branch b leaf.
(i -> branch -> [b] -> b)
-> (i -> leaf -> b) -> T i branch leaf -> b
foldLabelAlt i -> branch -> [b] -> b
branchF i -> leaf -> b
leafF =
   let recourse :: T i branch leaf -> b
recourse =
          forall i branch leaf b.
(i -> branch -> [T i branch leaf] -> b)
-> (i -> leaf -> b) -> T i branch leaf -> b
switchLabel (\i
i branch
x -> i -> branch -> [b] -> b
branchF i
i branch
x forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a b. (a -> b) -> [a] -> [b]
List.map T i branch leaf -> b
recourse) i -> leaf -> b
leafF
   in  T i branch leaf -> b
recourse

switchLabel ::
        (i -> branch -> [T i branch leaf] -> b)
     -> (i -> leaf -> b)
     -> (T i branch leaf -> b)
switchLabel :: forall i branch leaf b.
(i -> branch -> [T i branch leaf] -> b)
-> (i -> leaf -> b) -> T i branch leaf -> b
switchLabel i -> branch -> [T i branch leaf] -> b
branchF i -> leaf -> b
leafF (i
i,Elem i branch leaf
n) =
   forall branch i leaf a.
(branch -> [T i branch leaf] -> a)
-> (leaf -> a) -> Elem i branch leaf -> a
switchElem (i -> branch -> [T i branch leaf] -> b
branchF i
i) (i -> leaf -> b
leafF i
i) Elem i branch leaf
n

switchElem ::
        (branch -> [T i branch leaf] -> a)
     -> (leaf -> a)
     -> (Elem i branch leaf -> a)
switchElem :: forall branch i leaf a.
(branch -> [T i branch leaf] -> a)
-> (leaf -> a) -> Elem i branch leaf -> a
switchElem branch -> [T i branch leaf] -> a
branchF leaf -> a
_ (Branch branch
x [T i branch leaf]
subTrees) =
   branch -> [T i branch leaf] -> a
branchF branch
x [T i branch leaf]
subTrees
switchElem branch -> [T i branch leaf] -> a
_ leaf -> a
leafF (Leaf leaf
x) = leaf -> a
leafF leaf
x

allSubTrees :: T i branch leaf -> [T i branch leaf]
allSubTrees :: forall i branch leaf. T i branch leaf -> [T i branch leaf]
allSubTrees T i branch leaf
tree =
   T i branch leaf
tree forall a. a -> [a] -> [a]
:
   forall i a b branch leaf.
(i -> a -> b)
-> (branch -> [T i branch leaf] -> a)
-> (leaf -> a)
-> T i branch leaf
-> b
switch (forall a b c. (a -> b -> c) -> b -> a -> c
flip forall a b. a -> b -> a
const) (forall a b. a -> b -> a
const (forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
concatMap forall i branch leaf. T i branch leaf -> [T i branch leaf]
allSubTrees)) (forall a b. a -> b -> a
const []) T i branch leaf
tree



mapA :: Applicative m =>
       (branch0 -> m branch1)
    -> (leaf0 -> m leaf1)
    -> (T i branch0 leaf0 -> m (T i branch1 leaf1))
mapA :: forall (m :: * -> *) branch0 branch1 leaf0 leaf1 i.
Applicative m =>
(branch0 -> m branch1)
-> (leaf0 -> m leaf1) -> T i branch0 leaf0 -> m (T i branch1 leaf1)
mapA branch0 -> m branch1
branchF leaf0 -> m leaf1
leafF =
   forall (m :: * -> *) i a b branch leaf.
Applicative m =>
(i -> m (a -> b))
-> (branch -> m ([b] -> a))
-> (leaf -> m a)
-> T i branch leaf
-> m b
foldA
      (forall (f :: * -> *) a. Applicative f => a -> f a
App.pure forall b c a. (b -> c) -> (a -> b) -> a -> c
. (,))
      (forall (f :: * -> *) a b. Applicative f => (a -> b) -> f a -> f b
App.liftA forall i branch leaf.
branch -> [T i branch leaf] -> Elem i branch leaf
Branch forall b c a. (b -> c) -> (a -> b) -> a -> c
. branch0 -> m branch1
branchF)
      (forall (f :: * -> *) a b. Applicative f => (a -> b) -> f a -> f b
App.liftA forall i branch leaf. leaf -> Elem i branch leaf
Leaf forall b c a. (b -> c) -> (a -> b) -> a -> c
. leaf0 -> m leaf1
leafF)

mapCondA :: Applicative m =>
       (branch -> Bool)
    -> (branch -> m branch)
    -> (leaf -> m leaf)
    -> (T i branch leaf -> m (T i branch leaf))
mapCondA :: forall (m :: * -> *) branch leaf i.
Applicative m =>
(branch -> Bool)
-> (branch -> m branch)
-> (leaf -> m leaf)
-> T i branch leaf
-> m (T i branch leaf)
mapCondA branch -> Bool
descend branch -> m branch
branchF leaf -> m leaf
leafF =
   let recourse :: T i branch leaf -> m (T i branch leaf)
recourse =
          forall i a b branch leaf.
(i -> a -> b)
-> (branch -> [T i branch leaf] -> a)
-> (leaf -> a)
-> T i branch leaf
-> b
switch
             (forall (f :: * -> *) a b. Applicative f => (a -> b) -> f a -> f b
App.liftA forall b c a. (b -> c) -> (a -> b) -> a -> c
. (,))
             (\branch
branch -> forall (f :: * -> *) a b c.
Applicative f =>
(a -> b -> c) -> f a -> f b -> f c
App.liftA2 forall i branch leaf.
branch -> [T i branch leaf] -> Elem i branch leaf
Branch (branch -> m branch
branchF branch
branch) forall b c a. (b -> c) -> (a -> b) -> a -> c
.
                if branch -> Bool
descend branch
branch
                   then forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
traverse T i branch leaf -> m (T i branch leaf)
recourse
                   else forall (f :: * -> *) a. Applicative f => a -> f a
App.pure)
             (forall (f :: * -> *) a b. Applicative f => (a -> b) -> f a -> f b
App.liftA forall i branch leaf. leaf -> Elem i branch leaf
Leaf forall b c a. (b -> c) -> (a -> b) -> a -> c
. leaf -> m leaf
leafF)
   in  forall {i}. T i branch leaf -> m (T i branch leaf)
recourse

foldA :: Applicative m =>
        (i -> m (a -> b))
     -> (branch -> m ([b] -> a))
     -> (leaf -> m a)
     -> (T i branch leaf -> m b)
foldA :: forall (m :: * -> *) i a b branch leaf.
Applicative m =>
(i -> m (a -> b))
-> (branch -> m ([b] -> a))
-> (leaf -> m a)
-> T i branch leaf
-> m b
foldA i -> m (a -> b)
iF branch -> m ([b] -> a)
branchF leaf -> m a
leafF =
   let recourse :: T i branch leaf -> m b
recourse =
          forall i a b branch leaf.
(i -> a -> b)
-> (branch -> [T i branch leaf] -> a)
-> (leaf -> a)
-> T i branch leaf
-> b
switch (\i
i m a
x -> i -> m (a -> b)
iF i
i forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> m a
x)
             (\branch
x [T i branch leaf]
subTrees -> branch -> m ([b] -> a)
branchF branch
x forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
traverse T i branch leaf -> m b
recourse [T i branch leaf]
subTrees)
             leaf -> m a
leafF
   in  T i branch leaf -> m b
recourse

foldM :: Monad m =>
        (i -> a -> m b)
     -> (branch -> [b] -> m a)
     -> (leaf -> m a)
     -> (T i branch leaf -> m b)
foldM :: forall (m :: * -> *) i a b branch leaf.
Monad m =>
(i -> a -> m b)
-> (branch -> [b] -> m a)
-> (leaf -> m a)
-> T i branch leaf
-> m b
foldM i -> a -> m b
iF branch -> [b] -> m a
branchF leaf -> m a
leafF =
   let recourse :: T i branch leaf -> m b
recourse =
          forall i a b branch leaf.
(i -> a -> b)
-> (branch -> [T i branch leaf] -> a)
-> (leaf -> a)
-> T i branch leaf
-> b
switch (forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
(=<<) forall b c a. (b -> c) -> (a -> b) -> a -> c
. i -> a -> m b
iF)
                   (\branch
x -> branch -> [b] -> m a
branchF branch
x forall (m :: * -> *) b c a.
Monad m =>
(b -> m c) -> (a -> m b) -> a -> m c
<=< forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
Monad.mapM T i branch leaf -> m b
recourse)
                   leaf -> m a
leafF
   in  T i branch leaf -> m b
recourse