0 | module Data.Cotree
  1 |
  2 | import Data.Colist
  3 | import Data.Maybe
  4 | import Data.Tree
  5 |
  6 | %default total
  7 |
  8 | --------------------------------------------------------------------------------
  9 | --          Cotrees: Potentially infinte trees
 10 | --------------------------------------------------------------------------------
 11 |
 12 | ||| A potentially finite stream of trees
 13 | public export
 14 | Coforest : Type -> Type
 15 |
 16 | ||| A potentially infinite rose tree
 17 | public export
 18 | record Cotree (a : Type) where
 19 |   constructor MkCotree
 20 |   value  : a
 21 |   forest : Inf (Coforest a)
 22 |
 23 | Coforest = Colist . Cotree
 24 |
 25 | public export
 26 | singleton : a -> Cotree a
 27 | singleton a = MkCotree a Nil
 28 |
 29 | public export
 30 | unfold : (f : b -> (a,Colist b)) -> b -> Cotree a
 31 | unfold f vb =
 32 |   let (va,bs) := f vb
 33 |    in MkCotree va $ unfoldF bs
 34 |
 35 |   where
 36 |     unfoldF : Colist b -> Coforest a
 37 |     unfoldF []       = []
 38 |     unfoldF (h :: t) = unfold f h :: unfoldF t
 39 |
 40 | public export
 41 | iterate : (f : a -> Colist a) -> a -> Cotree a
 42 | iterate f a = unfold (\v => (v, f v)) a
 43 |
 44 | public export
 45 | expand : (a -> Colist a) -> Cotree a -> Cotree a
 46 | expand f (MkCotree v vs) =
 47 |   let MkCotree v2 vs2 := iterate f v
 48 |    in MkCotree v2 (run vs vs2)
 49 |
 50 |   where
 51 |     run : Coforest a -> Coforest a -> Coforest a
 52 |     run []        ys = ys
 53 |     run (x :: xs) ys = expand f x :: run xs ys
 54 |
 55 | --------------------------------------------------------------------------------
 56 | --          To and from Tree
 57 | --------------------------------------------------------------------------------
 58 |
 59 | public export
 60 | fromTree : Tree a -> Cotree a
 61 | fromTree (MkTree v fo) = MkCotree v (fromForest fo)
 62 |
 63 |   where
 64 |     fromForest : Forest a -> Coforest a
 65 |     fromForest []        = []
 66 |     fromForest (x :: xs) = fromTree x :: fromForest xs
 67 |
 68 | ||| Converts a Cotree to a tree of the given maximum depth and width.
 69 | ||| The maximum numbers of elements in the tree will be
 70 | ||| maxWidth ^ maxDepth.
 71 | public export
 72 | toTree : (maxDepth : Nat) -> (maxWidth : Nat) -> Cotree a -> Tree a
 73 | toTree 0     _  (MkCotree v fo) = MkTree v []
 74 | toTree (S k) mw (MkCotree v fo) = MkTree v (toForest mw fo)
 75 |
 76 |   where
 77 |     toForest : Nat -> Coforest a -> Forest a
 78 |     toForest 0     _         = []
 79 |     toForest (S n) []        = []
 80 |     toForest (S n) (t :: ts) = toTree k mw t :: toForest n ts
 81 |
 82 | --------------------------------------------------------------------------------
 83 | --          Functor and Applicative
 84 | --------------------------------------------------------------------------------
 85 |
 86 | public export
 87 | mapCotree : (a -> b) -> Cotree a -> Cotree b
 88 | mapCotree f (MkCotree v vs) = MkCotree (f v) (mapForest vs)
 89 |
 90 |   where
 91 |     mapForest : Coforest a -> Coforest b
 92 |     mapForest []       = []
 93 |     mapForest (h :: t) = mapCotree f h :: mapForest t
 94 |
 95 | public export
 96 | interleave : Cotree (a -> b) -> Cotree a -> Cotree b
 97 | interleave tf@(MkCotree vf fs) ta@(MkCotree va as) =
 98 |   MkCotree (vf va) (interleaveFs fs)
 99 |
