001    package fj.data;
002    
003    import fj.F;
004    import fj.F2;
005    import fj.P;
006    import fj.P1;
007    import fj.P2;
008    import static fj.F2W.$$;
009    import static fj.Function.*;
010    import static fj.data.Stream.*;
011    import fj.pre.Monoid;
012    import fj.pre.Show;
013    
014    import java.util.Collection;
015    import java.util.Iterator;
016    
017    /**
018     * Provides a lazy, immutable, non-empty, multi-way tree (a rose tree).
019     *
020     * @version %build.number%<br>
021     *          <ul>
022     *          <li>$LastChangedRevision: 196 $</li>
023     *          <li>$LastChangedDate: 2009-06-28 16:02:26 +1000 (Sun, 28 Jun 2009) $</li>
024     *          <li>Author: runar</li>
025     *          </ul>
026     */
027    public final class Tree<A> implements Iterable<A> {
028      /**
029       * Returns an iterator for this tree. This method exists to permit the use in a <code>for</code>-each loop.
030       *
031       * @return A iterator for this tree.
032       */
033      public Iterator<A> iterator() {
034        return flatten().iterator();
035      }
036    
037      private final A root;
038      private final P1<Stream<Tree<A>>> subForest;
039    
040      private Tree(final A root, final P1<Stream<Tree<A>>> subForest) {
041        this.root = root;
042        this.subForest = subForest;
043      }
044    
045      /**
046       * Creates a nullary tree.
047       *
048       * @param root The root element of the tree.
049       * @return A nullary tree with the root element in it.
050       */
051      public static <A> Tree<A> leaf(final A root) {
052        return node(root, Stream.<Tree<A>>nil());
053      }
054    
055      /**
056       * Creates a new tree given a root and a (potentially infinite) subforest.
057       *
058       * @param root   The root element of the tree.
059       * @param forest A stream of the tree's subtrees.
060       * @return A newly sprouted tree.
061       */
062      public static <A> Tree<A> node(final A root, final P1<Stream<Tree<A>>> forest) {
063        return new Tree<A>(root, forest);
064      }
065    
066      /**
067       * Creates a new tree given a root and a (potentially infinite) subforest.
068       *
069       * @param root   The root element of the tree.
070       * @param forest A stream of the tree's subtrees.
071       * @return A newly sprouted tree.
072       */
073      public static <A> Tree<A> node(final A root, final Stream<Tree<A>> forest) {
074        return new Tree<A>(root, P.p(forest));
075      }
076    
077      /**
078       * Creates a new n-ary given a root and a subforest of length n.
079       *
080       * @param root   The root element of the tree.
081       * @param forest A list of the tree's subtrees.
082       * @return A newly sprouted tree.
083       */
084      public static <A> Tree<A> node(final A root, final List<Tree<A>> forest) {
085        return node(root, forest.toStream());
086      }
087    
088      /**
089       * First-class constructor of trees.
090       *
091       * @return A function that constructs an n-ary tree given a root and a subforest or length n.
092       */
093      public static <A> F<A, F<P1<Stream<Tree<A>>>, Tree<A>>> node() {
094        return curry(new F2<A, P1<Stream<Tree<A>>>, Tree<A>>() {
095          public Tree<A> f(final A a, final P1<Stream<Tree<A>>> p1) {
096            return node(a, p1);
097          }
098        });
099      }
100    
101      /**
102       * Returns the root element of the tree.
103       *
104       * @return The root element of the tree.
105       */
106      public A root() {
107        return root;
108      }
109    
110      /**
111       * Returns a stream of the tree's subtrees.
112       *
113       * @return A stream of the tree's subtrees.
114       */
115      public P1<Stream<Tree<A>>> subForest() {
116        return subForest;
117      }
118    
119      /**
120       * Provides a transformation from a tree to its root.
121       *
122       * @return A transformation from a tree to its root.
123       */
124      public static <A> F<Tree<A>, A> root_() {
125        return new F<Tree<A>, A>() {
126          public A f(final Tree<A> a) {
127            return a.root();
128          }
129        };
130      }
131    
132      /**
133       * Provides a transformation from a tree to its subforest.
134       *
135       * @return A transformation from a tree to its subforest.
136       */
137      public static <A> F<Tree<A>, P1<Stream<Tree<A>>>> subForest_() {
138        return new F<Tree<A>, P1<Stream<Tree<A>>>>() {
139          public P1<Stream<Tree<A>>> f(final Tree<A> a) {
140            return a.subForest();
141          }
142        };
143      }
144    
145      /**
146       * Puts the elements of the tree into a Stream, in pre-order.
147       *
148       * @return The elements of the tree in pre-order.
149       */
150      public Stream<A> flatten() {
151        final F2<Tree<A>, P1<Stream<A>>, Stream<A>> squish = new F2<Tree<A>, P1<Stream<A>>, Stream<A>>() {
152          public Stream<A> f(final Tree<A> t, final P1<Stream<A>> xs) {
153            return cons(t.root(), t.subForest().map(Stream.<Tree<A>, Stream<A>>foldRight().f(curry(this)).f(xs._1())));
154          }
155        };
156        return squish.f(this, P.p(Stream.<A>nil()));
157      }
158    
159      /**
160       * flatten :: Tree a -> [a]
161       * flatten t = squish t []
162       * where squish (Node x ts) xs = x:Prelude.foldr squish xs ts
163       * Puts the elements of the tree into a Stream, in pre-order.
164       *
165       * @return The elements of the tree in pre-order.
166       */
167      public static <A> F<Tree<A>, Stream<A>> flatten_() {
168        return new F<Tree<A>, Stream<A>>() {
169          public Stream<A> f(final Tree<A> t) {
170            return t.flatten();
171          }
172        };
173      }
174    
175      /**
176       * Provides a stream of the elements of the tree at each level, in level order.
177       *
178       * @return The elements of the tree at each level.
179       */
180      public Stream<Stream<A>> levels() {
181        final F<Stream<Tree<A>>, Stream<Tree<A>>> flatSubForests =
182            Stream.<Tree<A>, Tree<A>>bind_().f(compose(P1.<Stream<Tree<A>>>__1(), Tree.<A>subForest_()));
183        final F<Stream<Tree<A>>, Stream<A>> roots = Stream.<Tree<A>, A>map_().f(Tree.<A>root_());
184        return iterateWhile(flatSubForests, Stream.<Tree<A>>isNotEmpty_(), single(this)).map(roots);
185      }
186    
187      /**
188       * Maps the given function over this tree.
189       *
190       * @param f The function to map over this tree.
191       * @return The new Tree after the function has been applied to each element in this Tree.
192       */
193      public <B> Tree<B> fmap(final F<A, B> f) {
194        return node(f.f(root()), subForest().map(Stream.<Tree<A>, Tree<B>>map_().f(Tree.<A, B>fmap_().f(f))));
195      }
196    
197      /**
198       * Provides a transformation to lift any function so that it maps over Trees.
199       *
200       * @return A transformation to lift any function so that it maps over Trees.
201       */
202      public static <A, B> F<F<A, B>, F<Tree<A>, Tree<B>>> fmap_() {
203        return new F<F<A, B>, F<Tree<A>, Tree<B>>>() {
204          public F<Tree<A>, Tree<B>> f(final F<A, B> f) {
205            return new F<Tree<A>, Tree<B>>() {
206              public Tree<B> f(final Tree<A> a) {
207                return a.fmap(f);
208              }
209            };
210          }
211        };
212      }
213    
214      /**
215       * Folds this tree using the given monoid.
216       *
217       * @param f A transformation from this tree's elements, to the monoid.
218       * @param m The monoid to fold this tree with.
219       * @return The result of folding the tree with the given monoid.
220       */
221      public <B> B foldMap(final F<A, B> f, final Monoid<B> m) {
222        return m.sum(f.f(root()), m.sumRight(subForest()._1().map(foldMap_(f, m)).toList()));
223      }
224    
225      /**
226       * Projects an immutable collection of this tree.
227       *
228       * @return An immutable collection of this tree.
229       */
230      public Collection<A> toCollection() {
231        return flatten().toCollection();
232      }
233    
234      /**
235       * Provides a function that folds a tree with the given monoid.
236       *
237       * @param f A transformation from a tree's elements to the monoid.
238       * @param m A monoid to fold the tree with.
239       * @return A function that, given a tree, folds it with the given monoid.
240       */
241      public static <A, B> F<Tree<A>, B> foldMap_(final F<A, B> f, final Monoid<B> m) {
242        return new F<Tree<A>, B>() {
243          public B f(final Tree<A> t) {
244            return t.foldMap(f, m);
245          }
246        };
247      }
248    
249      /**
250       * Builds a tree from a seed value.
251       *
252       * @param f A function with which to build the tree.
253       * @return A function which, given a seed value, yields a tree.
254       */
255      public static <A, B> F<B, Tree<A>> unfoldTree(final F<B, P2<A, P1<Stream<B>>>> f) {
256        return new F<B, Tree<A>>() {
257          public Tree<A> f(final B b) {
258            final P2<A, P1<Stream<B>>> p = f.f(b);
259            return node(p._1(), p._2().map(Stream.<B, Tree<A>>map_().f(unfoldTree(f))));
260          }
261        };
262      }
263    
264      /**
265       * Applies the given function to all subtrees of this tree, returning a tree of the results (comonad pattern).
266       *
267       * @param f A function to bind across all the subtrees of this tree.
268       * @return A new tree, with the results of applying the given function to each subtree of this tree. The result
269       *         of applying the function to the entire tree is the root label, and the results of applying to the
270       *         root's children are labels of the root's subforest, etc.
271       */
272      public <B> Tree<B> cobind(final F<Tree<A>, B> f) {
273        return unfoldTree(new F<Tree<A>, P2<B, P1<Stream<Tree<A>>>>>() {
274          public P2<B, P1<Stream<Tree<A>>>> f(final Tree<A> t) {
275            return P.p(f.f(t), t.subForest());
276          }
277        }).f(this);
278      }
279    
280      /**
281       * Expands this tree into a tree of trees, with this tree as the root label, and subtrees as the labels of
282       * child nodes (comonad pattern).
283       *
284       * @return A tree of trees, with this tree as its root label, and subtrees of this tree as the labels of
285       *         its child nodes.
286       */
287      public Tree<Tree<A>> cojoin() {
288        final F<Tree<A>, Tree<A>> id = identity();
289        return cobind(id);
290      }
291    
292      private static <A> Stream<String> drawSubTrees(final Show<A> s, final Stream<Tree<A>> ts) {
293        return ts.isEmpty() ? Stream.<String>nil()
294                            : ts.tail()._1().isEmpty() ? shift("`- ", "   ", ts.head().drawTree(s)).cons("|")
295                                                       : shift("+- ", "|  ", ts.head().drawTree(s))
296                                                           .append(drawSubTrees(s, ts.tail()._1()));
297      }
298    
299      private static Stream<String> shift(final String f, final String o, final Stream<String> s) {
300        return Stream.repeat(o).cons(f).zipWith(s, Monoid.stringMonoid.sum());
301      }
302    
303      private Stream<String> drawTree(final Show<A> s) {
304        return drawSubTrees(s, subForest._1()).cons(s.showS(root));
305      }
306    
307      /**
308       * Draws a 2-dimensional representation of a tree.
309       *
310       * @param s A show instance for the elements of the tree.
311       * @return a String showing this tree in two dimensions.
312       */
313      public String draw(final Show<A> s) {
314        return Monoid.stringMonoid.join(drawTree(s), "\n");
315      }
316    
317      /**
318       * Provides a show instance that draws a 2-dimensional representation of a tree.
319       *
320       * @param s A show instance for the elements of the tree.
321       * @return a show instance that draws a 2-dimensional representation of a tree.
322       */
323      public static <A> Show<Tree<A>> show2D(final Show<A> s) {
324        return Show.showS(new F<Tree<A>, String>() {
325          public String f(final Tree<A> tree) {
326            return tree.draw(s);
327          }
328        });
329      }
330    
331      /**
332       * Zips this tree with antother, using the given function. The resulting tree is the structural intersection
333       * of the two trees.
334       *
335       * @param bs A tree to zip this tree with.
336       * @param f  A function with which to zip together the two trees.
337       * @return A new tree of the results of applying the given function over this tree and the given tree, position-wise.
338       */
339      public <B, C> Tree<C> zipWith(final Tree<B> bs, final F2<A, B, C> f) {
340        return $$(f).zipTree().f(this, bs);
341      }
342    
343      /**
344       * Zips this tree with antother, using the given function. The resulting tree is the structural intersection
345       * of the two trees.
346       *
347       * @param bs A tree to zip this tree with.
348       * @param f  A function with which to zip together the two trees.
349       * @return A new tree of the results of applying the given function over this tree and the given tree, position-wise.
350       */
351      public <B, C> Tree<C> zipWith(final Tree<B> bs, final F<A, F<B, C>> f) {
352        return zipWith(bs, uncurryF2(f));
353      }
354    }