0 | module Data.Tree
  1 |
  2 | import Data.List
  3 | import Data.List1
  4 | import Data.String
  5 |
  6 | import Derive.Prelude
  7 |
  8 | %language ElabReflection
  9 | %default total
 10 |
 11 | --------------------------------------------------------------------------------
 12 | --          Finite Trees
 13 | --------------------------------------------------------------------------------
 14 |
 15 | ||| A finite rose tree
 16 | public export
 17 | record Tree (a : Type) where
 18 |   constructor MkTree
 19 |   value : a
 20 |   forest : List (Tree a)
 21 |
 22 | ||| A finite forest of trees
 23 | public export
 24 | Forest : Type -> Type
 25 | Forest = List . Tree
 26 |
 27 | %runElab derive "Tree" [Show,Eq]
 28 |
 29 | --------------------------------------------------------------------------------
 30 | --          Creating Trees
 31 | --------------------------------------------------------------------------------
 32 |
 33 | public export
 34 | singleton : a -> Tree a
 35 | singleton a = MkTree a []
 36 |
 37 | public export
 38 | replicate : (width : Nat) -> (depth : Nat) -> a -> Tree a
 39 | replicate _         0 x = MkTree x []
 40 | replicate width (S k) x = MkTree x $ replicate width (replicate width k x)
 41 |
 42 | ||| Unfold a tree up to the given depth.
 43 | public export
 44 | unfold : (depth : Nat) -> (f : s -> (a,List s)) -> s -> Tree a
 45 | unfold 0     f s = MkTree (fst $ f s) []
 46 | unfold (S k) f s =
 47 |   let (a,ss) := f s
 48 |    in MkTree a (map (unfold k f) ss)
 49 |
 50 | --------------------------------------------------------------------------------
 51 | --          Flattening Trees
 52 | --------------------------------------------------------------------------------
 53 |
 54 | zipWithKeep : (a -> a -> a) -> List a -> List a -> List a
 55 | zipWithKeep f [] ys = ys
 56 | zipWithKeep f xs [] = xs
 57 | zipWithKeep f (x :: xs) (y :: ys) = f x y :: zipWithKeep f xs ys
 58 |
 59 | public export
 60 | flatten : Tree a -> List a
 61 | flatten (MkTree v vs) = v :: flattenF vs
 62 |
 63 |   where
 64 |     flattenF : Forest a -> List a
 65 |     flattenF []        = Nil
 66 |     flattenF (x :: xs) = flatten x ++ flattenF xs
 67 |
 68 | public export
 69 | layers : Tree a -> List (List a)
 70 | layers (MkTree v vs) = [v] :: layersF vs
 71 |
 72 |   where
 73 |     layersF : Forest a -> List (List a)
 74 |     layersF []        = Nil
 75 |     layersF (x :: xs) = zipWithKeep (++) (layers x) (layersF xs)
 76 |
 77 | --------------------------------------------------------------------------------
 78 | --          Accessing Elements
 79 | --------------------------------------------------------------------------------
 80 |
 81 | public export
 82 | index : List Nat -> Tree a -> Maybe a
 83 | index []        x = Just x.value
 84 | index (y :: ys) x = ix y x.forest >>= index ys
 85 |
 86 |   where
 87 |     ix : Nat -> List b -> Maybe b
 88 |     ix _ []            = Nothing
 89 |     ix 0     (z :: _)  = Just z
 90 |     ix (S k) (_ :: zs) = ix k zs
 91 |
 92 | --------------------------------------------------------------------------------
 93 | --          Functor and Monad Implementations
 94 | --------------------------------------------------------------------------------
 95 |
 96 | -- All implementations are boilerplaty to satisfy the totality checker.
 97 | foldlTree : (a -> e -> a) -> a -> Tree e -> a
 98 | foldlTree f acc (MkTree v vs) = foldlF (f acc v) vs
 99 |
