/core/src/main/scala/scalaz/Tree.scala

  1package scalaz
  2
  3import scalaz.Free.Trampoline
  4import scalaz.Trampoline._
  5import scalaz.EphemeralStream._
  6
  7/**
  8 * A multi-way tree, also known as a rose tree. Also known as Cofree[Stream, A].
  9 */
 10sealed abstract class Tree[A] {
 11
 12  import Tree._
 13
 14  /** The label at the root of this tree. */
 15  def rootLabel: A
 16
 17  /** The child nodes of this tree. */
 18  def subForest: EStream[Tree[A]]
 19
 20  def foldMapTrampoline[B: Monoid](f: A => B): Trampoline[B] = {
 21    for {
 22      root <- delay(f(rootLabel))
 23      subForests <- Foldable[EStream].foldMap[Tree[A], Trampoline[B]](subForest)(_.foldMapTrampoline(f))
 24    } yield Monoid[B].append(root, subForests)
 25  }
 26
 27  /** Maps the elements of the Tree into a Monoid and folds the resulting Tree. */
 28  def foldMap[B: Monoid](f: A => B): B =
 29    foldMapTrampoline[B](f).run
 30
 31  def foldRight[B](z: => B)(f: (A, => B) => B): B =
 32    Foldable[EStream].foldRight(flatten, z)(f)
 33
 34  /** A 2D String representation of this Tree. */
 35  def drawTree(implicit sh: Show[A]): String = {
 36    val reversedLines = draw.run
 37    val first = new StringBuilder(reversedLines.head.toString.reverse)
 38    val rest = reversedLines.tail
 39    rest.foldLeft(first) { (acc, elem) =>
 40      acc.append("\n").append(elem.toString.reverse)
 41    }.append("\n").toString
 42  }
 43
 44  /** A histomorphic transform. Each element in the resulting tree
 45   * is a function of the corresponding element in this tree
 46   * and the histomorphic transform of its children.
 47   **/
 48  def scanr[B](g: (A, EStream[Tree[B]]) => B): Tree[B] = {
 49    val c = Need(subForest.map(_.scanr(g)))
 50    Node(g(rootLabel, c.value), c.value)
 51  }
 52
 53  /** A 2D String representation of this Tree, separated into lines.
 54    * Uses reversed StringBuilders for performance, because they are
 55    * prepended to.
 56    **/
 57  private def draw(implicit sh: Show[A]): Trampoline[Vector[StringBuilder]] = {
 58    import Trampoline._
 59    val branch = " -+" // "+- ".reverse
 60    val stem = " -`" // "`- ".reverse
 61    val trunk = "  |" // "|  ".reverse
 62
 63    def drawSubTrees(s: EStream[Tree[A]]): Trampoline[Vector[StringBuilder]] = s match {
 64      case ts if ts.isEmpty       => 
 65        done(Vector.empty[StringBuilder])
 66      case t ##:: ts if ts.isEmpty =>
 67        suspend(t.draw).map(subtree => new StringBuilder("|") +: shift(stem, "   ", subtree))
 68      case t ##:: ts               => for {
 69                                       subtree <- suspend(t.draw)
 70                                       otherSubtrees <- suspend(drawSubTrees(ts))
 71                                     } yield new StringBuilder("|") +: (shift(branch, trunk, subtree) ++ otherSubtrees)
 72    }
 73
 74    def shift(first: String, other: String, s: Vector[StringBuilder]): Vector[StringBuilder] = {
 75      var i = 0
 76      while (i < s.length) {
 77        if (i == 0) s(i).append(first)
 78        else s(i).append(other)
 79        i += 1
 80      }
 81      s
 82    }
 83
 84    drawSubTrees(subForest).map { subtrees =>
 85      new StringBuilder(sh.shows(rootLabel).reverse) +: subtrees
 86    }
 87  }
 88
 89  /** Pre-order traversal. */
 90  def flatten: EStream[A] = {
 91    def squish(tree: Tree[A], xs: EStream[A]): EStream[A] =
 92      EphemeralStream.cons(tree.rootLabel, Foldable[EStream].foldRight(tree.subForest, xs)(squish(_, _)))
 93
 94    squish(this, EphemeralStream.emptyEphemeralStream)
 95  }
 96
 97  /** Breadth-first traversal. */
 98  def levels: EStream[EStream[A]] = {
 99    val f = (s: EStream[Tree[A]]) => {
100      Foldable[EStream].foldMap(s)((_: Tree[A]).subForest)
101    }
102    EphemeralStream.iterate(EphemeralStream(this))(f) takeWhile (!_.isEmpty) map (_ map (_.rootLabel))
103  }
104
105  def toStrictTree: StrictTree[A] = {
106    import std.vector.vectorInstance
107
108    def trampolined(t: Tree[A]): Trampoline[StrictTree[A]] = {
109      t match {
110        case Tree.Leaf(root) =>
111          Trampoline.done(StrictTree.Leaf(root))
112        case Tree.Node(root, forest) =>
113          for {
114            strictForest <- Trampoline.suspend(Applicative[Trampoline].traverse(forest.toIList.toVector)(trampolined))
115          } yield StrictTree(root, strictForest)
116      }
117    }
118
119    trampolined(this).run
120  }
121
122  /** Binds the given function across all the subtrees of this tree. */
123  def cobind[B](f: Tree[A] => B): Tree[B] = unfoldTree(this)(t => (f(t), () => t.subForest))
124
125  /** A TreeLoc zipper of this tree, focused on the root node. */
126  def loc: TreeLoc[A] = TreeLoc.loc(this, emptyEphemeralStream, emptyEphemeralStream, emptyEphemeralStream)
127
128  /** Turns a tree of pairs into a pair of trees. */
129  def unzip[A1, A2](p: A => (A1, A2)): (Tree[A1], Tree[A2]) = {
130    val uz = Need(subForest.map(_.unzip(p)))
131    val fst = Need(uz.value map (_._1))
132    val snd = Need(uz.value map (_._2))
133    val (a, b) = p(rootLabel)
134    (Node(a, fst.value), Node(b, snd.value))
135  }
136
137  def foldNode[Z](f: A => EStream[Tree[A]] => Z): Z =
138    f(rootLabel)(subForest)
139
140  def map[B](f: A => B): Tree[B] =
141    Node(f(rootLabel), subForest map (_ map f))
142
143  def flatMap[B](f: A => Tree[B]): Tree[B] = {
144    val r: Tree[B] = f(rootLabel)
145    Node(r.rootLabel, r.subForest ++ subForest.map(_.flatMap(f)))
146  }
147
148  def traverse1[G[_] : Apply, B](f: A => G[B]): G[Tree[B]] = {
149    val G = Apply[G]
150    import EphemeralStream._
151
152    std.boolean.fold(subForest.isEmpty, 
153                 G.map(f(rootLabel))(Leaf(_)),
154                 G.apply2 (f(rootLabel),
155                           (OneAnd.oneAndTraverse[EStream]
156                              .traverse1
157                                (OneAnd(
158                                  subForest.headOption.get, subForest.tailOption.get))
159                                (_.traverse1(f))))
160                           { case (h, t) => Node(h, t.head ##:: t.tail) }
161                          )
162  }
163
164}
165
166sealed abstract class TreeInstances {
167  implicit val treeInstance: Traverse1[Tree] with Monad[Tree] with Comonad[Tree] with Align[Tree] with Zip[Tree] = new Traverse1[Tree] with Monad[Tree] with Comonad[Tree] with Align[Tree] with Zip[Tree] {
168    def point[A](a: => A): Tree[A] = Tree.Leaf(a)
169    def cobind[A, B](fa: Tree[A])(f: Tree[A] => B): Tree[B] = fa cobind f
170    def copoint[A](p: Tree[A]): A = p.rootLabel
171    override def map[A, B](fa: Tree[A])(f: A => B) = fa map f
172    def bind[A, B](fa: Tree[A])(f: A => Tree[B]): Tree[B] = fa flatMap f
173    def traverse1Impl[G[_]: Apply, A, B](fa: Tree[A])(f: A => G[B]): G[Tree[B]] = fa traverse1 f
174    override def foldRight[A, B](fa: Tree[A], z: => B)(f: (A, => B) => B): B = fa.foldRight(z)(f)
175
176    override def foldMapRight1[A, B](fa: Tree[A])(z: A => B)(f: (A, => B) => B) = (fa.flatten.reverse: @unchecked) match {
177      case h ##:: t => t.foldLeft(z(h))((b, a) => f(a, b))
178    }
179    
180    override def foldLeft[A, B](fa: Tree[A], z: B)(f: (B, A) => B): B =
181      fa.flatten.foldLeft(z)((b, a) => f(b, a))
182
183    override def foldMapLeft1[A, B](fa: Tree[A])(z: A => B)(f: (B, A) => B): B = fa.flatten match {
184      case h ##:: t => t.foldLeft(z(h))((b, a) => f(b, a))
185    }
186
187    override def foldMap[A, B](fa: Tree[A])(f: A => B)(implicit F: Monoid[B]): B = fa foldMap f
188    def alignWith[A, B, C](f: (\&/[A, B]) => C) = {
189      def align(ta: Tree[A], tb: Tree[B]): Tree[C] =
190        Tree.Node(f(\&/(ta.rootLabel, tb.rootLabel)), Align[EStream].alignWith[Tree[A], Tree[B], Tree[C]]({
191          case \&/.This(sta) => sta map {a => f(\&/.This(a))}
192          case \&/.That(stb) => stb map {b => f(\&/.That(b))}
193          case \&/(sta, stb) => align(sta, stb)
194        })(ta.subForest, tb.subForest))
195      align
196    }
197    def zip[A, B](aa: => Tree[A], bb: => Tree[B]) = {
198      val a = aa
199      val b = bb
200      Tree.Node(
201        (a.rootLabel, b.rootLabel),
202        Zip[EStream].zipWith(a.subForest, b.subForest)(zip(_, _))
203      )
204    }
205  }
206
207  implicit def treeEqual[A](implicit A0: Equal[A]): Equal[Tree[A]] =
208    new TreeEqual[A] { def A = A0 }
209
210  implicit def treeOrder[A](implicit A0: Order[A]): Order[Tree[A]] =
211    new Order[Tree[A]] with TreeEqual[A] {
212      override def A: Order[A] = A0
213      override def order(x: Tree[A], y: Tree[A]) =
214        A.order(x.rootLabel, y.rootLabel) match {
215          case Ordering.EQ =>
216            Order[EStream[Tree[A]]].order(x.subForest, y.subForest)
217          case x => x
218        }
219    }
220
221  /* TODO
222  def applic[A, B](f: Tree[A => B]) = a => Tree.node((f.rootLabel)(a.rootLabel), implicitly[Applic[newtypes.ZipStream]].applic(f.subForest.map(applic[A, B](_)).ʐ)(a.subForest ʐ).value)
223   */
224}
225
226object Tree extends TreeInstances {
227  /**
228   * Node represents a tree node that may have children.
229   *
230   * You can use Node for tree construction or pattern matching.
231   */
232  object Node {
233    def apply[A](root: => A, forest: => EStream[Tree[A]]): Tree[A] = {
234      new Tree[A] {
235        private[this] val rootc = Need(root)
236        private[this] val forestc = Need(forest)
237        def rootLabel = rootc.value
238        def subForest = forestc.value
239
240        override def toString = "<tree>"
241      }
242    }
243
244    def unapply[A](t: Tree[A]): Option[(A, EStream[Tree[A]])] = Some((t.rootLabel, t.subForest))
245  }
246
247  /**
248   *  Leaf represents a tree node with no children.
249   *
250   *  You can use Leaf for tree construction or pattern matching.
251   */
252  object Leaf {
253    def apply[A](root: => A): Tree[A] = {
254      Node(root, EphemeralStream.emptyEphemeralStream)
255    }
256
257    def unapply[A](t: Tree[A]): Option[A] = {
258      t match {
259        case Node(root, xs) if xs.isEmpty =>
260          Some(root)
261        case _ =>
262          None
263      }
264    }
265  }
266
267  def unfoldForest[A, B](s: EStream[A])(f: A => (B, () => EStream[A])): EStream[Tree[B]] =
268    s.map(unfoldTree(_)(f))
269
270  def unfoldTree[A, B](v: A)(f: A => (B, () => EStream[A])): Tree[B] =
271    f(v) match {
272      case (a, bs) => Node(a, unfoldForest(bs.apply())(f))
273    }
274}
275
276private trait TreeEqual[A] extends Equal[Tree[A]] {
277  def A: Equal[A]
278
279  override final def equal(a1: Tree[A], a2: Tree[A]) = {
280    def corresponds[B](a1: EStream[Tree[A]], a2: EStream[Tree[A]]): Trampoline[Boolean] = {
281      (a1.isEmpty, a2.isEmpty) match {
282        case (true, true) => Trampoline.done(true)
283        case (_, true) | (true, _) => Trampoline.done(false)
284        case _ =>
285          for {
286            heads <- trampolined(a1.headOption.get, a2.headOption.get)
287            tails <- corresponds(a1.tailOption.get, a2.tailOption.get)
288          } yield heads && tails
289      }
290    }
291
292    def trampolined(a1: Tree[A], a2: Tree[A]): Trampoline[Boolean] = {
293      for {
294        roots <- Trampoline.done(A.equal(a1.rootLabel, a2.rootLabel))
295        subForests <- corresponds(a1.subForest, a2.subForest)
296      } yield roots && subForests
297    }
298
299    trampolined(a1, a2).run
300  }
301
302}