perf: comprehension fuse scope+eval and inline BinaryOp(ValidId,ValidId) fast path

He-Pin · He-Pin · commit 9b5caef67c0a · 2026-04-05T12:00:42.000+08:00
Fuse comprehension scope building with body evaluation, eliminating redundant scope allocation in the innermost loop. When the body is BinaryOp(ValidId,ValidId), inline the scope lookups and binary-op dispatch entirely, avoiding 3× visitExpr overhead per iteration. Key changes: - ValScope.extendMutable(): creates a scope with one extra mutable slot for reuse across iterations (safe because results are eagerly evaluated, not captured in lazy thunks) - visitCompInline: split by rest (Nil vs non-Nil), with BinaryOp fast path for innermost loops - evalBinaryOpNumNum: @switch-dispatched Num×Num fast path covering all comparison, arithmetic, modulo, bitwise, and shift operators with full safety checks (overflow, division-by-zero, safe integer range) - visitBinaryOpValues: polymorphic fallback for non-Num operands covering string concat/format, object merge, array concat, equality, and 'in' Benchmark: comparison2 -53.1% (74.1 → 34.8 ms/op), zero regressions across 35 benchmarks. Upstream: jit branch commits 3466461 (fuse) + 71545ba (inline)
diff --git a/sjsonnet/src/sjsonnet/Evaluator.scala b/sjsonnet/src/sjsonnet/Evaluator.scala
@@ -189,11 +189,229 @@ class Evaluator(
     visitExpr(e.returned)(s)
   }
 
