Tree.hs

{-#LANGUAGE DeriveFunctor, FlexibleInstances, TypeFamilies #-}

module Tree
  ( TreeLike(..)
  , TreeF(..)
  , Tree(..)
  , Trie
  , toTrie
  , nthLevel
  , levels
  , levels'
  , levels''
  , concatTree
  , concatNonEmpty
  ) where

import Control.Arrow ((***), (&&&))
import Control.Comonad (Comonad(..))
import Data.Functor.Foldable (Base, Fix(..), Foldable(..), Prim(..),
  Unfoldable(..))
import Data.List (unfoldr)
import Data.Monoid ((<>), Monoid(..), mconcat, Sum(..))
import Data.Bifunctor (Bifunctor(..))
import Data.Bifunctor.Flip (Flip(..))
import Data.List.NonEmpty (NonEmpty(..), fromList)
import Data.Semigroup (First(..), sconcat)

unFix :: Fix f -> f (Fix f)
unFix (Fix f) = f

-- | Every @Comonad@ gives us @extract@, allowing us to get the value at the
-- \"root\". Adding the function 'children', which returns descendants of the
-- root, characterizes a 'TreeLike' structure.
class Comonad t => TreeLike t where
  children :: t a -> [t a]

-- | A one-dimensional 'TreeLike' structure.
instance TreeLike (NonEmpty) where
  children (_ :| []) = []
  children (_ :| as) = [fromList as]

data instance Prim (NonEmpty a) b = Cons' a [b] | Nil'
  deriving (Eq, Show)

instance Functor (Prim (NonEmpty a)) where
  fmap f (Cons' a bs) = Cons' a (map f bs)
  fmap _ Nil' = Nil'

type instance Base (NonEmpty a) = Prim (NonEmpty a)

instance Foldable (NonEmpty a) where
  project (a :| []) = Cons' a []
  project (a :| as) = Cons' a [fromList as]

-- | The base @Functor@ of a 'Tree'.
newtype TreeF a b = TreeF { unTreeF :: (a, [b]) }
  deriving (Eq, Functor, Show)

instance Bifunctor TreeF where
  bimap f s = TreeF . bimap f (map s) . unTreeF

instance Comonad (Flip TreeF b) where
  extract = fst . unTreeF . runFlip
  extend f t = fmap (const $ f t) t

-- | A rose 'Tree' defined in terms of the fixed point (@Fix@) of its base
-- @Functor@, 'TreeF'.
newtype Tree a = Tree { unTree :: Fix (TreeF a) }
  deriving (Eq, Show)

instance Functor Tree where
  fmap f = Tree . Fix . bimap f (unTree . fmap f . Tree) . unFix . unTree

instance TreeLike Tree where
  children = map Tree . snd . unTreeF . unFix . unTree

instance Comonad Tree where
  extract = extract . Flip . unFix . unTree
  extend f t = Tree (Fix (TreeF (f t, map (unTree . extend f) $ children t)))

type instance Base (Tree a) = TreeF a

instance Foldable (Tree a) where
  project = fmap Tree . unFix . unTree

-- | A 'Trie' defined in terms of 'Tree'.
newtype Trie a b = Trie { unTrie :: Tree (a, Maybe b) }
  deriving (Eq, Show, Functor)

instance Bifunctor Trie where
  bimap f s = Trie . Tree . Fix . TreeF
            . bimap (bimap f (fmap s))
                    (map (unTree . unTrie . bimap f s . Trie . Tree))
            . unTreeF . unFix . unTree . unTrie

-- | @extract@ gives us the root or leftmost leaf.
instance Comonad (Trie a) where
  extract (Trie (Tree (Fix (TreeF ((_, Just b),    _))))) = b
  extract (Trie (Tree (Fix (TreeF ((_, Nothing), mus))))) =
    getFirst . sconcat . fromList $ map (First . extract . Trie . Tree) mus
  extend f t = fmap (const $ f t) t

-- | @extract@ gives us the edge label.
instance Comonad (Flip Trie b) where
  extract = fst . extract . unTrie . runFlip
  extend f t@(Flip (Trie (Tree (Fix (TreeF ((a, b), mus)))))) =
    Flip (Trie (Tree (Fix (TreeF ((f t, b),
      map (unTree . unTrie . runFlip . extend f) $ children t)))))

-- | This instance seperates levels by leaves.
instance TreeLike (Trie a) where
  children = map Trie . children . unTrie

-- | This instance separates levels by edges.
instance TreeLike (Flip Trie b) where
  children (Flip (Trie (Tree (Fix (TreeF ((_, Just b),  mus)))))) =
      map (Flip . Trie . Tree) mus
  children (Flip (Trie (Tree (Fix (TreeF ((_, Nothing), mus)))))) =
      map (Flip . Trie . Tree) $ concatMap go mus
    where
      go (Fix (TreeF ((_, Just _),    _))) = []
      go (Fix (TreeF ((_, Nothing), mus))) = concatMap go mus

toTrie :: String -> Trie Char String
toTrie str = Trie . Tree $ toTrie' str
  where
    toTrie' [c] = Fix (TreeF ((c, Just str), []))
    toTrie' (c:cs) = Fix (TreeF ((c, Nothing), [toTrie' cs]))

-- ...

split :: (Comonad f, TreeLike f) => [f a] -> ([a], [f a])
split = second concat . unzip . map (extract &&& children)

-- | Return the nth level of a 'TreeLike' structure.
nthLevel :: (Comonad f, TreeLike f) => f a -> Integer -> [a]
nthLevel = go . (return . extract &&& children)
  where
    go (as, _) 0 = as
    go (as, []) n = []
    go (as, ts) n = go (split ts) (n-1)

-- | Return the levels of a 'TreeLike' structure.
levels :: (Comonad f, TreeLike f) => f a -> [[a]]
levels = go . (return . extract &&& children)
  where
    go (as, []) = [as]
    go (as, ts) = as : go (split ts)

-- | 'levels' rewritten in terms of @unfoldr@.
levels' :: (Comonad f, TreeLike f) => f a -> [[a]]
levels' =
    unfoldr (go . second concat . unzip . map (extract &&& children)) . return
  where
    go ([], _) = Nothing
    go pair    = Just pair

-- | 'levels' rewritten as an anamorphism (in terms of @ana@).
levels'' :: (Comonad f, TreeLike f) => f a -> [[a]]
levels'' =
    ana (go . second concat . unzip . map (extract &&& children)) . return
  where
    go ([], _) = Nil
    go pair    = uncurry Cons $ pair

-- | Concatenate a 'Tree' of @Monoid@s (written in terms of @cata@).
concatTree :: Monoid m => Tree m -> m
concatTree = cata (uncurry (<>) . second mconcat . unTreeF)

-- | Concatenate a 'NonEmpty' list of @Monoid@s (written in terms of @cata@).
concatNonEmpty :: Monoid m => NonEmpty m -> m
concatNonEmpty = cata fn
  where
    fn (Cons' a as) = a <> mconcat as
    fn Nil' = mempty

tree :: Tree Int
tree = Tree (Fix (TreeF (1, [ Fix (TreeF (2, [ Fix (TreeF (3, []))
                                             , Fix (TreeF (4, []))
                                             ]
                                         )
                                  )
                            , Fix (TreeF (5, [ Fix (TreeF (6, []))
                                             , Fix (TreeF (7, [ Fix (TreeF (8, [])) ]))
                                             , Fix (TreeF (9, [ Fix (TreeF (10, [])) ]))
                                             ]
                                         )
                                  )
                            ]
                        )
                 )
            )