Skip to content

Commit 95fbfb0

Browse files
ehighamclaude
ehigham
and
claude
authored
[query] Walk children through ChildEnv
The walkers in BaseIR take [E[_]: Scopes, A] and a v in place of the GenericBindingEnv bound or update, and call ChildEnv.get. v makes the value each name is bound to. TypeCheck, Parser.annotateTypes, PruneDeadFields.rebuildIR and NormalizeNames pass their own v. VisitIR.withEnv takes the same parameters. FreeVariableBindingEnv gets a Scopes instance, through a phantom type E[_], in place of extend, and its classes shrink. FreeVariables finds the same names: a fresh scope empties eval, agg and scan, as the old extend did. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
1 parent 730740d commit 95fbfb0

7 files changed

Lines changed: 95 additions & 101 deletions

File tree

‎hail/hail/src/is/hail/expr/ir/BaseIR.scala‎

Lines changed: 26 additions & 19 deletions
Original file line numberDiff line numberDiff line change
@@ -65,25 +65,29 @@ abstract class BaseIR {
6565
}
6666
}
6767

68-
def forEachChildWithEnv[E <: GenericBindingEnv[E, Type]](env: E)(f: (BaseIR, E) => Unit): Unit =
68+
def forEachChildWithEnv[E[_]: Scopes, A](
69+
env: E[A]
70+
)(
71+
v: (Name, Type) => A
72+
)(
73+
f: (BaseIR, E[A]) => Unit
74+
): Unit =
6975
childrenSeq.view.zipWithIndex.foreach { case (child, i) =>
70-
val childEnv = env.extend(Bindings.get(this, i))
76+
val childEnv = ChildEnv.get(env, this, i)(v)
7177
f(child, childEnv)
7278
}
7379

