{-# LANGUAGE DeriveGeneric #-}
module TreeExtra 
  where
import GHC.Generics (Generic)
import Control.DeepSeq
import Data.Foldable
data FTree a
    = FBranch (a -> a) (FTree a)
    | FBranches [FTree a]
    | FLeaf a
instance Foldable FTree where
    foldMap g (FBranch f t) = foldMap (g . f) t
    foldMap g (FBranches ts) = mconcat $ map (foldMap g) ts
    foldMap g (FLeaf x) = g x
    {-# INLINE foldMap #-}
instance Functor FTree where
    fmap g (FBranch f t) = fmap (g . f) t
    fmap f (FBranches ts) = FBranches $ fmap (fmap f) ts
    fmap f (FLeaf x) = FLeaf (f x)
    {-# INLINE fmap #-}

{- Tree with values at and only at the leaves. -}
data LTree a
    = LBranches [LTree a]
    | LLeaf a
  deriving (Generic)
instance NFData a => NFData (LTree a)
instance Foldable LTree where
    foldMap g (LBranches ts) = mconcat $ map (foldMap g) ts
    foldMap g (LLeaf a) = g a
    {-# INLINE foldMap #-}
    foldr _ x (LBranches []) = x
    --foldr g x (LBranches (t:ts)) = foldr g x t `seq` foldr g (foldr g x t) (LBranches ts)
    foldr g x (LBranches (t:ts)) = foldr g (foldr g x t) (LBranches ts)
    --foldr g x (LBranches (t:ts)) = foldr g (foldr g x (LBranches ts)) t
    foldr g x (LLeaf y) = g y x
    {-# INLINE foldr #-}
    foldl' _ x (LBranches []) = x
    foldl' g x (LBranches (t:ts)) = foldl' g x t `seq` foldl' g (foldl' g x t) (LBranches ts)
    foldl' g x (LLeaf y) = g x y
    {-# INLINE foldl' #-}
instance Functor LTree where
    fmap f (LBranches ts) = LBranches $ fmap (fmap f) ts
    fmap f (LLeaf x) = LLeaf (f x)
    {-# INLINE fmap #-}

data RTree a b
    = RBranches a [RTree a b]
    | RLeaf b
instance Foldable (RTree a) where
    foldMap g (RBranches _ ts) = mconcat $ map (foldMap g) ts
    foldMap g (RLeaf x)      = g x
    {-# INLINE foldMap #-}
instance Functor (RTree a) where
    fmap f (RBranches i ts) = RBranches i $ fmap (fmap f) ts
    fmap f (RLeaf x)      = RLeaf (f x)