Data/TreeSeq/Strict.hs

   1 {-# LANGUAGE BangPatterns #-}
   2 {-# LANGUAGE FlexibleInstances #-}
   3 {-# LANGUAGE OverloadedStrings #-}
   4 module Data.TreeSeq.Strict where
   5
   6 import Control.Applicative (Applicative(..))
   7 import Control.Monad (Monad(..))
   8 import Data.Bool
   9 import Data.Eq (Eq(..))
  10 import Data.Foldable (Foldable(..))
  11 import Data.Foldable (foldr)
  12 import Data.Function (($), (.))
  13 import Data.Functor (Functor(..), (<$>))
  14 import Data.Monoid (Monoid(..))
  15 import Data.Ord (Ord(..))
  16 import Data.Semigroup (Semigroup(..))
  17 import Data.Sequence (Seq, ViewL(..))
  18 import Data.Traversable (Traversable(..))
  19 import Text.Show (Show(..))
  20 import qualified Data.List as List
  21 import qualified Data.Sequence as Seq
  22 import qualified Data.Text.Lazy as TL
  23
  24 -- * Type 'Tree'
  25 data Tree a
  26  =   Tree { unTree   :: !a
  27           , subTrees :: !(Trees a)
  28           }
  29  deriving (Eq, Ord, Show)
  30 instance Functor Tree where
  31         fmap f (Tree a ts) = Tree (f a) (fmap (fmap f) ts)
  32 instance Applicative Tree where
  33         pure a = Tree a mempty
  34         Tree f tfs <*> ta@(Tree a tas) =
  35                 Tree (f a) (fmap (f <$>) tas <> fmap (<*> ta) tfs)
  36 instance Monad Tree where
  37         return = pure
  38         Tree a ts >>= f =
  39                 Tree a' (ts' <> fmap (>>= f) ts)
  40                 where Tree a' ts' = f a
  41 instance Foldable Tree where
  42         foldMap f (Tree a ts) = f a `mappend` foldMap (foldMap f) ts
  43 instance Traversable Tree where
  44         traverse f (Tree a ts) = Tree <$> f a <*> traverse (traverse f) ts
  45         sequenceA  (Tree a ts) = Tree <$>   a <*> traverse sequenceA ts
  46
  47 tree0 :: a -> Tree a
  48 tree0 a = Tree a mempty
  49
  50 isTree0 :: Tree a -> Bool
  51 isTree0 (Tree _ ts) = null ts
  52
  53 isTreeN :: Tree a -> Bool
  54 isTreeN (Tree _ ts) = not (null ts)
  55
  56 {-
  57 mapWithNode :: (Maybe k -> a -> b) -> Tree k a -> Tree k b
  58 mapWithNode = go Nothing
  59         where
  60         go _k f (Tree k ts) = Tree k (go (Just k) f <$> ts)
  61         go k  f (Tree0 a)    = Tree0 (f k a)
  62
  63 mapAlsoNode :: (k -> l) -> (Maybe k -> a -> b) -> Tree k a -> Tree l b
  64 mapAlsoNode fk fv = go Nothing
  65         where
  66         go _k (Tree k ts) = Tree (fk k) $ go (Just k) <$> ts
  67         go k  (Tree0 a)    = Tree0 (fv k a)
  68
  69 traverseWithNode :: Applicative f => (Maybe k -> a -> f b) -> Tree k a -> f (Tree k b)
  70 traverseWithNode = go Nothing
  71         where
  72         go _p f (Tree k ts) = Tree k <$> traverse (go (Just k) f) ts
  73         go p  f (Tree0 a)    = Tree0 <$> f p a
  74
  75 foldlWithTree :: (b -> Tree k a -> b) -> b -> Tree k a -> b
  76 foldlWithTree f b t =
  77         case t of
  78          Tree _k ts -> foldl' (foldlWithTree f) (f b t) ts
  79          Tree0{} -> f b t
  80
  81 bindTree :: Tree k a -> (Tree k a -> Tree l b) -> Tree l b
  82 bindTree t f =
  83         case t of
  84          Tree0{} -> f t
  85          Tree _k ks ->
  86                 case f t of
  87                  u@Tree0{}  -> u
  88                  Tree l ls -> Tree l $ ls <> ((`bindTree` f) <$> ks)
  89
  90 bindTrees :: Tree k a -> (Tree k a -> Trees l b) -> Trees l b
  91 bindTrees t f =
  92         case t of
  93          Tree0{} -> f t
  94          Tree _k ks ->
  95                 f t >>= \fs ->
  96                         case fs of
  97                          Tree0 b  -> Seq.singleton $ Tree0 b
  98                          Tree l ls -> pure $ Tree l $ ls <> (ks >>= (`bindTrees` f))
  99
 100 joinTrees :: Trees k (Trees k a) -> Trees k a
 101 joinTrees ts =
 102         ts >>= \case
 103          Tree0 s -> s
 104          Tree k ks -> Seq.singleton $ Tree k $ joinTrees ks
 105 -}
 106
 107 -- * Type 'Trees'
 108 type Trees a = Seq (Tree a)
 109
 110 -- * Type 'Pretty'
 111 newtype Pretty a = Pretty a
 112 instance Show a => Show (Pretty (Trees a)) where
 113         show (Pretty t) = TL.unpack $ prettyTrees t
 114 instance Show a => Show (Pretty (Tree a)) where
 115         show (Pretty t) = TL.unpack $ prettyTree t
 116
 117 prettyTree :: Show a => Tree a -> TL.Text
 118 prettyTree = TL.unlines . pretty
 119
 120 prettyTrees :: Show a => Trees a -> TL.Text
 121 prettyTrees = foldr (\t acc -> prettyTree t <> "\n" <> acc) ""
 122
 123 pretty :: Show a => Tree a -> [TL.Text]
 124 pretty (Tree a ts0) = TL.pack (show a) : prettySubTrees ts0
 125         where
 126         prettySubTrees s =
 127                 case Seq.viewl s of
 128                  Seq.EmptyL -> []
 129                  t:<ts | Seq.null ts -> "|" : shift "`- " "   " (pretty t)
 130                        | otherwise   -> "|" : shift "+- " "|  " (pretty t) <> prettySubTrees ts
 131         shift first other = List.zipWith (<>) (first : List.repeat other)