-  def visitComp(e: Comp)(implicit scope: ValScope): Val =
-    Val.Arr(
-      e.pos,
-      visitComp(e.first :: e.rest.toList, Array(scope)).map(s => visitAsLazy(e.value)(s))
-    )
+  def visitComp(e: Comp)(implicit scope: ValScope): Val = {
+    val body = e.value
+    val results = new collection.mutable.ArrayBuilder.ofRef[Eval]
+    results.sizeHint(256)
+    visitCompInline(e.first :: e.rest.toList, scope, body, results)
+    Val.Arr(e.pos, results.result())
+  }
+
+  /**
+   * Fused scope-building + body evaluation: eliminates intermediate scope array allocation. For
+   * innermost ForSpec with BinaryOp(ValidId,ValidId) body, inlines scope lookups and binary-op
+   * dispatch to avoid 3× visitExpr overhead per iteration.
+   */
+  private def visitCompInline(
+      specs: List[CompSpec],
+      scope: ValScope,
+      body: Expr,
+      results: collection.mutable.ArrayBuilder.ofRef[Eval]): Unit = specs match {
+    case (spec @ ForSpec(_, name, expr)) :: rest =>
+      visitExpr(expr)(scope) match {
+        case a: Val.Arr =>
+          if (debugStats != null) debugStats.arrayCompIterations += a.length
+          val lazyArr = a.asLazyArray
+          rest match {
+            case Nil =>
+              // Innermost loop — try BinaryOp(ValidId,ValidId) fast path
+              body match {
+                case binOp: BinaryOp
+                    if binOp.lhs.tag == ExprTags.ValidId
+                      && binOp.rhs.tag == ExprTags.ValidId =>
+                  // Fast path: reuse mutable scope, inline scope lookups + binary-op dispatch.
+                  // Avoids 3× visitExpr dispatch overhead per iteration.
+                  // NOTE: This evaluates the body eagerly (not lazy). This differs from the
+                  // normal path which uses visitAsLazy. The practical impact is minimal —
+                  // only observable if errors in comprehension elements are never accessed.
+                  // Both go-jsonnet and jrsonnet also evaluate comprehensions eagerly.
+                  val mutableScope = scope.extendMutable()
+                  val slot = scope.bindings.length
+                  val bindings = mutableScope.bindings
+                  val lhsIdx = binOp.lhs.asInstanceOf[ValidId].nameIdx
+                  val rhsIdx = binOp.rhs.asInstanceOf[ValidId].nameIdx
+                  val op = binOp.op
+                  val bpos = binOp.pos
+                  var j = 0
+                  while (j < lazyArr.length) {
+                    bindings(slot) = lazyArr(j)
+                    val l = bindings(lhsIdx).value
+                    val r = bindings(rhsIdx).value
+                    (l, r) match {
+                      case (ln: Val.Num, rn: Val.Num) =>
+                        results += evalBinaryOpNumNum(op, ln, rn, bpos)
+                      case _ =>
+                        results += visitBinaryOpValues(op, l, r, bpos)
+                    }
+                    j += 1
+                  }
+                case _ =>
+                  var j = 0
+                  while (j < lazyArr.length) {
+                    val newScope = scope.extendSimple(lazyArr(j))
+                    results += visitAsLazy(body)(newScope)
+                    j += 1
+                  }
+              }
+            case _ =>
+              var j = 0
+              while (j < lazyArr.length) {
+                val newScope = scope.extendSimple(lazyArr(j))
+                visitCompInline(rest, newScope, body, results)
+                j += 1
+              }
+          }
+        case r =>
+          Error.fail(
+            "In comprehension, can only iterate over array, not " + r.prettyName,
+            spec
+          )
+      }
+    case (spec @ IfSpec(offset, expr)) :: rest =>
+      visitExpr(expr)(scope) match {
+        case Val.True(_) =>
+          rest match {
+            case Nil => results += visitAsLazy(body)(scope)
+            case _   => visitCompInline(rest, scope, body, results)
+          }
+        case Val.False(_) => // filtered out
+        case other        =>
+          Error.fail(
+            "Condition must be boolean, got " + other.prettyName,
+            spec
+          )
+      }
+    case Nil =>
+      results += visitAsLazy(body)(scope)
+  }
+
+  /**
+   * Fast-path binary op evaluation for Num×Num operands. Handles comparison, arithmetic, modulo,
+   * bitwise, and shift operations without visitExpr dispatch overhead. Includes the same safety
+   * checks (overflow, division-by-zero) as visitBinaryOp.
+   */
+  @inline private def evalBinaryOpNumNum(op: Int, ln: Val.Num, rn: Val.Num, pos: Position): Val = {
+    val ld = ln.asDouble
+    val rd = rn.asDouble
+    (op: @switch) match {
+      case Expr.BinaryOp.OP_<  => Val.bool(pos, ld < rd)
+      case Expr.BinaryOp.OP_>  => Val.bool(pos, ld > rd)
+      case Expr.BinaryOp.OP_<= => Val.bool(pos, ld <= rd)
+      case Expr.BinaryOp.OP_>= => Val.bool(pos, ld >= rd)
+      case Expr.BinaryOp.OP_== => Val.bool(pos, ld == rd)
+      case Expr.BinaryOp.OP_!= => Val.bool(pos, ld != rd)
+      case Expr.BinaryOp.OP_+  =>
+        val r = ld + rd
+        if (r.isInfinite) Error.fail("overflow", pos)
+        Val.Num(pos, r)
+      case Expr.BinaryOp.OP_- =>
+        val r = ld - rd
+        if (r.isInfinite) Error.fail("overflow", pos)
+        Val.Num(pos, r)
+      case Expr.BinaryOp.OP_* =>
+        val r = ld * rd
+        if (r.isInfinite) Error.fail("overflow", pos)
+        Val.Num(pos, r)
+      case Expr.BinaryOp.OP_/ =>
+        if (rd == 0) Error.fail("division by zero", pos)
+        val r = ld / rd
+        if (r.isInfinite) Error.fail("overflow", pos)
+        Val.Num(pos, r)
+      case Expr.BinaryOp.OP_%  => Val.Num(pos, ld % rd)
+      case Expr.BinaryOp.OP_<< =>
+        val ll = ld.toSafeLong(pos); val rr = rd.toSafeLong(pos)
+        if (rr < 0) Error.fail("shift by negative exponent", pos)
+        if (rr >= 1 && math.abs(ll) >= (1L << (63 - rr)))
+          Error.fail("numeric value outside safe integer range for bitwise operation", pos)
+        Val.Num(pos, (ll << rr).toDouble)
+      case Expr.BinaryOp.OP_>> =>
+        val ll = ld.toSafeLong(pos); val rr = rd.toSafeLong(pos)
+        if (rr < 0) Error.fail("shift by negative exponent", pos)
+        Val.Num(pos, (ll >> rr).toDouble)
+      case Expr.BinaryOp.OP_& =>
+        Val.Num(pos, (ld.toSafeLong(pos) & rd.toSafeLong(pos)).toDouble)
+      case Expr.BinaryOp.OP_^ =>
+        Val.Num(pos, (ld.toSafeLong(pos) ^ rd.toSafeLong(pos)).toDouble)
+      case Expr.BinaryOp.OP_| =>
+        Val.Num(pos, (ld.toSafeLong(pos) | rd.toSafeLong(pos)).toDouble)
+      case _ => visitBinaryOpValues(op, ln, rn, pos)
+    }
+  }
+
+  /**
+   * Evaluate a binary op on two pre-evaluated values. Covers polymorphic operators (comparison,
+   * equality, string concat/format, object merge, array concat, in) as fallback when the Num×Num
+   * fast path doesn't apply.
+   */
+  private def visitBinaryOpValues(op: Int, l: Val, r: Val, pos: Position): Val = {
+    (op: @switch) match {
+      case Expr.BinaryOp.OP_< =>
+        (l, r) match {
+          case (ln: Val.Num, rn: Val.Num) => Val.bool(pos, ln.asDouble < rn.asDouble)
+          case (ls: Val.Str, rs: Val.Str) =>
+            Val.bool(pos, Util.compareStringsByCodepoint(ls.str, rs.str) < 0)
+          case (la: Val.Arr, ra: Val.Arr) => Val.bool(pos, compare(la, ra) < 0)
+          case _                          => failBinOp(l, op, r, pos)
+        }
+      case Expr.BinaryOp.OP_> =>
+        (l, r) match {
+          case (ln: Val.Num, rn: Val.Num) => Val.bool(pos, ln.asDouble > rn.asDouble)
+          case (ls: Val.Str, rs: Val.Str) =>
+            Val.bool(pos, Util.compareStringsByCodepoint(ls.str, rs.str) > 0)
+          case (la: Val.Arr, ra: Val.Arr) => Val.bool(pos, compare(la, ra) > 0)
+          case _                          => failBinOp(l, op, r, pos)
+        }
+      case Expr.BinaryOp.OP_<= =>
+        (l, r) match {
+          case (ln: Val.Num, rn: Val.Num) => Val.bool(pos, ln.asDouble <= rn.asDouble)
+          case (ls: Val.Str, rs: Val.Str) =>
+            Val.bool(pos, Util.compareStringsByCodepoint(ls.str, rs.str) <= 0)
+          case (la: Val.Arr, ra: Val.Arr) => Val.bool(pos, compare(la, ra) <= 0)
+          case _                          => failBinOp(l, op, r, pos)
+        }
+      case Expr.BinaryOp.OP_>= =>
+        (l, r) match {
+          case (ln: Val.Num, rn: Val.Num) => Val.bool(pos, ln.asDouble >= rn.asDouble)
+          case (ls: Val.Str, rs: Val.Str) =>
+            Val.bool(pos, Util.compareStringsByCodepoint(ls.str, rs.str) >= 0)
+          case (la: Val.Arr, ra: Val.Arr) => Val.bool(pos, compare(la, ra) >= 0)
+          case _                          => failBinOp(l, op, r, pos)
+        }
+      case Expr.BinaryOp.OP_== =>
+        if (l.isInstanceOf[Val.Func] && r.isInstanceOf[Val.Func])
+          Error.fail("cannot test equality of functions", pos)
+        Val.bool(pos, equal(l, r))
+      case Expr.BinaryOp.OP_!= =>
+        if (l.isInstanceOf[Val.Func] && r.isInstanceOf[Val.Func])
+          Error.fail("cannot test equality of functions", pos)
+        Val.bool(pos, !equal(l, r))
+      case Expr.BinaryOp.OP_+ =>
+        (l, r) match {
+          case (ln: Val.Num, rn: Val.Num) =>
+            val res = ln.asDouble + rn.asDouble
+            if (res.isInfinite) Error.fail("overflow", pos)
+            Val.Num(pos, res)
+          case (ls: Val.Str, rs: Val.Str) => Val.Str(pos, ls.str + rs.str)
+          case (ls: Val.Str, _)           => Val.Str(pos, ls.str + Materializer.stringify(r))
+          case (_, rs: Val.Str)           => Val.Str(pos, Materializer.stringify(l) + rs.str)
+          case (lo: Val.Obj, ro: Val.Obj) => ro.addSuper(pos, lo)
+          case (la: Val.Arr, ra: Val.Arr) => la.concat(pos, ra)
+          case _                          => failBinOp(l, op, r, pos)
+        }
+      case Expr.BinaryOp.OP_% =>
+        (l, r) match {
+          case (ln: Val.Num, rn: Val.Num) => Val.Num(pos, ln.asDouble % rn.asDouble)
+          case (ls: Val.Str, _)           => Val.Str(pos, Format.format(ls.str, r, pos))
+          case _                          => failBinOp(l, op, r, pos)
+        }
+      case Expr.BinaryOp.OP_in =>
+        (l, r) match {
+          case (ls: Val.Str, o: Val.Obj) => Val.bool(pos, o.containsKey(ls.str))
+          case _                         => failBinOp(l, op, r, pos)
+        }
+      case _ => failBinOp(l, op, r, pos)
+    }
+  }
 
   def visitArr(e: Arr)(implicit scope: ValScope): Val =
     Val.Arr(e.pos, e.value.map(visitAsLazy))
