Data/TreeSeq/Strict.hs

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