1 {-# LANGUAGE BangPatterns #-}
2 {-# LANGUAGE DeriveFunctor #-}
3 {-# LANGUAGE OverloadedStrings #-}
4 module Data.TreeSeq.Strict where
6 import Control.Applicative (Applicative(..))
7 import Control.Monad (Monad(..), ap)
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
27 = TreeN !k !(Trees k a)
29 deriving (Eq, Ord, Show, Functor)
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
42 instance Monad (Tree k) where
46 TreeN k $ (>>= f) <$> ts
48 isTree0 :: Tree k a -> Bool
49 isTree0 Tree0{} = True
52 isTreeN :: Tree k a -> Bool
53 isTreeN TreeN{} = True
56 unTree :: Tree a a -> a
57 unTree (TreeN k _) = k
60 mapWithKey :: (Maybe k -> a -> b) -> Tree k a -> Tree k b
61 mapWithKey = go Nothing
63 go _k f (TreeN k ts) = TreeN k (go (Just k) f <$> ts)
64 go k f (Tree0 a) = Tree0 (f k a)
66 mapAlsoKey :: (k -> l) -> (Maybe k -> a -> b) -> Tree k a -> Tree l b
67 mapAlsoKey fk fv = go Nothing
69 go _k (TreeN k ts) = TreeN (fk k) $ go (Just k) <$> ts
70 go k (Tree0 a) = Tree0 (fv k a)
72 traverseWithNode :: Applicative f => (Maybe k -> a -> f b) -> Tree k a -> f (Tree k b)
73 traverseWithNode = go Nothing
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
78 foldlWithTree :: (b -> Tree k a -> b) -> b -> Tree k a -> b
81 TreeN _k ts -> foldl' (foldlWithTree f) (f b t) ts
84 bindTree :: Tree k a -> (Tree k a -> Tree l b) -> Tree l b
91 TreeN l ls -> TreeN l $ ls <> ((`bindTree` f) <$> ks)
93 bindTrees :: Tree k a -> (Tree k a -> Trees l b) -> Trees l b
100 Tree0 b -> Seq.singleton $ Tree0 b
101 TreeN l ls -> pure $ TreeN l $ ls <> (ks >>= (`bindTrees` f))
103 joinTrees :: Trees k (Trees k a) -> Trees k a
107 TreeN k ks -> Seq.singleton $ TreeN k $ joinTrees ks
110 type Trees k a = Seq (Tree k a)
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
117 prettyTree :: (Show k, Show a) => Tree k a -> Text
118 prettyTree = Text.unlines . pretty
120 prettyTrees :: (Show k, Show a) => Trees k a -> Text
121 prettyTrees = foldr (\t acc -> prettyTree t <> "\n" <> acc) ""
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
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)