100 |   where
101 |     interleaveAs : Coforest a -> Coforest b
102 |     interleaveAs []       = []
103 |     interleaveAs (h :: t) = interleave tf h :: interleaveAs t
104 |
105 |     interleaveFs : Coforest (a -> b) -> Coforest b
106 |     interleaveFs []       = interleaveAs as
107 |     interleaveFs (h :: t) = interleave h ta :: interleaveFs t
108 |
109 | public export
110 | bind : Cotree a -> (a -> Cotree b) -> Cotree b
111 | bind (MkCotree v vs) f =
112 |   let MkCotree w ws := f v
113 |    in MkCotree w (run vs ws)
114 |
115 |   where
116 |     run : Coforest a -> Coforest b -> Coforest b
117 |     run []        ys = ys
118 |     run (x :: xs) ys = bind x f :: run xs ys
119 |
120 | public export
121 | bindMaybe : Cotree (Maybe a) -> (a -> Cotree (Maybe b)) -> Cotree (Maybe b)
122 | bindMaybe (MkCotree mv tas) f =
123 |   case map f mv of
124 |     Nothing => MkCotree Nothing (run tas Nil)
125 |     Just (MkCotree mb tbs) => MkCotree mb (run tas tbs)
126 |
127 |   where
128 |     run : Coforest (Maybe a) -> Coforest (Maybe b) -> Coforest (Maybe b)
129 |     run [] ys        = ys
130 |     run (x :: xs) ys = bindMaybe x f :: run xs ys
131 |
132 | --------------------------------------------------------------------------------
133 | --          Shrinking
134 | --------------------------------------------------------------------------------
135 |
136 | public export
137 | shrink : (maxSteps : Nat) -> Cotree (Maybe a) -> List a
138 | shrink maxSteps x = run maxSteps [x]
139 |
140 |   where
141 |     run : Nat -> Coforest (Maybe a) -> List a
142 |     run _ Nil          = Nil
143 |     run 0 _            = Nil
144 |     run (S k) (h :: t) =
145 |       case h.value of
146 |         Just a  => a :: run k h.forest
147 |         Nothing => run k t
148 |
149 | public export
150 | mapShrink : (maxSteps : Nat) -> (a -> Maybe b) -> Cotree a -> List b
151 | mapShrink ms f = shrink ms . mapCotree f
152 |
153 | public export
154 | shrinkIf : (maxSteps : Nat) -> (a -> Bool) -> Cotree a -> List a
155 | shrinkIf ms p = mapShrink ms (\a => if p a then Just a else Nothing)
156 |
157 | ||| Prunes a cotree up to the given depth and width.
158 | public export
159 | pruneTo : (width : Nat) -> (depth : Nat) -> Cotree a -> Cotree a
160 | pruneTo _ 0     (MkCotree v _ ) = MkCotree v Nil
161 | pruneTo w (S d) (MkCotree v vs) = MkCotree v $ (map (pruneTo w d) $ keep w vs)
162 |
163 |   where
164 |     keep : Nat -> Colist t -> Colist t
165 |     keep _ [] = []
166 |     keep 0 _  = []
167 |     keep (S k) (x :: xs) = x :: keep k xs
168 |
169 | ||| Removes all children from a cotree
170 | public export
171 | prune : Cotree a -> Cotree a
172 | prune = pruneTo 0 0
173 |
174 | --------------------------------------------------------------------------------
175 | --          Interfaces
176 | --------------------------------------------------------------------------------
177 |
178 | public export
179 | Functor Cotree where
180 |   map = mapCotree
181 |
182 | public export
183 | Applicative Cotree where
184 |   pure  = singleton
185 |   (<*>) = interleave
186 |
187 | --------------------------------------------------------------------------------
188 | --          Filtering
189 | --------------------------------------------------------------------------------
190 |
191 | public export
192 | takeUntil : (a -> Bool) -> Cotree a -> Cotree a
193 | takeUntil f (MkCotree v vs) =
194 |   if f v
195 |      then MkCotree v []
196 |      else MkCotree v (takeUntilF vs)
197 |
198 |   where
199 |     takeUntilF : Coforest a -> Coforest a
200 |     takeUntilF []        = vs
201 |     takeUntilF (MkCotree x xs :: ts) =
202 |       if f x
203 |          then [MkCotree x []]
204 |          else MkCotree x (takeUntilF xs) :: takeUntilF ts
205 |
206 | public export
207 | takeBeforeNothing : Cotree (Maybe a) -> Maybe (Cotree a)
208 | takeBeforeNothing (MkCotree Nothing _) = Nothing
209 | takeBeforeNothing (MkCotree (Just v) vs) = Just (MkCotree v (run vs))
210 |
211 |   where
212 |     run : Coforest (Maybe a) -> Coforest a
213 |     run []                             = []
214 |     run ((MkCotree Nothing _) :: _)    = []
215 |     run ((MkCotree (Just v) vs) :: ts) = MkCotree v (run vs) :: run ts
216 |
217 | public export
218 | takeBefore : (a -> Bool) -> Cotree a -> Maybe (Cotree a)
219 | takeBefore f = takeBeforeNothing . map (\a => toMaybe (f a) a)
220 |
221 | public export %inline
222 | takeWhile : (a -> Bool) -> Cotree a -> Maybe (Cotree a)
223 | takeWhile f = takeBefore (not . f)
224 |
225 | public export %inline
226 | mapMaybe : (a -> Maybe b) -> Cotree a -> Maybe (Cotree b)
227 | mapMaybe f = takeBeforeNothing . map f
228 |