Haskell Data.Rose

Rose.hs

{-# LANGUAGE RebindableSyntax, NoMonomorphismRestriction, ConstraintKinds #-}
{-# LANGUAGE DeriveDataTypeable, StandaloneDeriving #-}
{-# LANGUAGE Safe #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE UndecidableInstances #-}
{-# LANGUAGE ViewPatterns #-}
{-# LANGUAGE ScopedTypeVariables #-}

-----------------------------------------------------------------------------
-- |
-- Module      :  Data.Rose
-- License     :  GPLv3 (see the file src/LICENSE)
--
-- Maintainer  :  [email protected]
-- Stability   :  experimental
-- Portability :  GHC
--
-- Generalized rose trees, called Roses.
--
-----------------------------------------------------------------------------

module Data.Rose(
    Rose(Rose, rootLabel, subBush), Bush,
    -- * Two-dimensional drawing
    drawRose, drawBush,
    -- * Extraction
    flatten, levels,
    -- * Building trees
    unfoldRose, unfoldBush,
    unfoldRoseM, unfoldBushM,
    unfoldRoseM_BF, unfoldBushM_BF,
    ) where

import YAPP

import qualified Data.Functor (Functor (..))
import qualified Data.Functor.Apply (Apply (..))
import qualified Data.Functor.Bind (Bind (..))
import qualified Control.Applicative (Applicative (..))
import qualified Control.Monad (Monad (..))
import qualified Data.Semigroup (Monoid (mappend))
import Data.Sequence (Seq, empty, singleton, (<|), (|>), fromList,
            ViewL(..), ViewR(..), viewl, viewr)
--import Data.Foldable (Foldable(foldMap), toList)
--import Data.Traversable (Traversable(traverse))
import Data.Typeable
import Control.DeepSeq (NFData(rnf))

import Data.Data (Data)

-- | Multi-way trees, also known as /rose trees/.
data Rose t a = Rose {
        rootLabel :: a,       -- ^ label value
        subBush :: Bush t a   -- ^ zero or more child trees
    }
type Bush t a = t (Rose t a)
deriving instance (Eq   a, Eq   (Bush t a)) => Eq   (Rose t a)
deriving instance (Read a, Read (Bush t a)) => Read (Rose t a)
deriving instance (Show a, Show (Bush t a)) => Show (Rose t a)
deriving instance Typeable Rose
deriving instance (Data a, Typeable t, Data (Bush t a)) => Data (Rose t a)

instance Functor f => Functor (Rose f) where
    fmap f (Rose x ts) = Rose (f x) (map (map f) ts)

instance (Apply f, Alt f) => Apply (Rose f) where
    Rose f tfs <.> tx@(Rose x txs) =
        Rose (f x) ((map f) <$> txs <|> ((<*> tx) <$> tfs))

instance (Apply f, Plus f) => Applicative (Rose f) where
    pure x = Rose x mzero 
    (<*>) = (<*>)

instance (Apply f, Alt f) => Bind (Rose f) where
    Rose x ts >>- f = Rose x' (ts' <|> ((>>= f) <$> ts))
        where Rose x' ts' = f x

instance (Apply m, Plus m) => Monad (Rose m) where
    return = return
    (>>=) = (>>=)

instance (Semigroup a, Semigroup (Bush m a)) => Semigroup (Rose m a) where
    (Rose x tx) <> (Rose y ty) = Rose (x ++ y) (tx ++ ty)

instance (Monoid a, Monoid (Bush m a)) => Monoid (Rose m a) where
    mempty = Rose mempty mempty
    mappend = (~++~)

instance (Functor t, Traversable t) => Traversable (Rose t) where
    traverse ((WrapApplicative .) -> f) (Rose x ts) = unwrapApplicative $
        Rose <$> f x <*> traverse (traverse f) ts

instance Foldable t => Foldable (Rose t) where
    foldMap f (Rose x ts) = f x ~++~ foldMap (foldMap f) ts

instance (NFData a, NFData (Bush t a)) => NFData (Rose t a) where
    rnf (Rose x ts) = rnf x `seq` rnf ts

-- | Neat 2-dimensional drawing of a tree.
drawRose :: (Foldable t, Show a) => Rose t a -> String
drawRose r = unlines $ (draw r :: [String])

-- | Neat 2-dimensional drawing of a forest.
drawBush :: (Foldable t, Functor t, Show a) => Bush t a -> String
drawBush  = unlines . map drawRose

-- In these type signatures, here be dragons. Fight them at your own peril.
draw (Rose x ts0) = return (show x) ++ drawSubRoses ts0
    where
    drawSubRoses = foldMap (\t ->return "|" ++ shift "+- " "|  " (draw t))

    shift _     _    (length -> 0) = mempty
    shift first rest t             = return (first ++ head t) ++
                                     map (rest ++) (tail t)

-- | The elements of a tree in pre-order.
flatten :: (ApplicSMonoid f a, Foldable t) => Rose t a -> f a
flatten t = squish t mempty
  where squish (Rose x ts) xs = return x ++ foldr squish xs ts

-- | Lists of nodes at each level of the tree.
levels :: (Monoid (Bush t b), Foldable t, Applicative t) => Rose t b -> [t b]
levels t =
    map (map rootLabel) $
        takeWhile (not . null) $
        iterate (foldMap subBush) (return t)

-- | Build a tree from a seed value
unfoldRose :: Functor t => (b -> (a, t b)) -> b -> Rose t a
unfoldRose f b = let (a, bs) = f b in Rose a (unfoldBush f bs)

-- | Build a forest from a list of seed values
unfoldBush :: Functor t => (b -> (a, t b)) -> t b -> Bush t a
unfoldBush f = map (unfoldRose f)

-- | Monadic tree builder, in depth-first order
unfoldRoseM :: (ABMonad m, Traversable t) =>
               (b -> m (a, t b)) -> b -> m (Rose t a)
unfoldRoseM f b = do
    (a, bs) <- f b
    ts <- unfoldBushM f bs
    return (Rose a ts)

-- | Monadic forest builder, in depth-first order
unfoldBushM :: (ABMonad m, Traversable t) =>
               (a -> m (b, t a)) -> t a -> m (Bush t b)
unfoldBushM f = traverse (unfoldRoseM f)

-- | Monadic tree builder, in breadth-first order,
-- using an algorithm adapted from
-- /Breadth-First Numbering: Lessons from a Small Exercise in Algorithm Design/,
-- by Chris Okasaki, /ICFP'00/.
--unfoldRoseM_BF :: Monad m => (b -> m (a, [b])) -> b -> m (Rose a)
--unfoldRoseM_BF :: ABMonad m => (b -> m (a, [] b)) -> b -> m (Rose [] a)
unfoldRoseM_BF f b = getElement <$> unfoldBushQ f (return b)
  where
    getElement xs = case viewl xs of
        x :< _ -> x
        EmptyL -> error "unfoldRoseM_BF"

-- | Monadic forest builder, in breadth-first order,
-- using an algorithm adapted from
-- /Breadth-First Numbering: Lessons from a Small Exercise in Algorithm Design/,
-- by Chris Okasaki, /ICFP'00/.
--unfoldBushM_BF :: Monad m => (b -> m (a, [b])) -> [b] -> m (Bush a)
unfoldBushM_BF f = map toList . unfoldBushQ f . fromList

-- takes a sequence (queue) of seeds
-- produces a sequence (reversed queue) of trees of the same length
--unfoldBushQ :: Monad m => (b -> m (a, [b])) -> Seq b -> m (Seq (Rose a))
unfoldBushQ f aQ = case viewl aQ of
    EmptyL -> return mempty
    a :< aQ' -> do
        (b, as) <- f a
        tQ <- unfoldBushQ f (foldl (|>) aQ' as)
        let (tQ', ts) = splitOnto [] as tQ
        return (Rose b ts <| tQ')
  where
    --splitOnto :: [a'] -> [b'] -> Seq a' -> (Seq a', [a'])
    splitOnto as (null -> True) q = (q, as)
    splitOnto as (tail -> bs) q = case viewr q of
        q' :> a -> splitOnto (return a ++ as) bs q'
        EmptyR -> error "unfoldBushQ"