Skip to content

Commit 05f8dee

Browse files
authored
chore: Add more static optimizer (#628)
Motivation: more aggressive optimization, but behind a flag, the default is false. master: ``` 125] Benchmark (path) Mode Cnt Score Error Units 125] RegressionBenchmark.main bench/resources/bug_suite/assertions.jsonnet avgt 0.326 ms/op 125] RegressionBenchmark.main bench/resources/cpp_suite/bench.01.jsonnet avgt 0.073 ms/op 125] RegressionBenchmark.main bench/resources/cpp_suite/bench.02.jsonnet avgt 40.602 ms/op 125] RegressionBenchmark.main bench/resources/cpp_suite/bench.03.jsonnet avgt 13.352 ms/op 125] RegressionBenchmark.main bench/resources/cpp_suite/bench.04.jsonnet avgt 32.769 ms/op 125] RegressionBenchmark.main bench/resources/cpp_suite/bench.06.jsonnet avgt 0.448 ms/op 125] RegressionBenchmark.main bench/resources/cpp_suite/bench.07.jsonnet avgt 3.253 ms/op 125] RegressionBenchmark.main bench/resources/cpp_suite/bench.08.jsonnet avgt 0.062 ms/op 125] RegressionBenchmark.main bench/resources/cpp_suite/bench.09.jsonnet avgt 0.068 ms/op 125] RegressionBenchmark.main bench/resources/cpp_suite/gen_big_object.jsonnet avgt 1.012 ms/op 125] RegressionBenchmark.main bench/resources/cpp_suite/large_string_join.jsonnet avgt 2.059 ms/op 125] RegressionBenchmark.main bench/resources/cpp_suite/large_string_template.jsonnet avgt 2.439 ms/op 125] RegressionBenchmark.main bench/resources/cpp_suite/realistic1.jsonnet avgt 3.298 ms/op 125] RegressionBenchmark.main bench/resources/cpp_suite/realistic2.jsonnet avgt 125.995 ms/op 125] RegressionBenchmark.main bench/resources/go_suite/base64.jsonnet avgt 0.843 ms/op 125] RegressionBenchmark.main bench/resources/go_suite/base64Decode.jsonnet avgt 0.651 ms/op 125] RegressionBenchmark.main bench/resources/go_suite/base64DecodeBytes.jsonnet avgt 9.457 ms/op 125] RegressionBenchmark.main bench/resources/go_suite/base64_byte_array.jsonnet avgt 1.515 ms/op 125] RegressionBenchmark.main bench/resources/go_suite/comparison.jsonnet avgt 23.599 ms/op 125] RegressionBenchmark.main bench/resources/go_suite/comparison2.jsonnet avgt 75.864 ms/op 125] RegressionBenchmark.main bench/resources/go_suite/escapeStringJson.jsonnet avgt 0.050 ms/op 125] RegressionBenchmark.main bench/resources/go_suite/foldl.jsonnet avgt 9.695 ms/op 125] RegressionBenchmark.main bench/resources/go_suite/lstripChars.jsonnet avgt 0.650 ms/op 125] RegressionBenchmark.main bench/resources/go_suite/manifestJsonEx.jsonnet avgt 0.072 ms/op 125] RegressionBenchmark.main bench/resources/go_suite/manifestTomlEx.jsonnet avgt 0.089 ms/op 125] RegressionBenchmark.main bench/resources/go_suite/manifestYamlDoc.jsonnet avgt 0.076 ms/op 125] RegressionBenchmark.main bench/resources/go_suite/member.jsonnet avgt 0.731 ms/op 125] RegressionBenchmark.main bench/resources/go_suite/parseInt.jsonnet avgt 0.052 ms/op 125] RegressionBenchmark.main bench/resources/go_suite/reverse.jsonnet avgt 11.621 ms/op 125] RegressionBenchmark.main bench/resources/go_suite/rstripChars.jsonnet avgt 0.647 ms/op 125] RegressionBenchmark.main bench/resources/go_suite/stripChars.jsonnet avgt 0.632 ms/op 125] RegressionBenchmark.main bench/resources/go_suite/substr.jsonnet avgt 0.173 ms/op 125] RegressionBenchmark.main bench/resources/sjsonnet_suite/setDiff.jsonnet avgt 0.493 ms/op 125] RegressionBenchmark.main bench/resources/sjsonnet_suite/setInter.jsonnet avgt 0.443 ms/op 125] RegressionBenchmark.main bench/resources/sjsonnet_suite/setUnion.jsonnet avgt 0.792 ms/op 125/125, SUCCESS] ./mill bench.runRegressions 438s ``` this branch with `aggressiveStaticOptimization` = `true`: ``` 125] Benchmark (path) Mode Cnt Score Error Units 125] RegressionBenchmark.main bench/resources/bug_suite/assertions.jsonnet avgt 0.329 ms/op 125] RegressionBenchmark.main bench/resources/cpp_suite/bench.01.jsonnet avgt 0.072 ms/op 125] RegressionBenchmark.main bench/resources/cpp_suite/bench.02.jsonnet avgt 40.614 ms/op 125] RegressionBenchmark.main bench/resources/cpp_suite/bench.03.jsonnet avgt 13.026 ms/op 125] RegressionBenchmark.main bench/resources/cpp_suite/bench.04.jsonnet avgt 31.369 ms/op 125] RegressionBenchmark.main bench/resources/cpp_suite/bench.06.jsonnet avgt 0.433 ms/op 125] RegressionBenchmark.main bench/resources/cpp_suite/bench.07.jsonnet avgt 3.043 ms/op 125] RegressionBenchmark.main bench/resources/cpp_suite/bench.08.jsonnet avgt 0.059 ms/op 125] RegressionBenchmark.main bench/resources/cpp_suite/bench.09.jsonnet avgt 0.067 ms/op 125] RegressionBenchmark.main bench/resources/cpp_suite/gen_big_object.jsonnet avgt 0.987 ms/op 125] RegressionBenchmark.main bench/resources/cpp_suite/large_string_join.jsonnet avgt 1.985 ms/op 125] RegressionBenchmark.main bench/resources/cpp_suite/large_string_template.jsonnet avgt 2.290 ms/op 125] RegressionBenchmark.main bench/resources/cpp_suite/realistic1.jsonnet avgt 3.191 ms/op 125] RegressionBenchmark.main bench/resources/cpp_suite/realistic2.jsonnet avgt 70.614 ms/op 125] RegressionBenchmark.main bench/resources/go_suite/base64.jsonnet avgt 0.838 ms/op 125] RegressionBenchmark.main bench/resources/go_suite/base64Decode.jsonnet avgt 0.641 ms/op 125] RegressionBenchmark.main bench/resources/go_suite/base64DecodeBytes.jsonnet avgt 9.197 ms/op 125] RegressionBenchmark.main bench/resources/go_suite/base64_byte_array.jsonnet avgt 1.468 ms/op 125] RegressionBenchmark.main bench/resources/go_suite/comparison.jsonnet avgt 22.964 ms/op 125] RegressionBenchmark.main bench/resources/go_suite/comparison2.jsonnet avgt 78.190 ms/op 125] RegressionBenchmark.main bench/resources/go_suite/escapeStringJson.jsonnet avgt 0.050 ms/op 125] RegressionBenchmark.main bench/resources/go_suite/foldl.jsonnet avgt 9.088 ms/op 125] RegressionBenchmark.main bench/resources/go_suite/lstripChars.jsonnet avgt 0.653 ms/op 125] RegressionBenchmark.main bench/resources/go_suite/manifestJsonEx.jsonnet avgt 0.072 ms/op 125] RegressionBenchmark.main bench/resources/go_suite/manifestTomlEx.jsonnet avgt 0.090 ms/op 125] RegressionBenchmark.main bench/resources/go_suite/manifestYamlDoc.jsonnet avgt 0.078 ms/op 125] RegressionBenchmark.main bench/resources/go_suite/member.jsonnet avgt 0.695 ms/op 125] RegressionBenchmark.main bench/resources/go_suite/parseInt.jsonnet avgt 0.052 ms/op 125] RegressionBenchmark.main bench/resources/go_suite/reverse.jsonnet avgt 11.731 ms/op 125] RegressionBenchmark.main bench/resources/go_suite/rstripChars.jsonnet avgt 0.611 ms/op 125] RegressionBenchmark.main bench/resources/go_suite/stripChars.jsonnet avgt 0.597 ms/op 125] RegressionBenchmark.main bench/resources/go_suite/substr.jsonnet avgt 0.165 ms/op 125] RegressionBenchmark.main bench/resources/sjsonnet_suite/setDiff.jsonnet avgt 0.462 ms/op 125] RegressionBenchmark.main bench/resources/sjsonnet_suite/setInter.jsonnet avgt 0.457 ms/op 125] RegressionBenchmark.main bench/resources/sjsonnet_suite/setUnion.jsonnet avgt 0.715 ms/op 125/125, SUCCESS] ./mill bench.runRegressions 437s ``` refs: google/go-jsonnet#858
1 parent 07dc5be commit 05f8dee

File tree

6 files changed

+1399
-215
lines changed

6 files changed

+1399
-215
lines changed

build.mill

Lines changed: 3 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -78,8 +78,9 @@ trait SjsonnetCrossModule extends CrossScalaModule with ScalafmtModule {
7878
"-feature",
7979
"-opt-inline-from:sjsonnet.*,sjsonnet.**",
8080
"-Xsource:3",
81-
"-Xlint:_",
82-
) ++ (if (scalaVersion().startsWith("2.13")) Seq("-Wopt", "-Wconf:origin=scala.collection.compat.*:s")
81+
"-Xlint:_"
82+
) ++ (if (scalaVersion().startsWith("2.13"))
83+
Seq("-Wopt", "-Wconf:origin=scala.collection.compat.*:s")
8384
else Seq("-Xfatal-warnings", "-Ywarn-unused:-nowarn"))
8485
else Seq[String]("-Wconf:origin=scala.collection.compat.*:s", "-Xlint:all")
8586
)

