{-# LANGUAGE DeriveDataTypeable #-}
{-# LANGUAGE DeriveGeneric #-}
module ELynx.Data.Topology.Rooted
(
Topology (..),
Forest,
fromTree,
fromLabeledTree,
degree,
leaves,
flatten,
identify,
prune,
dropLeavesWith,
zipTreesWith,
zipTrees,
duplicateLeaves,
)
where
import Control.Applicative
import Control.DeepSeq
import Control.Monad
import Data.Aeson
import Data.Data
import Data.Foldable
import Data.List.NonEmpty (NonEmpty)
import qualified Data.List.NonEmpty as N
import Data.Maybe
import qualified Data.Set as S
import Data.Traversable
import qualified Data.Tree as T
import qualified ELynx.Data.Tree.Rooted as R
import GHC.Generics
singleton :: NonEmpty a -> Bool
singleton xs = 1 == length (N.take 2 xs)
data Topology a
= Node {forest :: Forest a}
| Leaf {label :: a}
deriving (Eq, Read, Show, Data, Generic)
type Forest a = NonEmpty (Topology a)
instance Functor Topology where
fmap f (Node ts) = Node $ fmap (fmap f) ts
fmap f (Leaf lb) = Leaf $ f lb
instance Foldable Topology where
foldMap f (Node ts) = foldMap (foldMap f) ts
foldMap f (Leaf lb) = f lb
null _ = False
{-# INLINE null #-}
toList = flatten
{-# INLINE toList #-}
instance Traversable Topology where
traverse g (Node ts) = Node <$> traverse (traverse g) ts
traverse g (Leaf lb) = Leaf <$> g lb
instance Applicative Topology where
pure = Leaf
(Node tsF) <*> tx = Node $ fmap (<*> tx) tsF
(Leaf lbF) <*> tx = lbF <$> tx
liftA2 f (Node tsX) ty = Node $ fmap (\tx -> liftA2 f tx ty) tsX
liftA2 f (Leaf lbX) (Node tsY) = Node $ fmap (f lbX <$>) tsY
liftA2 f (Leaf lbX) (Leaf lbY) = Leaf $ f lbX lbY
(Node tsX) *> ty@(Node tsY) = Node $ tsY <> fmap (*> ty) tsX
(Leaf _) *> (Node tsY) = Node tsY
_ *> (Leaf lbY) = Leaf lbY
(Node tsX) <* ty = Node $ fmap (<* ty) tsX
(Leaf lbX) <* _ = Leaf lbX
instance Monad Topology where
(Node ts) >>= f = Node $ fmap (>>= f) ts
(Leaf lb) >>= f = case f lb of
Node ts' -> Node ts'
Leaf lb' -> Leaf lb'
instance NFData a => NFData (Topology a) where
rnf (Node ts) = rnf ts
rnf (Leaf lb) = rnf lb
instance ToJSON a => ToJSON (Topology a)
instance FromJSON a => FromJSON (Topology a)
degree :: Topology a -> Int
degree (Node ts) = (+ 1) $ length ts
degree (Leaf _) = 1
leaves :: Ord a => Topology a -> [a]
leaves (Leaf lb) = [lb]
leaves (Node ts) = concatMap leaves ts
flatten :: Topology a -> [a]
flatten t = squish t []
where
squish (Node ts) xs = foldr squish xs ts
squish (Leaf lb) xs = lb : xs
fromTree :: T.Tree a -> Topology a
fromTree (T.Node lb []) = Leaf lb
fromTree (T.Node _ xs) = Node $ fromTree <$> N.fromList xs
fromLabeledTree :: R.Tree e a -> Topology a
fromLabeledTree (R.Node _ lb []) = Leaf lb
fromLabeledTree (R.Node _ _ xs) = Node $ fromLabeledTree <$> N.fromList xs
identify :: Traversable t => t a -> t Int
identify = snd . mapAccumL (\i _ -> (i + 1, i)) (0 :: Int)
prune :: Topology a -> Topology a
prune (Node ts)
| singleton ts = Node $ fmap prune $ forest $ N.head ts
| otherwise = Node $ fmap prune ts
prune (Leaf lb) = Leaf lb
dropLeavesWith :: (a -> Bool) -> Topology a -> Maybe (Topology a)
dropLeavesWith p (Leaf lb)
| p lb = Nothing
| otherwise = Just $ Leaf lb
dropLeavesWith p (Node ts) =
if null ts'
then Nothing
else
Just $ Node $ N.fromList ts'
where
ts' = catMaybes $ N.toList $ fmap (dropLeavesWith p) ts
zipTreesWith :: (a1 -> a2 -> a) -> Topology a1 -> Topology a2 -> Maybe (Topology a)
zipTreesWith f (Node tsL) (Node tsR) =
if N.length tsL == N.length tsR
then
zipWithM (zipTreesWith f) (N.toList tsL) (N.toList tsR) >>= Just . Node . N.fromList
else Nothing
zipTreesWith f (Leaf lbL) (Leaf lbR) = Just $ Leaf $ f lbL lbR
zipTreesWith _ _ _ = Nothing
zipTrees :: Topology a1 -> Topology a2 -> Maybe (Topology (a1, a2))
zipTrees = zipTreesWith (,)
duplicates :: Ord a => [a] -> Bool
duplicates = go S.empty
where
go _ [] = False
go seen (x : xs) = x `S.member` seen || go (S.insert x seen) xs
duplicateLeaves :: Ord a => Topology a -> Bool
duplicateLeaves = duplicates . leaves