Automatic tail-call optimization (auto-TCO) in StaticOptimizer

sjsonnet/src/sjsonnet/Evaluator.scala

@@ -618,11 +618,18 @@ class Evaluator(
     checkStackDepth(e.pos, e)
     try {
       val lhs = visitExpr(e.value)
+      // Auto-TCO'd calls should normally be intercepted by visitExprWithTailCallSupport,
+      // but we handle them defensively here to preserve lazy semantics if this path is ever reached.
       implicit val tailstrictMode: TailstrictMode =
-        if (e.tailstrict) TailstrictModeEnabled else TailstrictModeDisabled
+        if (!e.strict && e.tailstrict) TailstrictModeAutoTCO
+        else if (e.tailstrict) TailstrictModeEnabled
+        else TailstrictModeDisabled
 
       if (e.tailstrict) {
-        TailCall.resolve(lhs.cast[Val.Func].apply(e.args.map(visitExpr(_)), e.namedNames, e.pos))
+        val args: Array[Eval] =
+          if (!e.strict) e.args.map(visitAsLazy(_))
+          else e.args.map(visitExpr(_)).asInstanceOf[Array[Eval]]
+        TailCall.resolve(lhs.cast[Val.Func].apply(args, e.namedNames, e.pos))
       } else {
         lhs.cast[Val.Func].apply(e.args.map(visitAsLazy(_)), e.namedNames, e.pos)
       }
@@ -635,7 +642,9 @@ class Evaluator(
     try {
       val lhs = visitExpr(e.value)
       implicit val tailstrictMode: TailstrictMode =
-        if (e.tailstrict) TailstrictModeEnabled else TailstrictModeDisabled
+        if (!e.strict && e.tailstrict) TailstrictModeAutoTCO
+        else if (e.tailstrict) TailstrictModeEnabled
+        else TailstrictModeDisabled
       if (e.tailstrict) {
         TailCall.resolve(lhs.cast[Val.Func].apply0(e.pos))
       } else {
@@ -650,9 +659,12 @@ class Evaluator(
     try {
       val lhs = visitExpr(e.value)
       implicit val tailstrictMode: TailstrictMode =
-        if (e.tailstrict) TailstrictModeEnabled else TailstrictModeDisabled
+        if (!e.strict && e.tailstrict) TailstrictModeAutoTCO
+        else if (e.tailstrict) TailstrictModeEnabled
+        else TailstrictModeDisabled
       if (e.tailstrict) {
-        TailCall.resolve(lhs.cast[Val.Func].apply1(visitExpr(e.a1), e.pos))
+        val arg: Eval = if (!e.strict) visitAsLazy(e.a1) else visitExpr(e.a1)
+        TailCall.resolve(lhs.cast[Val.Func].apply1(arg, e.pos))
       } else {
         val l1 = visitAsLazy(e.a1)
         lhs.cast[Val.Func].apply1(l1, e.pos)
@@ -666,10 +678,18 @@ class Evaluator(
     try {
       val lhs = visitExpr(e.value)
       implicit val tailstrictMode: TailstrictMode =
-        if (e.tailstrict) TailstrictModeEnabled else TailstrictModeDisabled
+        if (!e.strict && e.tailstrict) TailstrictModeAutoTCO
+        else if (e.tailstrict) TailstrictModeEnabled
+        else TailstrictModeDisabled
 
       if (e.tailstrict) {
-        TailCall.resolve(lhs.cast[Val.Func].apply2(visitExpr(e.a1), visitExpr(e.a2), e.pos))
+        if (!e.strict) {
+          TailCall.resolve(
+            lhs.cast[Val.Func].apply2(visitAsLazy(e.a1), visitAsLazy(e.a2), e.pos)
+          )
+        } else {
+          TailCall.resolve(lhs.cast[Val.Func].apply2(visitExpr(e.a1), visitExpr(e.a2), e.pos))
+        }
       } else {
         val l1 = visitAsLazy(e.a1)
         val l2 = visitAsLazy(e.a2)
@@ -684,12 +704,22 @@ class Evaluator(
     try {
       val lhs = visitExpr(e.value)
       implicit val tailstrictMode: TailstrictMode =
-        if (e.tailstrict) TailstrictModeEnabled else TailstrictModeDisabled
+        if (!e.strict && e.tailstrict) TailstrictModeAutoTCO
+        else if (e.tailstrict) TailstrictModeEnabled
+        else TailstrictModeDisabled
 
       if (e.tailstrict) {
-        TailCall.resolve(
-          lhs.cast[Val.Func].apply3(visitExpr(e.a1), visitExpr(e.a2), visitExpr(e.a3), e.pos)
-        )
+        if (!e.strict) {
+          TailCall.resolve(
+            lhs
+              .cast[Val.Func]
+              .apply3(visitAsLazy(e.a1), visitAsLazy(e.a2), visitAsLazy(e.a3), e.pos)
+          )
+        } else {
+          TailCall.resolve(
+            lhs.cast[Val.Func].apply3(visitExpr(e.a1), visitExpr(e.a2), visitExpr(e.a3), e.pos)
+          )
+        }
       } else {
         val l1 = visitAsLazy(e.a1)
         val l2 = visitAsLazy(e.a2)
@@ -1158,13 +1188,36 @@ class Evaluator(
       }
     }
 
+  // And/Or rhs tail-position helpers — extracted to preserve @tailrec on visitExprWithTailCallSupport.
+  // TailCall sentinels pass through without a boolean type check: this is a deliberate semantic
+  // relaxation matching google/jsonnet behavior (where `&&` is simply `if a then b else false`).
+  // Direct non-boolean rhs values (e.g. `true && "hello"`) are still caught.
+  private def visitAndRhsTailPos(rhs: Expr, andPos: Position)(implicit scope: ValScope): Val = {
+    visitExprWithTailCallSupport(rhs) match {
+      case b: Val.Bool  => b
+      case tc: TailCall => tc
+      case unknown      =>
+        Error.fail(s"binary operator && does not operate on ${unknown.prettyName}s.", andPos)
+    }
+  }
+
+  private def visitOrRhsTailPos(rhs: Expr, orPos: Position)(implicit scope: ValScope): Val = {
+    visitExprWithTailCallSupport(rhs) match {
+      case b: Val.Bool  => b
+      case tc: TailCall => tc
+      case unknown      =>
+        Error.fail(s"binary operator || does not operate on ${unknown.prettyName}s.", orPos)
+    }
+  }
+
   /**
    * Evaluate an expression with tail-call support. When a `tailstrict` call is encountered at a
    * potential tail position, returns a [[TailCall]] sentinel instead of recursing, enabling
    * `TailCall.resolve` in `visitApply*` to iterate rather than grow the JVM stack.
    *
    * Potential tail positions are propagated through: IfElse (both branches), LocalExpr (returned),
-   * and AssertExpr (returned). All other expression types delegate to normal `visitExpr`.
+   * AssertExpr (returned), And (rhs), Or (rhs), and Expr.Error (value). All other expression types
+   * delegate to normal `visitExpr`.
    */
   @tailrec
   private def visitExprWithTailCallSupport(e: Expr)(implicit scope: ValScope): Val = e match {
@@ -1208,6 +1261,26 @@ class Evaluator(
         }
       }
       visitExprWithTailCallSupport(e.returned)
+    case e: And =>
+      // rhs of && is in tail position: when lhs is true, rhs is returned directly.
+      // Type check via helper to preserve @tailrec on this method.
+      visitExpr(e.lhs) match {
+        case _: Val.True  => visitAndRhsTailPos(e.rhs, e.pos)
+        case _: Val.False => Val.staticFalse
+        case unknown      =>
+          Error.fail(s"binary operator && does not operate on ${unknown.prettyName}s.", e.pos)
+      }
+    case e: Or =>
+      // rhs of || is in tail position: when lhs is false, rhs is returned directly.
+      // Type check via helper to preserve @tailrec on this method.
+      visitExpr(e.lhs) match {
+        case _: Val.True  => Val.staticTrue
+        case _: Val.False => visitOrRhsTailPos(e.rhs, e.pos)
+        case unknown      =>
+          Error.fail(s"binary operator || does not operate on ${unknown.prettyName}s.", e.pos)
+      }
+    case e: Expr.Error =>
+      Error.fail(materializeError(visitExpr(e.value)), e.pos)
     // Tail-position tailstrict calls: match TailstrictableExpr to unify the tailstrict guard,
     // then dispatch by concrete type.
     //
@@ -1222,33 +1295,41 @@ class Evaluator(
       e match {
         case e: Apply =>
           try {
+            val isStrict = e.isStrict
             val func = visitExpr(e.value).cast[Val.Func]
-            new TailCall(func, e.args.map(visitExpr(_)).asInstanceOf[Array[Eval]], e.namedNames, e)
+            val args: Array[Eval] =
+              if (!isStrict) e.args.map(visitAsLazy(_))
+              else e.args.map(visitExpr(_)).asInstanceOf[Array[Eval]]
+            new TailCall(func, args, e.namedNames, e, strict = isStrict)
           } catch Error.withStackFrame(e)
         case e: Apply0 =>
           try {
             val func = visitExpr(e.value).cast[Val.Func]
-            new TailCall(func, Evaluator.emptyLazyArray, null, e)
+            new TailCall(func, Evaluator.emptyLazyArray, null, e, strict = e.isStrict)
           } catch Error.withStackFrame(e)
         case e: Apply1 =>
           try {
+            val isStrict = e.isStrict
             val func = visitExpr(e.value).cast[Val.Func]
-            new TailCall(func, Array[Eval](visitExpr(e.a1)), null, e)
+            val arg: Eval = if (!isStrict) visitAsLazy(e.a1) else visitExpr(e.a1)
+            new TailCall(func, Array[Eval](arg), null, e, strict = isStrict)
           } catch Error.withStackFrame(e)
         case e: Apply2 =>
           try {
+            val isStrict = e.isStrict
             val func = visitExpr(e.value).cast[Val.Func]
-            new TailCall(func, Array[Eval](visitExpr(e.a1), visitExpr(e.a2)), null, e)
+            val a1: Eval = if (!isStrict) visitAsLazy(e.a1) else visitExpr(e.a1)
+            val a2: Eval = if (!isStrict) visitAsLazy(e.a2) else visitExpr(e.a2)
+            new TailCall(func, Array[Eval](a1, a2), null, e, strict = isStrict)
           } catch Error.withStackFrame(e)
         case e: Apply3 =>
           try {
+            val isStrict = e.isStrict
             val func = visitExpr(e.value).cast[Val.Func]
-            new TailCall(
-              func,
-              Array[Eval](visitExpr(e.a1), visitExpr(e.a2), visitExpr(e.a3)),
-              null,
-              e
-            )
+            val a1: Eval = if (!isStrict) visitAsLazy(e.a1) else visitExpr(e.a1)
+            val a2: Eval = if (!isStrict) visitAsLazy(e.a2) else visitExpr(e.a2)
+            val a3: Eval = if (!isStrict) visitAsLazy(e.a3) else visitExpr(e.a3)
+            new TailCall(func, Array[Eval](a1, a2, a3), null, e, strict = isStrict)
           } catch Error.withStackFrame(e)
         case _ => visitExpr(e)
       }

sjsonnet/src/sjsonnet/Expr.scala

@@ -52,6 +52,13 @@ trait Expr {
  */
 trait TailstrictableExpr extends Expr {
   def tailstrict: Boolean
+
+  /**
+   * True when this call was marked as strict (eager argument evaluation) by an explicit
+   * `tailstrict` annotation. False when auto-TCO'd (lazy argument evaluation to preserve Jsonnet's
+   * standard lazy semantics).
+   */
+  def isStrict: Boolean = false
 }
 
 object Expr {
@@ -231,36 +238,58 @@ object Expr {
       value: Expr,
       args: Array[Expr],
       namedNames: Array[String],
-      tailstrict: Boolean)
+      tailstrict: Boolean,
+      strict: Boolean = true)
       extends TailstrictableExpr {
     final override private[sjsonnet] def tag = ExprTags.Apply
     override def exprErrorString: String = Expr.callTargetName(value)
+    override def isStrict: Boolean = strict
   }
-  final case class Apply0(var pos: Position, value: Expr, tailstrict: Boolean)
+  final case class Apply0(
+      var pos: Position,
+      value: Expr,
+      tailstrict: Boolean,
+      strict: Boolean = true)
       extends TailstrictableExpr {
     final override private[sjsonnet] def tag = ExprTags.Apply0
     override def exprErrorString: String = Expr.callTargetName(value)
+    override def isStrict: Boolean = strict
   }
-  final case class Apply1(var pos: Position, value: Expr, a1: Expr, tailstrict: Boolean)
+  final case class Apply1(
+      var pos: Position,
+      value: Expr,
+      a1: Expr,
+      tailstrict: Boolean,
+      strict: Boolean = true)
       extends TailstrictableExpr {
     final override private[sjsonnet] def tag = ExprTags.Apply1
     override def exprErrorString: String = Expr.callTargetName(value)
+    override def isStrict: Boolean = strict
   }
-  final case class Apply2(var pos: Position, value: Expr, a1: Expr, a2: Expr, tailstrict: Boolean)
+  final case class Apply2(
+      var pos: Position,
+      value: Expr,
+      a1: Expr,
+      a2: Expr,
+      tailstrict: Boolean,
+      strict: Boolean = true)
       extends TailstrictableExpr {
     final override private[sjsonnet] def tag = ExprTags.Apply2
     override def exprErrorString: String = Expr.callTargetName(value)
+    override def isStrict: Boolean = strict
   }
   final case class Apply3(
       var pos: Position,
       value: Expr,
       a1: Expr,
       a2: Expr,
       a3: Expr,
-      tailstrict: Boolean)
+      tailstrict: Boolean,
+      strict: Boolean = true)
       extends TailstrictableExpr {
     final override private[sjsonnet] def tag = ExprTags.Apply3
     override def exprErrorString: String = Expr.callTargetName(value)
+    override def isStrict: Boolean = strict
   }
   final case class ApplyBuiltin(
       var pos: Position,

sjsonnet/src/sjsonnet/ExprTransform.scala

@@ -14,37 +14,37 @@ abstract class ExprTransform {
         if (x2 eq x) expr
         else Select(pos, x2, name)
 
-      case Apply(pos, x, y, namedNames, tailstrict) =>
+      case Apply(pos, x, y, namedNames, tailstrict, strict) =>
         val x2 = transform(x)
         val y2 = transformArr(y)
         if ((x2 eq x) && (y2 eq y)) expr
-        else Apply(pos, x2, y2, namedNames, tailstrict)
+        else Apply(pos, x2, y2, namedNames, tailstrict, strict)
 
-      case Apply0(pos, x, tailstrict) =>
+      case Apply0(pos, x, tailstrict, strict) =>
         val x2 = transform(x)
         if (x2 eq x) expr
-        else Apply0(pos, x2, tailstrict)
+        else Apply0(pos, x2, tailstrict, strict)
 
-      case Apply1(pos, x, y, tailstrict) =>
+      case Apply1(pos, x, y, tailstrict, strict) =>
         val x2 = transform(x)
         val y2 = transform(y)
         if ((x2 eq x) && (y2 eq y)) expr
-        else Apply1(pos, x2, y2, tailstrict)
+        else Apply1(pos, x2, y2, tailstrict, strict)
 
-      case Apply2(pos, x, y, z, tailstrict) =>
+      case Apply2(pos, x, y, z, tailstrict, strict) =>
         val x2 = transform(x)
         val y2 = transform(y)
         val z2 = transform(z)
         if ((x2 eq x) && (y2 eq y) && (z2 eq z)) expr
-        else Apply2(pos, x2, y2, z2, tailstrict)
+        else Apply2(pos, x2, y2, z2, tailstrict, strict)
 
-      case Apply3(pos, x, y, z, a, tailstrict) =>
+      case Apply3(pos, x, y, z, a, tailstrict, strict) =>
         val x2 = transform(x)
         val y2 = transform(y)
         val z2 = transform(z)
         val a2 = transform(a)
         if ((x2 eq x) && (y2 eq y) && (z2 eq z) && (a2 eq a)) expr
-        else Apply3(pos, x2, y2, z2, a2, tailstrict)
+        else Apply3(pos, x2, y2, z2, a2, tailstrict, strict)
 
       case ApplyBuiltin(pos, func, x, tailstrict) =>
         val x2 = transformArr(x)