diff --git a/sjsonnet/src/sjsonnet/ValScope.scala b/sjsonnet/src/sjsonnet/ValScope.scala
@@ -59,6 +59,16 @@ final class ValScope private (val bindings: Array[Eval]) extends AnyVal {
     b(bindings.length + 2) = l3
     new ValScope(b)
   }
+
+  /**
+   * Extend scope with one mutable slot for in-place reuse in tight comprehension loops. The caller
+   * can mutate `bindings(parentLength)` without allocating a new scope per iteration. Only safe
+   * when the body does not create lazy values that capture the scope array.
+   */
+  def extendMutable(): ValScope = {
+    val b = util.Arrays.copyOf(bindings, bindings.length + 1)
+    new ValScope(b)
+  }
 }
 
 object ValScope {
diff --git a/sjsonnet/test/resources/new_test_suite/comprehension_binop_types.jsonnet b/sjsonnet/test/resources/new_test_suite/comprehension_binop_types.jsonnet
@@ -0,0 +1,55 @@
+// Regression test: all binary operators in comprehensions with ValidId operands
+local strs = ["hello", "world"];
+local nums = [1, 2, 3];
+local arrs = [[1, 2], [3, 4]];
+
+// String concatenation
+local str_concat = [a + b for a in strs for b in strs];
+
+// Numeric arithmetic
+local num_add = [a + b for a in nums for b in nums];
+local num_sub = [a - b for a in [10, 20] for b in [3, 5]];
+local num_mul = [a * b for a in [2, 3] for b in [4, 5]];
+local num_div = [a / b for a in [10, 20] for b in [2, 5]];
+local num_mod = [a % b for a in [10, 7] for b in [3, 4]];
+
+// Comparison operators
+local cmp_lt = [a < b for a in nums for b in nums];
+local cmp_eq = [a == b for a in nums for b in nums];
+local cmp_ne = [a != b for a in nums for b in nums];
+
+// Bitwise operators
+local bw_and = [a & b for a in [3, 5] for b in [6, 7]];
+local bw_or  = [a | b for a in [3, 5] for b in [6, 7]];
+local bw_xor = [a ^ b for a in [3, 5] for b in [6, 7]];
+local bw_shl = [a << b for a in [1, 2] for b in [1, 2]];
+local bw_shr = [a >> b for a in [8, 16] for b in [1, 2]];
+
+// String formatting
+local str_fmt = [a % b for a in ["val=%d", "x=%d"] for b in [42, 99]];
+
+// Array concatenation
+local arr_concat = [a + b for a in arrs for b in arrs];
+
+// 'in' operator
+local objs = [{a: 1}, {b: 2}];
+local in_test = [a in b for a in ["a", "b"] for b in objs];
+
+std.assertEqual(str_concat, ["hellohello", "helloworld", "worldhello", "worldworld"]) &&
+std.assertEqual(num_add, [2, 3, 4, 3, 4, 5, 4, 5, 6]) &&
+std.assertEqual(num_sub, [7, 5, 17, 15]) &&
+std.assertEqual(num_mul, [8, 10, 12, 15]) &&
+std.assertEqual(num_div, [5, 2, 10, 4]) &&
+std.assertEqual(num_mod, [1, 2, 1, 3]) &&
+std.assertEqual(cmp_lt, [false, true, true, false, false, true, false, false, false]) &&
+std.assertEqual(cmp_eq, [true, false, false, false, true, false, false, false, true]) &&
+std.assertEqual(cmp_ne, [false, true, true, true, false, true, true, true, false]) &&
+std.assertEqual(bw_and, [2, 3, 4, 5]) &&
+std.assertEqual(bw_or, [7, 7, 7, 7]) &&
+std.assertEqual(bw_xor, [5, 4, 3, 2]) &&
+std.assertEqual(bw_shl, [2, 4, 4, 8]) &&
+std.assertEqual(bw_shr, [4, 2, 8, 4]) &&
+std.assertEqual(str_fmt, ["val=42", "val=99", "x=42", "x=99"]) &&
+std.assertEqual(arr_concat, [[1, 2, 1, 2], [1, 2, 3, 4], [3, 4, 1, 2], [3, 4, 3, 4]]) &&
+std.assertEqual(in_test, [true, false, false, true]) &&
+true
diff --git a/sjsonnet/test/resources/new_test_suite/comprehension_binop_types.jsonnet.golden b/sjsonnet/test/resources/new_test_suite/comprehension_binop_types.jsonnet.golden
@@ -0,0 +1 @@
+true
