Data/TreeSeq/Strict.hs

   1 module Data.TreeSeq.Strict where
   2
   3 import Control.Applicative (Applicative(..))
   4 import Control.Monad (Monad(..))
   5 import Data.Eq (Eq(..))
   6 import Data.Foldable (Foldable(..))
   7 import Data.Function (($), (.))
   8 import Data.Functor (Functor(..), (<$>))
   9 import Data.Monoid (Monoid(..))
  10 import Data.Ord (Ord(..))
  11 import Data.Semigroup (Semigroup(..))
  12 import Data.Sequence (Seq, ViewL(..))
  13 import Data.Traversable (Traversable(..))
  14 import Text.Show (Show(..))
  15 import qualified Data.List as List
  16 import qualified Data.Sequence as Seq
  17
  18 -- * Type 'Tree'
  19 data Tree a
  20  =   Tree { unTree   :: !a
  21           , subTrees :: !(Trees a)
  22           }
  23  deriving (Eq, Ord, Show)
  24 instance Functor Tree where
  25         fmap f (Tree a ts) = Tree (f a) (fmap (fmap f) ts)
  26 instance Applicative Tree where
  27         pure a = Tree a mempty
  28         Tree f tfs <*> ta@(Tree a tas) =
  29                 Tree (f a) (fmap (f <$>) tas <> fmap (<*> ta) tfs)
  30 instance Monad Tree where
  31         return = pure
  32         Tree a ts >>= f =
  33                 Tree a' (ts' <> fmap (>>= f) ts)
  34                 where Tree a' ts' = f a
  35 instance Foldable Tree where
  36         foldMap f (Tree a ts) = f a `mappend` foldMap (foldMap f) ts
  37 instance Traversable Tree where
  38         traverse f (Tree a ts) = Tree <$> f a <*> traverse (traverse f) ts
  39         sequenceA  (Tree a ts) = Tree <$>   a <*> traverse sequenceA ts
  40
  41 tree0 :: a -> Tree a
  42 tree0 a = Tree a mempty
  43
  44 isTree0 :: Tree a -> Bool
  45 isTree0 (Tree _ ts) = null ts
  46
  47 isTreeN :: Tree a -> Bool
  48 isTreeN (Tree _ ts) = not (null ts)
  49
  50 -- * Type 'Trees'
  51 type Trees a = Seq (Tree a)
  52
  53 prettyTree :: Show a => Tree a -> String
  54 prettyTree = List.unlines . pretty
  55
  56 prettyTrees :: Show a => Trees a -> String
  57 prettyTrees = foldr (\t acc -> prettyTree t <> "\n" <> acc) ""
  58
  59 pretty :: Show a => Tree a -> [String]
  60 pretty (Tree a ts0) = show a : prettySubTrees ts0
  61         where
  62         shift first other = List.zipWith (<>) $ first : List.repeat other
  63         prettySubTrees s =
  64                 case Seq.viewl s of
  65                  Seq.EmptyL -> []
  66                  t:<ts | Seq.null ts -> "|" : shift "`- " "   " (pretty t)
  67                        | otherwise   -> "|" : shift "+- " "|  " (pretty t) <> prettySubTrees ts