74-
def mapChildrenWithEnv(env: BindingEnv[Type])(f: (BaseIR, BindingEnv[Type]) => BaseIR): BaseIR =
75-
mapChildrenWithEnv[BindingEnv[Type]](env, (env, bindings) => env.extend(bindings))(f)
76-
77-
def mapChildrenWithEnv[E](
78-
env: E,
79-
update: (E, Bindings[Type]) => E,
80+
def mapChildrenWithEnv[E[_]: Scopes, A](
81+
env: E[A]
8082
)(
81-
f: (BaseIR, E) => BaseIR
83+
v: (Name, Type) => A
84+
)(
85+
f: (BaseIR, E[A]) => BaseIR
8286
): BaseIR = {
8387
val newChildren = childrenSeq.toArray
8488
var res = this
8589
for (i <- newChildren.indices) {
86-
val childEnv = update(env, Bindings.get(res, i))
90+
val childEnv = ChildEnv.get(env, res, i)(v)
8791
val child = newChildren(i)
8892
val newChild = f(child, childEnv)
8993
if (!(newChild eq child)) {
@@ -94,16 +98,17 @@ abstract class BaseIR {
9498
res
9599
}
96100

97-
def mapChildrenWithEnvStackSafe[E](
98-
env: E,
99-
update: (E, Bindings[Type]) => E,
101+
def mapChildrenWithEnvStackSafe[E[_]: Scopes, A](
102+
env: E[A]
103+
)(
104+
v: (Name, Type) => A
100105
)(
101-
f: (BaseIR, E) => StackFrame[BaseIR]
106+
f: (BaseIR, E[A]) => StackFrame[BaseIR]
102107
): StackFrame[BaseIR] = {
103108
val newChildren = childrenSeq.toArray
104109
var res = this
105110
newChildren.indices.foreachRecur { i =>
106-
val childEnv = update(env, Bindings.get(res, i))
111+
val childEnv = ChildEnv.get(env, res, i)(v)
107112
val child = newChildren(i)
108113
f(child, childEnv).map { newChild =>
109114
if (!(newChild eq child)) {
@@ -114,13 +119,15 @@ abstract class BaseIR {
114119
}.map(_ => res)
115120
}
116121

117-
def forEachChildWithEnvStackSafe[E <: GenericBindingEnv[E, Type]](
118-
env: E
122+
def forEachChildWithEnvStackSafe[E[_]: Scopes, A](
123+
env: E[A]
124+
)(
125+
v: (Name, Type) => A
119126
)(
120-
f: (BaseIR, Int, E) => StackFrame[Unit]
127+
f: (BaseIR, Int, E[A]) => StackFrame[Unit]
121128
): StackFrame[Unit] =
122129
childrenSeq.view.zipWithIndex.foreachRecur { case (child, i) =>
123-
val childEnv = env.extend(Bindings.get(this, i))
130+
val childEnv = ChildEnv.get(env, this, i)(v)
124131
f(child, i, childEnv)
125132
}
126133
}

‎hail/hail/src/is/hail/expr/ir/FreeVariables.scala‎

Lines changed: 48 additions & 68 deletions
Original file line numberDiff line numberDiff line change
@@ -1,60 +1,26 @@
11
package is.hail.expr.ir
22

33
import is.hail.expr.ir.defs.Ref
4-
import is.hail.types.virtual.Type
54

65
import scala.collection.mutable
76

8-
case class FreeVariableEnv(boundVars: Env[Unit], freeVars: mutable.Set[Name]) {
9-
def this(boundVars: Env[Unit]) =
10-
this(boundVars, mutable.Set.empty)
11-
12-
private def copy(boundVars: Env[Unit]): FreeVariableEnv =
13-
new FreeVariableEnv(boundVars, freeVars)
14-
7+
case class FreeVariableEnv(boundVars: Set[Name], freeVars: mutable.Set[Name]) {
158
def visitRef(name: Name): Unit =
169
if (!boundVars.contains(name))
1710
freeVars += name
1811

19-
def bindIterable(bindings: Seq[(Name, Type)]): FreeVariableEnv =
20-
copy(boundVars.bindIterable(bindings.view.map(b => (b._1, ()))))
12+
def bindIterable(bindings: Seq[(Name, Any)]): FreeVariableEnv =
13+
copy(boundVars = boundVars ++ bindings.view.map(_._1))
2114

2215
def getFreeVars: Env[Unit] = new Env[Unit].bindIterable(freeVars.view.map(n => (n, ())))
2316
}
2417

18+
// The free variables of eval, agg and scan.
2519
case class FreeVariableBindingEnv(
2620
evalVars: Option[FreeVariableEnv],
2721
aggVars: Option[FreeVariableEnv],
2822
scanVars: Option[FreeVariableEnv],
29-
) extends GenericBindingEnv[FreeVariableBindingEnv, Type] {
30-
override def extend(bindings: Bindings[Type]): FreeVariableBindingEnv = {
31-
val Bindings(all, eval, agg, scan, relational, dropEval) = bindings
32-
var newEnv = this
33-
if (dropEval) newEnv = newEnv.noEval
34-
agg match {
35-
case AggEnv.Drop => newEnv = newEnv.noAgg
36-
case AggEnv.Promote => newEnv = newEnv.promoteAgg
37-
case AggEnv.Create(bindings) =>
38-
newEnv = newEnv.createAgg.bindAgg(bindings.map(all): _*)
39-
case AggEnv.Bind(bindings) =>
40-
newEnv = newEnv.bindAgg(bindings.map(all): _*)
41-
case _ =>
42-
}
43-
scan match {
44-
case AggEnv.Drop => newEnv = newEnv.noScan
45-
case AggEnv.Promote => newEnv = newEnv.promoteScan
46-
case AggEnv.Create(bindings) =>
47-
newEnv = newEnv.createScan.bindScan(bindings.map(all): _*)
48-
case AggEnv.Bind(bindings) =>
49-
newEnv = newEnv.bindScan(bindings.map(all): _*)
50-
case _ =>
51-
}
52-
if (eval.nonEmpty) newEnv = newEnv.bindEval(eval.map(all): _*)
53-
if (relational.nonEmpty)
54-
newEnv = newEnv.bindRelational(relational.map(all): _*)
55-
newEnv
56-
}
57-
23+
) {
5824
def visitRef(name: Name): Unit =
5925
evalVars.foreach(_.visitRef(name))
6026

@@ -63,44 +29,58 @@ case class FreeVariableBindingEnv(
6329
aggVars.map(_.getFreeVars),
6430
scanVars.map(_.getFreeVars),
6531
)
32+
}
6633

67-
override def promoteAgg: FreeVariableBindingEnv =
68-
copy(evalVars = aggVars, aggVars = None)
69-
70-
override def promoteScan: FreeVariableBindingEnv =
71-
copy(evalVars = scanVars, scanVars = None)
72-
73-
override def bindEval(bindings: (Name, Type)*): FreeVariableBindingEnv =
74-
copy(evalVars = evalVars.map(_.bindIterable(bindings)))
75-
76-
override def noEval: FreeVariableBindingEnv = copy(evalVars = None)
77-
78-
override def bindAgg(bindings: (Name, Type)*): FreeVariableBindingEnv =
79-
copy(aggVars = aggVars.map(_.bindIterable(bindings)))
80-
81-
override def bindScan(bindings: (Name, Type)*): FreeVariableBindingEnv =
82-
copy(scanVars = scanVars.map(_.bindIterable(bindings)))
83-
84-
override def createAgg: FreeVariableBindingEnv = copy(aggVars = evalVars)
34+
object FreeVariableBindingEnv {
35+
type E[_] = FreeVariableBindingEnv
36+
37+
// As BindingEnv's, but binding or promoting a missing scope does nothing, and relational
38+
// names aren't tracked
39+
implicit object freeVariableScopes extends Scopes[E] {
40+
override def bind[A](e: E[A], s: EnvScope, bs: Seq[(Name, A)]): E[A] =
41+
s match {
42+
case Scope.EVAL => e.copy(evalVars = e.evalVars.map(_.bindIterable(bs)))
43+
case Scope.AGG => e.copy(aggVars = e.aggVars.map(_.bindIterable(bs)))
44+
case Scope.SCAN => e.copy(scanVars = e.scanVars.map(_.bindIterable(bs)))
45+
case Scope.REL => e
46+
}
8547

86-
override def createScan: FreeVariableBindingEnv = copy(scanVars = evalVars)
48+
override def create[A](e: E[A], s: Scope): E[A] =
49+
s match {
50+
case Scope.AGG => e.copy(aggVars = e.evalVars)
51+
case Scope.SCAN => e.copy(scanVars = e.evalVars)
52+
case Scope.EVAL => throw new IllegalArgumentException("can't create eval")
53+
}
8754

88-
override def noAgg: FreeVariableBindingEnv = copy(aggVars = None)
55+
override def drop[A](e: E[A], s: Scope): E[A] =
56+
s match {
57+
case Scope.EVAL => e.copy(evalVars = None)
58+
case Scope.AGG => e.copy(aggVars = None)
59+
case Scope.SCAN => e.copy(scanVars = None)
60+
}
8961

90-
override def noScan: FreeVariableBindingEnv = copy(scanVars = None)
62+
override def promote[A](e: E[A], s: Scope): E[A] =
63+
s match {
64+
case Scope.EVAL => e
65+
case Scope.AGG => e.copy(evalVars = e.aggVars, aggVars = None)
66+
case Scope.SCAN => e.copy(evalVars = e.scanVars, scanVars = None)
67+
}
9168

92-
override def bindRelational(bindings: (Name, Type)*): FreeVariableBindingEnv =
93-
this
69+
override def fresh[A](e: E[A]): E[A] =
70+
FreeVariableBindingEnv(None, None, None)
71+
}
9472
}
9573

9674
object FreeVariables {
9775
def apply(ir: IR, supportsAgg: Boolean, supportsScan: Boolean): BindingEnv[Unit] = {
98-
val env = FreeVariableBindingEnv(
99-
Some(new FreeVariableEnv(Env.empty)),
100-
if (supportsAgg) Some(new FreeVariableEnv(Env.empty)) else None,
101-
if (supportsScan) Some(new FreeVariableEnv(Env.empty)) else None,
102-
)
103-
VisitIR.withEnv(ir, env) { (ir, env) =>
76+
val env =
77+
FreeVariableBindingEnv(
78+
Some(FreeVariableEnv(Set.empty, mutable.Set.empty)),
79+
Option.when(supportsAgg)(FreeVariableEnv(Set.empty, mutable.Set.empty)),
80+
Option.when(supportsScan)(FreeVariableEnv(Set.empty, mutable.Set.empty)),
81+
)
82+
83+
VisitIR.withEnv[FreeVariableBindingEnv.E, Unit](ir, env)((_, _) => ()) { (ir, env) =>
10484
ir match {
10585
case Ref(name, _) => env.visitRef(name)
10686
case _ =>

‎hail/hail/src/is/hail/expr/ir/MapIR.scala‎

Lines changed: 9 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -22,8 +22,15 @@ object VisitIR {
2222
ir.children.foreach(apply(_)(f))
2323
}
2424

25-
def withEnv[E <: GenericBindingEnv[E, Type]](ir: BaseIR, env: E)(f: (BaseIR, E) => Unit): Unit = {
25+
def withEnv[E[_]: Scopes, A](
26+
ir: BaseIR,
27+
env: E[A],
28+
)(
29+
v: (Name, Type) => A
30+
)(
31+
f: (BaseIR, E[A]) => Unit
32+
): Unit = {
2633
f(ir, env)
27-
ir.forEachChildWithEnv(env)(withEnv(_, _)(f))
34+
ir.forEachChildWithEnv(env)(v)(withEnv(_, _)(v)(f))
2835
}
2936
}

‎hail/hail/src/is/hail/expr/ir/NormalizeNames.scala‎

Lines changed: 7 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,9 @@
11
package is.hail.expr.ir
22

33
import is.hail.backend.ExecuteContext
4+
import is.hail.collection.FastSeq
45
import is.hail.expr.ir.NormalizeNames.needsRenaming
6+
import is.hail.expr.ir.Scopes.ScopesSyntax
57
import is.hail.expr.ir.defs._
68
import is.hail.expr.ir.lowering.invariant.{NoSharedNodes, UniquelyNamed}
79
import is.hail.types.virtual.Type
@@ -74,7 +76,7 @@ class NormalizeNames(freeVariables: Set[Name]) {
7476
}
7577
for (newValue <- normalizeIR(value, valueEnv)) yield {
7678
val newName = gen()
77-
newEnv = newEnv.bindInScope(name, newName, scope)
79+
newEnv = newEnv.bind(scope, FastSeq(name -> newName))
7880
Binding(newName, newValue.asInstanceOf[IR], scope)
7981
}
8082
}
@@ -108,12 +110,10 @@ class NormalizeNames(freeVariables: Set[Name]) {
108110
Recur(newName, args, typ).mapChildrenStackSafe(normalizeIR(_, env))
109111
case ir =>
110112
val bindingsMap = mutable.HashMap.empty[Name, Name]
111-
val updateEnv: (BindingEnv[Name], Bindings[Type]) => BindingEnv[Name] =
112-
if (needsRenaming(ir)) { (env, bindings) =>
113-
val bindingsNames = bindings.map((name, _) => bindingsMap.getOrElseUpdate(name, gen()))
114-
env.extend(bindingsNames)
115-
} else { (env, bindings) => env.extend(bindings.map((name, _) => name)) }
116-
ir.mapChildrenWithEnvStackSafe(env, updateEnv)(normalizeIR).map(modifyNames(_, bindingsMap))
113+
val v: (Name, Type) => Name =
114+
if (needsRenaming(ir)) (name, _) => bindingsMap.getOrElseUpdate(name, gen())
115+
else (name, _) => name
116+
ir.mapChildrenWithEnvStackSafe(env)(v)(normalizeIR).map(modifyNames(_, bindingsMap))
117117
}
118118

119119
private def modifyNames(ir: BaseIR, map: collection.Map[Name, Name]): BaseIR = {

‎hail/hail/src/is/hail/expr/ir/Parser.scala‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -2044,7 +2044,7 @@ object IRParser {
20442044

20452045
def annotateTypes(ctx: ExecuteContext, ir: BaseIR, env: BindingEnv[Type]): BaseIR = {
20462046
def run(ir: BaseIR, env: BindingEnv[Type]): BaseIR = {
2047-
val rw = ir.mapChildrenWithEnv(env)(run)
2047+
val rw = ir.mapChildrenWithEnv(env)((_, t) => t)(run)
20482048
rw match {
20492049
case x: Ref =>
20502050
x._typ = env.eval(x.name)

‎hail/hail/src/is/hail/expr/ir/PruneDeadFields.scala‎

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -2327,11 +2327,11 @@ object PruneDeadFields extends Logging {
23272327
))
23282328
}
23292329
case x: ApplyAggOp =>
2330-
x.mapChildrenWithEnv(env) { (child, childEnv) =>
2330+
x.mapChildrenWithEnv(env)((_, t) => t) { (child, childEnv) =>
23312331
rebuildIR(ctx, child.asInstanceOf[IR], childEnv, memo)
23322332
}.asInstanceOf[ApplyAggOp]
23332333
case x: ApplyScanOp =>
2334-
x.mapChildrenWithEnv(env) { (child, childEnv) =>
2334+
x.mapChildrenWithEnv(env)((_, t) => t) { (child, childEnv) =>
23352335
rebuildIR(ctx, child.asInstanceOf[IR], childEnv, memo)
23362336
}.asInstanceOf[ApplyScanOp]
23372337
case CollectDistributedArray(contexts, globals, cname, gname, body, dynamicID, staticID,
@@ -2355,7 +2355,7 @@ object PruneDeadFields extends Logging {
23552355
val dynamicID2 = rebuildIR(ctx, dynamicID, env, memo)
23562356
CollectDistributedArray(contexts2, globals2, cname, gname, body2, dynamicID2, staticID, tsd)
23572357
case _ =>
2358-
ir.mapChildrenWithEnv(env) { (child, childEnv) =>
2358+
ir.mapChildrenWithEnv(env)((_, t) => t) { (child, childEnv) =>
23592359
child match {
23602360
case valueIR: IR => rebuildIR(ctx, valueIR, childEnv, memo)
23612361
case mir: MatrixIR => rebuild(ctx, mir, memo)

‎hail/hail/src/is/hail/expr/ir/TypeCheck.scala‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -32,7 +32,7 @@ object TypeCheck extends Logging {
3232

3333
def check(ctx: ExecuteContext, ir: BaseIR, env: BindingEnv[Type]): StackFrame[Unit] = {
3434
for {
35-
_ <- ir.forEachChildWithEnvStackSafe(env) { (child, i, childEnv) =>
35+
_ <- ir.forEachChildWithEnvStackSafe(env)((_, t) => t) { (child, i, childEnv) =>
3636
for {
3737
_ <- call(check(ctx, child, childEnv))
3838
} yield

0 commit comments

Comments
 (0)