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 -- * Type 'Trees'
  70 type Trees k a = Seq (Tree k a)
  71
  72 -- * Type 'Pretty'
  73 newtype Pretty k a = Pretty (Trees k a)
  74 instance (Show k, Show a) => Show (Pretty k a) where
  75         show (Pretty t) = Text.unpack $ prettyTrees t
  76
  77 prettyTree :: (Show k, Show a) => Tree k a -> Text
  78 prettyTree = Text.unlines . pretty
  79
  80 prettyTrees :: (Show k, Show a) => Trees k a -> Text
  81 prettyTrees = foldr (\t acc -> prettyTree t <> "\n" <> acc) ""
  82
  83 pretty :: (Show k, Show a) => Tree k a -> [Text]
  84 pretty (Tree0 a)     = [Text.pack (show a)]
  85 pretty (TreeN k ts0) = Text.pack (show k) : prettySubTrees ts0
  86         where
  87         prettySubTrees s =
  88                 case Seq.viewl s of
  89                  Seq.EmptyL -> []
  90                  t:<ts | Seq.null ts -> "|" : shift "`- " "   " (pretty t)
  91                        | otherwise   -> "|" : shift "+- " "|  " (pretty t) <> prettySubTrees ts
  92         shift first other = List.zipWith (<>) (first : List.repeat other)