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 Data.Bool
   8 import Data.Eq (Eq(..))
   9 import Data.Foldable (Foldable(..))
  10 import Data.Foldable (foldr)
  11 import Data.Function (($), (.))
  12 import Data.Functor (Functor, (<$>))
  13 import Data.Maybe (Maybe(..))
  14 import Data.Ord (Ord(..))
  15 import Data.Semigroup (Semigroup(..))
  16 import Data.Sequence (Seq, ViewL(..))
  17 import Data.Text (Text)
  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 as Text
  23
  24 -- * Type 'Tree'
  25 data Tree k a
  26  =   TreeN !k !(Trees k a)
  27  |   Tree0 !a
  28  deriving (Eq, Ord, Show, Functor)
  29
  30 instance Traversable (Tree k) where
  31         traverse f (Tree0 a)    = Tree0 <$> f a
  32         traverse f (TreeN k ts) = TreeN k <$> traverse (traverse f) ts
  33         sequenceA (Tree0 a)     = Tree0 <$> a
  34         sequenceA (TreeN k ts)  = TreeN k <$> traverse sequenceA ts
  35 instance Foldable (Tree k) where
  36         foldMap f (TreeN _k ts) = foldMap (foldMap f) ts
  37         foldMap f (Tree0 k)     = f k
  38
  39 isTree0 :: Tree k a -> Bool
  40 isTree0 Tree0{} = True
  41 isTree0 _       = False
  42
  43 isTreeN :: Tree k a -> Bool
  44 isTreeN TreeN{} = True
  45 isTreeN _       = False
  46
  47 unTree :: Tree a a -> a
  48 unTree (TreeN k _) = k
  49 unTree (Tree0 a)   = a
  50
  51 mapWithKey :: (Maybe k -> a -> b) -> Tree k a -> Tree k b
  52 mapWithKey = go Nothing
  53         where
  54         go _k f (TreeN k ts) = TreeN k (go (Just k) f <$> ts)
  55         go k  f (Tree0 a)    = Tree0 (f k a)
  56
  57 mapAlsoKey :: (k -> l) -> (Maybe k -> a -> b) -> Tree k a -> Tree l b
  58 mapAlsoKey fk fv = go Nothing
  59         where
  60         go _k (TreeN k ts) = TreeN (fk k) $ go (Just k) <$> ts
  61         go k  (Tree0 a)    = Tree0 (fv k a)
  62
  63 traverseWithKey :: Applicative f => (Maybe k -> a -> f b) -> Tree k a -> f (Tree k b)
  64 traverseWithKey = go Nothing
  65         where
  66         go _p f (TreeN k ts) = TreeN k <$> traverse (go (Just k) f) ts
  67         go p  f (Tree0 a)    = Tree0 <$> f p a
  68
  69 foldlWithTree :: (b -> Tree k a -> b) -> b -> Tree k a -> b
  70 foldlWithTree f b t =
  71         case t of
  72          TreeN _k ts -> foldl' (foldlWithTree f) (f b t) ts
  73          Tree0{} -> f b t
  74
  75 -- * Type 'Trees'
  76 type Trees k a = Seq (Tree k a)
  77
  78 -- * Type 'Pretty'
  79 newtype Pretty k a = Pretty (Trees k a)
  80 instance (Show k, Show a) => Show (Pretty k a) where
  81         show (Pretty t) = Text.unpack $ prettyTrees t
  82
  83 prettyTree :: (Show k, Show a) => Tree k a -> Text
  84 prettyTree = Text.unlines . pretty
  85
  86 prettyTrees :: (Show k, Show a) => Trees k a -> Text
  87 prettyTrees = foldr (\t acc -> prettyTree t <> "\n" <> acc) ""
  88
  89 pretty :: (Show k, Show a) => Tree k a -> [Text]
  90 pretty (Tree0 a)     = [Text.pack (show a)]
  91 pretty (TreeN k ts0) = Text.pack (show k) : prettySubTrees ts0
  92         where
  93         prettySubTrees s =
  94                 case Seq.viewl s of
  95                  Seq.EmptyL -> []
  96                  t:<ts | Seq.null ts -> "|" : shift "`- " "   " (pretty t)
  97                        | otherwise   -> "|" : shift "+- " "|  " (pretty t) <> prettySubTrees ts
  98         shift first other = List.zipWith (<>) (first : List.repeat other)