forked from scalaz/scalaz
-
Notifications
You must be signed in to change notification settings - Fork 0
/
Tree.scala
166 lines (133 loc) · 6.04 KB
/
Tree.scala
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
package scalaz
import std.stream.{streamInstance, streamMonoid}
import std.string.stringInstance
/**
* A multi-way tree, also known as a rose tree. Also known as Cofree[Stream, A].
*/
sealed abstract class Tree[A] {
import Tree._
/** The label at the root of this tree. */
def rootLabel: A
/** The child nodes of this tree. */
def subForest: Stream[Tree[A]]
/** Maps the elements of the Tree into a Monoid and folds the resulting Tree. */
def foldMap[B: Monoid](f: A => B): B =
Monoid[B].append(f(rootLabel), Foldable[Stream].foldMap[Tree[A], B](subForest)((_: Tree[A]).foldMap(f)))
def foldRight[B](z: => B)(f: (A, => B) => B): B =
Foldable[Stream].foldRight(flatten, z)(f)
/** A 2D String representation of this Tree. */
def drawTree(implicit sh: Show[A]): String =
Foldable[Stream].foldMap(draw)((_: String) + "\n")
/** A histomorphic transform. Each element in the resulting tree
* is a function of the corresponding element in this tree
* and the histomorphic transform of its children.
**/
def scanr[B](g: (A, Stream[Tree[B]]) => B): Tree[B] = {
lazy val c = subForest.map(_.scanr(g))
node(g(rootLabel, c), c)
}
/** A 2D String representation of this Tree, separated into lines. */
def draw(implicit sh: Show[A]): Stream[String] = {
def drawSubTrees(s: Stream[Tree[A]]): Stream[String] = s match {
case Stream.Empty => Stream.Empty
case Stream(t) => "|" #:: shift("`- ", " ", t.draw)
case t #:: ts => "|" #:: shift("+- ", "| ", t.draw) append drawSubTrees(ts)
}
def shift(first: String, other: String, s: Stream[String]): Stream[String] =
(first #:: Stream.continually(other)).zip(s).map {
case (a, b) => a + b
}
sh.shows(rootLabel) #:: drawSubTrees(subForest)
}
/** Pre-order traversal. */
def flatten: Stream[A] = {
def squish(tree: Tree[A], xs: Stream[A]): Stream[A] =
Stream.cons(tree.rootLabel, Foldable[Stream].foldr[Tree[A], Stream[A]](tree.subForest, xs)(a => b => squish(a, b)))
squish(this, Stream.Empty)
}
/** Breadth-first traversal. */
def levels: Stream[Stream[A]] = {
val f = (s: Stream[Tree[A]]) => {
Foldable[Stream].foldMap(s)((_: Tree[A]).subForest)
}
Stream.iterate(Stream(this))(f) takeWhile (!_.isEmpty) map (_ map (_.rootLabel))
}
/** Binds the given function across all the subtrees of this tree. */
def cobind[B](f: Tree[A] => B): Tree[B] = unfoldTree(this)(t => (f(t), () => t.subForest))
/** A TreeLoc zipper of this tree, focused on the root node. */
def loc: TreeLoc[A] = TreeLoc.loc(this, Stream.Empty, Stream.Empty, Stream.Empty)
/** Turns a tree of pairs into a pair of trees. */
def unzip[A1, A2](implicit p: A => (A1, A2)): (Tree[A1], Tree[A2]) = {
lazy val uz = subForest.map(_.unzip)
lazy val fst = uz map (_._1)
lazy val snd = uz map (_._2)
(node(rootLabel._1, fst), node(rootLabel._2, snd))
}
def foldNode[Z](f: A => Stream[Tree[A]] => Z): Z =
f(rootLabel)(subForest)
def map[B](f: A => B): Tree[B] =
node(f(rootLabel), subForest map (_ map f))
def flatMap[B](f: A => Tree[B]): Tree[B] = {
val r: Tree[B] = f(rootLabel)
Tree.node(r.rootLabel, r.subForest #::: subForest.map(_.flatMap(f)))
}
def traverse1[G[_] : Apply, B](f: A => G[B]): G[Tree[B]] = {
val G = Apply[G]
import Stream._
subForest match {
case Empty => G.map(f(rootLabel))(Tree(_))
case x #:: xs => G.apply2(f(rootLabel), NonEmptyList.nel(x, xs.toList).traverse1(_.traverse1(f))) {
case (h, t) => Tree.node(h, t.list.toStream)
}
}
}
}
object Tree extends TreeInstances with TreeFunctions {
/** Construct a tree node with no children. */
def apply[A](root: => A): Tree[A] = leaf(root)
object Node {
def unapply[A](t: Tree[A]): Option[(A, Stream[Tree[A]])] = Some((t.rootLabel, t.subForest))
}
}
sealed abstract class TreeInstances {
implicit val treeInstance: Traverse1[Tree] with Monad[Tree] with Comonad[Tree] = new Traverse1[Tree] with Monad[Tree] with Comonad[Tree] {
def point[A](a: => A): Tree[A] = Tree.leaf(a)
def cobind[A, B](fa: Tree[A])(f: Tree[A] => B): Tree[B] = fa cobind f
def copoint[A](p: Tree[A]): A = p.rootLabel
override def map[A, B](fa: Tree[A])(f: A => B) = fa map f
def bind[A, B](fa: Tree[A])(f: A => Tree[B]): Tree[B] = fa flatMap f
def traverse1Impl[G[_]: Apply, A, B](fa: Tree[A])(f: A => G[B]): G[Tree[B]] = fa traverse1 f
override def foldRight[A, B](fa: Tree[A], z: => B)(f: (A, => B) => B): B = fa.foldRight(z)(f)
override def foldMapRight1[A, B](fa: Tree[A])(z: A => B)(f: (A, => B) => B): B = fa.subForest.foldRight(z(fa.rootLabel))((t, b) => treeInstance.foldRight(t, b)(f))
override def foldLeft[A, B](fa: Tree[A], z: B)(f: (B, A) => B): B =
fa.flatten.foldLeft(z)(f)
override def foldMapLeft1[A, B](fa: Tree[A])(z: A => B)(f: (B, A) => B): B = fa.flatten match {
case h #:: t => t.foldLeft(z(h))(f)
}
override def foldMap[A, B](fa: Tree[A])(f: A => B)(implicit F: Monoid[B]): B = fa foldMap f
}
implicit def treeEqual[A](implicit A: Equal[A]): Equal[Tree[A]] = new Equal[Tree[A]] {
def equal(a1: Tree[A], a2: Tree[A]): Boolean = {
A.equal(a1.rootLabel, a2.rootLabel) && a1.subForest.corresponds(a2.subForest)(equal _)
}
}
/* TODO
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)
*/
}
trait TreeFunctions {
/** Construct a new Tree node. */
def node[A](root: => A, forest: => Stream[Tree[A]]): Tree[A] = new Tree[A] {
lazy val rootLabel = root
lazy val subForest = forest
override def toString = "<tree>"
}
/** Construct a tree node with no children. */
def leaf[A](root: => A): Tree[A] = node(root, Stream.empty)
def unfoldForest[A, B](s: Stream[A])(f: A => (B, () => Stream[A])): Stream[Tree[B]] =
s.map(unfoldTree(_)(f))
def unfoldTree[A, B](v: A)(f: A => (B, () => Stream[A])): Tree[B] =
f(v) match {
case (a, bs) => node(a, unfoldForest(bs.apply())(f))
}
}