100 |   where
101 |     foldlF : a -> Forest e -> a
102 |     foldlF y []        = y
103 |     foldlF y (x :: xs) = foldlF (foldlTree f y x) xs
104 |
105 | foldrTree : (e -> a -> a) -> a -> Tree e -> a
106 | foldrTree f acc (MkTree v vs) = f v (foldrF acc vs)
107 |
108 |   where
109 |     foldrF : a -> Forest e -> a
110 |     foldrF y []        = y
111 |     foldrF y (x :: xs) = foldrTree f (foldrF y xs) x
112 |
113 | traverseTree : Applicative f => (a -> f b) -> Tree a -> f (Tree b)
114 | traverseTree fun (MkTree v vs) = [| MkTree (fun v) (traverseF vs) |]
115 |
116 |   where
117 |     traverseF : Forest a -> f (Forest b)
118 |     traverseF []        = pure []
119 |     traverseF (x :: xs) = [| traverseTree fun x :: traverseF xs |]
120 |
121 | mapTree : (a -> b) -> Tree a -> Tree b
122 | mapTree f (MkTree v vs) = MkTree (f v) (mapF vs)
123 |
124 |   where
125 |     mapF : Forest a -> Forest b
126 |     mapF []       = []
127 |     mapF (h :: t) = mapTree f h :: mapF t
128 |
129 | bindTree : Tree a -> (a -> Tree b) -> Tree b
130 | bindTree (MkTree va tas) f =
131 |   let MkTree vb tbs := f va
132 |    in MkTree vb (tbs ++ bindF tas)
133 |
134 |   where
135 |     bindF : Forest a -> Forest b
136 |     bindF []        = []
137 |     bindF (x :: xs) = bindTree x f :: bindF xs
138 |
139 | apTree : Tree (a -> b) -> Tree a -> Tree b
140 | apTree tf ta = bindTree tf $ \f => mapTree (apply f) ta
141 |
142 | joinTree : Tree (Tree a) -> Tree a
143 | joinTree (MkTree (MkTree va tas) ftas) =
144 |   MkTree va $ tas ++ joinF ftas
145 |
146 |   where
147 |     joinF : Forest (Tree a) -> Forest a
148 |     joinF []        = []
149 |     joinF (x :: xs) = joinTree x :: joinF xs
150 |
151 | --------------------------------------------------------------------------------
152 | --          Visualizing Trees
153 | --------------------------------------------------------------------------------
154 |
155 | export
156 | drawTree : Tree String -> String
157 | drawTree  = unlines . draw
158 |
159 |   where
160 |     drawForest : Forest String -> String
161 |     drawForest  = unlines . map drawTree
162 |
163 |     draw : Tree String -> List String
164 |     draw (MkTree x ts0) = lines x ++ subTrees ts0
165 |
166 |       where
167 |         shift : String -> String -> List String -> List String
168 |         shift first other tails =
169 |           zipWith (++) (first :: replicate (length tails) other) tails
170 |
171 |         subTrees : Forest String -> List String
172 |         subTrees []      = []
173 |         subTrees [t]     = "│" :: shift "└╼" "  " (draw t)
174 |         subTrees (t::ts) = "│" :: shift "├╼" "│ " (draw t) ++ subTrees ts
175 |
176 | --------------------------------------------------------------------------------
177 | --          Interfaces
178 | --------------------------------------------------------------------------------
179 |
180 | public export %inline
181 | Foldable Tree where
182 |   foldl  = foldlTree
183 |   foldr  = foldrTree
184 |   null _ = False
185 |
186 | public export %inline
187 | Functor Tree where
188 |   map = mapTree
189 |
190 | public export %inline
191 | Applicative Tree where
192 |   pure a = MkTree a Nil
193 |   (<*>)  = apTree
194 |
195 | public export %inline
196 | Monad Tree where
197 |   (>>=) = bindTree
198 |   join  = joinTree
199 |
200 | public export %inline
201 | Traversable Tree where
202 |   traverse = traverseTree
203 |