/
Util.scala
140 lines (116 loc) · 4.15 KB
/
Util.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
package scala.slick.ast
import scala.language.implicitConversions
import scala.collection.mutable.{HashSet, ArrayBuffer}
/**
* Utility methods for AST manipulation.
*/
object Util {
def mapOrNone[A <: AnyRef](c: Traversable[A])(f: A => A): Option[IndexedSeq[A]] = {
val b = new ArrayBuffer[A]
var changed = false
c.foreach { x =>
val n = f(x)
b += n
if(n ne x) changed = true
}
if(changed) Some(b.result()) else None
}
@inline implicit def nodeToNodeOps(n: Node): NodeOps = new NodeOps(n)
}
/** A scope for scoped traversal */
case class Scope(m: Map[Symbol, (Node, Scope)]) {
def get(s: Symbol) = m.get(s)
def + (s: Symbol, n: Node) = Scope(m + (s -> (n, this)))
}
object Scope {
val empty = Scope(Map())
}
/** Extra methods for Nodes. */
final class NodeOps(val tree: Node) extends AnyVal {
import Util._
@inline def collect[T](pf: PartialFunction[Node, T]): Iterable[T] = NodeOps.collect(tree, pf)
def collectAll[T](pf: PartialFunction[Node, Seq[T]]): Iterable[T] = collect[Seq[T]](pf).flatten
def replace(f: PartialFunction[Node, Node], keepType: Boolean = false): Node = NodeOps.replace(tree, f, keepType)
def foreach[U](f: (Node => U)) {
def g(n: Node) {
f(n)
n.nodeChildren.foreach(g)
}
g(tree)
}
def mapChildrenWithScope(f: (Option[Symbol], Node, Scope) => Node, scope: Scope): Node = tree match {
case d: DefNode =>
var local = scope
d.nodeMapScopedChildren { (symO, ch) =>
val r = f(symO, ch, local)
symO.foreach(sym => local = local + (sym, r))
r
}
case n => n.nodeMapChildren(ch => f(None, ch, scope))
}
def findNode(p: Node => Boolean): Option[Node] = {
if(p(tree)) Some(tree)
else {
val it = tree.nodeChildren.iterator.map(_.findNode(p)).dropWhile(_.isEmpty)
if(it.hasNext) it.next() else None
}
}
def select(field: Symbol): Node = (field, tree) match {
case (s: AnonSymbol, StructNode(ch)) => ch.find{ case (s2,_) => s == s2 }.get._2
case (s: FieldSymbol, StructNode(ch)) => ch.find{ case (s2,_) => s == s2 }.get._2
case (s: ElementSymbol, ProductNode(ch)) => ch(s.idx-1)
}
}
object NodeOps {
import Util._
// These methods should be in the class but 2.10.0-RC1 took away the ability
// to use closures in value classes
def collect[T](tree: Node, pf: PartialFunction[Node, T]): Iterable[T] = {
val b = new ArrayBuffer[T]
tree.foreach(pf.andThen[Unit]{ case t => b += t }.orElse[Node, Unit]{ case _ => () })
b
}
def replace(tree: Node, f: PartialFunction[Node, Node], keepType: Boolean): Node =
f.applyOrElse(tree, ({ case n: Node => n.nodeMapChildren(_.replace(f, keepType), keepType) }): PartialFunction[Node, Node])
}
/** Some less general but still useful methods for the code generators. */
object ExtraUtil {
def findPaths(startingAt: Set[Symbol], n: Node) = {
val b = new HashSet[Node]
def f(n: Node): Unit = n match {
case p @ Path(syms) if startingAt contains syms.last => b += p
case n => n.nodeChildren.foreach(f)
}
f(n)
b.toSet
}
def hasRowNumber(n: Node): Boolean = n match {
case c: Comprehension => false
case r: RowNumber => true
case n => n.nodeChildren.exists(hasRowNumber)
}
def replaceRowNumber(n: Node)(f: RowNumber => Node): Node = n match {
case c: Comprehension => c
case r: RowNumber => f(r)
case n => n.nodeMapChildren(ch => replaceRowNumber(ch)(f))
}
def linearizeFieldRefs(n: Node): IndexedSeq[Node] = {
val sels = new ArrayBuffer[Node]
def f(n: Node): Unit = n match {
case Path(_) => sels += n
case _: ProductNode | _: OptionApply | _: GetOrElse | _: TypeMapping | _: ClientSideOp =>
n.nodeChildren.foreach(f)
}
f(n)
sels
}
}
object ProductOfCommonPaths {
def unapply(n: ProductNode): Option[(Symbol, Vector[List[Symbol]])] = if(n.nodeChildren.isEmpty) None else
n.nodeChildren.foldLeft(null: Option[(Symbol, Vector[List[Symbol]])]) {
case (None, _) => None
case (null, FwdPath(sym :: rest)) => Some((sym, Vector(rest)))
case (Some((sym0, v)), FwdPath(sym :: rest)) if sym == sym0 => Some((sym, v :+ rest))
case _ => None
}
}