sjsonnet/src/sjsonnet/Settings.scala

Lines changed: 8 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -13,7 +13,14 @@ final case class Settings(
1313
brokenAssertionLogic: Boolean = false,
1414
maxMaterializeDepth: Int = 1000,
1515
materializeRecursiveDepthLimit: Int = 128,
16-
maxStack: Int = 500
16+
maxStack: Int = 500,
17+
/**
18+
* Enable aggressive static optimizations in the optimization phase, including: constant folding
19+
* for arithmetic, comparison, bitwise, and shift operators; branch elimination for if-else with
20+
* constant conditions; short-circuit elimination for And/Or with constant lhs. These reduce AST
21+
* node count, benefiting long-running Jsonnet programs.
22+
*/
23+
aggressiveStaticOptimization: Boolean = false
1724
)
1825

1926
object Settings {

sjsonnet/src/sjsonnet/StaticOptimizer.scala

Lines changed: 222 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -9,6 +9,13 @@ import ScopedExprTransform.*
99
* StaticOptimizer performs necessary transformations for the evaluator (assigning ValScope indices)
1010
* plus additional optimizations (post-order) and static checking (pre-order).
1111
*
12+
* When `aggressiveStaticOptimization` is enabled, the optimizer additionally performs during the
13+
* optimization phase:
14+
* - Constant folding for arithmetic (+, -, *, /, %), comparison (<, >, <=, >=, ==, !=), bitwise
15+
* (&, ^, |), shift (<<, >>), and unary (!, -, ~, +) operators.
16+
* - Branch elimination for if-else with constant conditions.
17+
* - Short-circuit elimination for And/Or with constant lhs operands.
18+
*
1219
* @param variableResolver
1320
* a function that resolves variable names to expressions, only called if the variable is not
1421
* found in the scope.
@@ -25,6 +32,8 @@ class StaticOptimizer(
2532
extends ScopedExprTransform {
2633
def optimize(e: Expr): Expr = transform(e)
2734

35+
private val aggressiveOptimization = ev.settings.aggressiveStaticOptimization
36+
2837
override def transform(_e: Expr): Expr = super.transform(check(_e)) match {
2938
case a: Apply => transformApply(a)
3039

@@ -96,7 +105,42 @@ class StaticOptimizer(
96105
Val.staticObject(pos, fields, internedStaticFieldSets, internedStrings)
97106
else m
98107

99-
case e => e
108+
// Aggressive optimizations: constant folding, branch elimination, short-circuit elimination.
109+
// These reduce AST node count at parse time, benefiting long-running Jsonnet programs.
110+
case e => if (aggressiveOptimization) tryAggressiveOptimize(e) else e
111+
}
112+
113+
/**
114+
* Aggressive static optimizations that benefit long-running programs by reducing AST size.
115+
* Includes: branch elimination, short-circuit elimination, constant folding for arithmetic,
116+
* comparison, bitwise, and shift operators.
117+
*/
118+
private def tryAggressiveOptimize(e: Expr): Expr = e match {
119+
// Constant folding: BinaryOp with two constant operands (most common case first)
120+
case e @ BinaryOp(pos, lhs: Val, op, rhs: Val) => tryFoldBinaryOp(pos, lhs, op, rhs, e)
121+
122+
// Constant folding: UnaryOp with constant operand
123+
case e @ UnaryOp(pos, op, v: Val) => tryFoldUnaryOp(pos, op, v, e)
124+
125+
// Branch elimination: constant condition in if-else
126+
case IfElse(_, _: Val.True, thenExpr, _) => thenExpr
127+
case IfElse(pos, _: Val.False, _, elseExpr) =>
128+
if (elseExpr == null) Val.Null(pos) else elseExpr
129+
130+
// Short-circuit elimination for And/Or with constant lhs.
131+
//
132+
// IMPORTANT: rhs MUST be guarded as `Val.Bool` — do NOT relax this to arbitrary Expr.
133+
// The Evaluator's visitAnd/visitOr enforces that rhs evaluates to Bool, throwing
134+
// "binary operator && does not operate on <type>s" otherwise. If we fold `true && rhs`
135+
// into just `rhs` without the Bool guard, we silently remove that runtime type check,
136+
// causing programs like `true && "hello"` to return "hello" instead of erroring.
137+
// See: Evaluator.visitAnd / Evaluator.visitOr for the authoritative runtime semantics.
138+
case And(_, _: Val.True, rhs: Val.Bool) => rhs
139+
case And(pos, _: Val.False, _) => Val.False(pos)
140+
case Or(pos, _: Val.True, _) => Val.True(pos)
141+
case Or(_, _: Val.False, rhs: Val.Bool) => rhs
142+
143+
case _ => e
100144
}
101145

102146
private object ValidSuper {
@@ -258,4 +302,181 @@ class StaticOptimizer(
258302
}
259303
target
260304
}
305+
306+
private def tryFoldUnaryOp(pos: Position, op: Int, v: Val, fallback: Expr): Expr =
307+
try {
308+
op match {
309+
case Expr.UnaryOp.OP_! =>
310+
v match {
311+
case _: Val.True => Val.False(pos)
312+
case _: Val.False => Val.True(pos)
313+
case _ => fallback
314+
}
315+
case Expr.UnaryOp.OP_- =>
316+
v match {
317+
case Val.Num(_, n) => Val.Num(pos, -n)
318+
case _ => fallback
319+
}
320+
case Expr.UnaryOp.OP_~ =>
321+
v match {
322+
case Val.Num(_, n) => Val.Num(pos, (~n.toLong).toDouble)
323+
case _ => fallback
324+
}
325+
case Expr.UnaryOp.OP_+ =>
326+
v match {
327+
case Val.Num(_, n) => Val.Num(pos, n)
328+
case _ => fallback
329+
}
330+
case _ => fallback
331+
}
332+
} catch { case _: Exception => fallback }
333+
334+
private def tryFoldBinaryOp(pos: Position, lhs: Val, op: Int, rhs: Val, fallback: Expr): Expr =
335+
try {
336+
op match {
337+
case BinaryOp.OP_+ =>
338+
(lhs, rhs) match {
339+
case (Val.Num(_, l), Val.Num(_, r)) => Val.Num(pos, l + r)
340+
case (Val.Str(_, l), Val.Str(_, r)) => Val.Str(pos, l + r)
341+
case (l: Val.Arr, r: Val.Arr) => l.concat(pos, r)
342+
case _ => fallback
343+
}
344+
case BinaryOp.OP_- =>
345+
(lhs, rhs) match {
346+
case (Val.Num(_, l), Val.Num(_, r)) => Val.Num(pos, l - r)
347+
case _ => fallback
348+
}
349+
case BinaryOp.OP_* =>
350+
(lhs, rhs) match {
351+
case (Val.Num(_, l), Val.Num(_, r)) => Val.Num(pos, l * r)
352+
case _ => fallback
353+
}
354+
case BinaryOp.OP_/ =>
355+
(lhs, rhs) match {
356+
case (Val.Num(_, l), Val.Num(_, r)) if r != 0 => Val.Num(pos, l / r)
357+
case _ => fallback
358+
}
359+
case BinaryOp.OP_% =>
360+
(lhs, rhs) match {
361+
case (Val.Num(_, l), Val.Num(_, r)) => Val.Num(pos, l % r)
362+
case _ => fallback
363+
}
364+
case BinaryOp.OP_< =>
365+
tryFoldComparison(pos, lhs, BinaryOp.OP_<, rhs, fallback)
366+
case BinaryOp.OP_> =>
367+
tryFoldComparison(pos, lhs, BinaryOp.OP_>, rhs, fallback)
368+
case BinaryOp.OP_<= =>
369+
tryFoldComparison(pos, lhs, BinaryOp.OP_<=, rhs, fallback)
370+
case BinaryOp.OP_>= =>
371+
tryFoldComparison(pos, lhs, BinaryOp.OP_>=, rhs, fallback)
372+
case BinaryOp.OP_== =>
373+
tryFoldEquality(pos, lhs, rhs, negate = false, fallback)
374+
case BinaryOp.OP_!= =>
375+
tryFoldEquality(pos, lhs, rhs, negate = true, fallback)
376+
case BinaryOp.OP_in =>
377+
(lhs, rhs) match {
378+
case (Val.Str(_, l), o: Val.Obj) => Val.bool(pos, o.containsKey(l))
379+
case _ => fallback
380+
}
381+
case BinaryOp.OP_<< =>
382+
(lhs, rhs) match {
383+
case (Val.Num(_, l), Val.Num(_, r)) =>
384+
val ll = lhs.asInstanceOf[Val.Num].asSafeLong
385+
val rr = rhs.asInstanceOf[Val.Num].asSafeLong
386+
if (rr < 0) fallback // negative shift → runtime error
387+
else if (rr >= 1 && math.abs(ll) >= (1L << (63 - rr)))
388+
fallback // overflow → runtime error
389+
else Val.Num(pos, (ll << rr).toDouble)
390+
case _ => fallback
391+
}
392+
case BinaryOp.OP_>> =>
393+
(lhs, rhs) match {
394+
case (Val.Num(_, l), Val.Num(_, r)) =>
395+
val ll = lhs.asInstanceOf[Val.Num].asSafeLong
396+
val rr = rhs.asInstanceOf[Val.Num].asSafeLong
397+
if (rr < 0) fallback // negative shift → runtime error
398+
else Val.Num(pos, (ll >> rr).toDouble)
399+
case _ => fallback
400+
}
401+
case BinaryOp.OP_& =>
402+
(lhs, rhs) match {
403+
case (Val.Num(_, _), Val.Num(_, _)) =>
404+
Val.Num(
405+
pos,
406+
(lhs.asInstanceOf[Val.Num].asSafeLong & rhs
407+
.asInstanceOf[Val.Num]
408+
.asSafeLong).toDouble
409+
)
410+
case _ => fallback
411+
}
412+
case BinaryOp.OP_^ =>
413+
(lhs, rhs) match {
414+
case (Val.Num(_, _), Val.Num(_, _)) =>
415+
Val.Num(pos, (lhs.asLong ^ rhs.asLong).toDouble)
416+
case _ => fallback
417+
}
418+
case BinaryOp.OP_| =>
419+
(lhs, rhs) match {
420+
case (Val.Num(_, _), Val.Num(_, _)) =>
421+
Val.Num(pos, (lhs.asLong | rhs.asLong).toDouble)
422+
case _ => fallback
423+
}
424+
case _ => fallback
425+
}
426+
} catch { case _: Exception => fallback }
427+
428+
private def tryFoldComparison(
429+
pos: Position,
430+
lhs: Val,
431+
op: Int,
432+
rhs: Val,
433+
fallback: Expr): Expr = {
434+
// Use IEEE 754 operators directly for Num, not java.lang.Double.compare,
435+
// because compare(-0.0, 0.0) == -1 while IEEE 754 treats -0.0 == 0.0.
436+
(lhs, rhs) match {
437+
case (Val.Num(_, l), Val.Num(_, r)) if !l.isNaN && !r.isNaN =>
438+
val result = op match {
439+
case BinaryOp.OP_< => l < r
440+
case BinaryOp.OP_> => l > r
441+
case BinaryOp.OP_<= => l <= r
442+
case BinaryOp.OP_>= => l >= r
443+
case _ => return fallback
444+
}
445+
Val.bool(pos, result)
446+
case (Val.Str(_, l), Val.Str(_, r)) =>
447+
val cmp = Util.compareStringsByCodepoint(l, r)
448+
val result = op match {
449+
case BinaryOp.OP_< => cmp < 0
450+
case BinaryOp.OP_> => cmp > 0
451+
case BinaryOp.OP_<= => cmp <= 0
452+
case BinaryOp.OP_>= => cmp >= 0
453+
case _ => return fallback
454+
}
455+
Val.bool(pos, result)
456+
case _ => fallback
457+
}
458+
}
459+
460+
private def tryFoldEquality(
461+
pos: Position,
462+
lhs: Val,
463+
rhs: Val,
464+
negate: Boolean,
465+
fallback: Expr): Expr = {
466+
def isSimpleLiteral(v: Val): Boolean = v match {
467+
case _: Val.Bool | _: Val.Null | _: Val.Str | _: Val.Num => true
468+
case _ => false
469+
}
470+
if (!isSimpleLiteral(lhs) || !isSimpleLiteral(rhs)) return fallback
471+
val result = (lhs, rhs) match {
472+
case (_: Val.True, _: Val.True) | (_: Val.False, _: Val.False) | (_: Val.Null, _: Val.Null) =>
473+
true
474+
case (Val.Num(_, l), Val.Num(_, r)) if !l.isNaN && !r.isNaN =>
475+
l == r
476+
case (Val.Str(_, l), Val.Str(_, r)) =>
477+
l == r
478+
case _ => false // different simple types are never equal
479+
}
480+
Val.bool(pos, if (negate) !result else result)
481+
}
261482
}

0 commit comments

Comments
 (0)