chore: TCO for tailstrict and Materializer.

Author: He-Pin

sjsonnet/src/sjsonnet/Error.scala

@@ -17,16 +17,20 @@ class Error(msg: String, stack: List[Error.Frame] = Nil, underlying: Option[Thro
   def addFrame(pos: Position, expr: Expr = null)(implicit ev: EvalErrorScope): Error = {
     if (stack.isEmpty || alwaysAddPos(expr)) {
       val exprErrorString = if (expr == null) null else expr.exprErrorString
-      val newFrame = new Error.Frame(pos, exprErrorString)
-      stack match {
-        case s :: ss if s.pos == pos =>
-          if (s.exprErrorString == null && exprErrorString != null) copy(stack = newFrame :: ss)
-          else this
-        case _ => copy(stack = newFrame :: stack)
-      }
+      addFrameString(pos, exprErrorString)
     } else this
   }
 
+  def addFrameString(pos: Position, exprErrorString: String)(implicit ev: EvalErrorScope): Error = {
+    val newFrame = new Error.Frame(pos, exprErrorString)
+    stack match {
+      case s :: ss if s.pos == pos =>
+        if (s.exprErrorString == null && exprErrorString != null) copy(stack = newFrame :: ss)
+        else this
+      case _ => copy(stack = newFrame :: stack)
+    }
+  }
+
   def asSeenFrom(ev: EvalErrorScope): Error =
     copy(stack = stack.map(_.asSeenFrom(ev)))
 

sjsonnet/src/sjsonnet/Evaluator.scala

@@ -202,13 +202,22 @@ class Evaluator(
     }
   }
 
+  /**
+   * Function application entry points (visitApply/visitApply0-3 for user functions,
+   * visitApplyBuiltin/visitApplyBuiltin0-4 for built-in functions).
+   *
+   * When `e.tailstrict` is true, the result is wrapped in `TailCall.resolve()` which iteratively
+   * resolves any [[TailCall]] chain. When false, arguments are wrapped as lazy thunks to preserve
+   * Jsonnet's default lazy evaluation semantics, and `Val.Func.apply` resolves any TailCall
+   * internally via `TailCall.resolve` before returning.
+   */
   protected def visitApply(e: Apply)(implicit scope: ValScope): Val = {
     val lhs = visitExpr(e.value)
     implicit val tailstrictMode: TailstrictMode =
       if (e.tailstrict) TailstrictModeEnabled else TailstrictModeDisabled
 
     if (e.tailstrict) {
-      lhs.cast[Val.Func].apply(e.args.map(visitExpr(_)), e.namedNames, e.pos)
+      TailCall.resolve(lhs.cast[Val.Func].apply(e.args.map(visitExpr(_)), e.namedNames, e.pos))
     } else {
       lhs.cast[Val.Func].apply(e.args.map(visitAsLazy(_)), e.namedNames, e.pos)
     }
@@ -218,15 +227,19 @@ class Evaluator(
     val lhs = visitExpr(e.value)
     implicit val tailstrictMode: TailstrictMode =
       if (e.tailstrict) TailstrictModeEnabled else TailstrictModeDisabled
-    lhs.cast[Val.Func].apply0(e.pos)
+    if (e.tailstrict) {
+      TailCall.resolve(lhs.cast[Val.Func].apply0(e.pos))
+    } else {
+      lhs.cast[Val.Func].apply0(e.pos)
+    }
   }
 
   protected def visitApply1(e: Apply1)(implicit scope: ValScope): Val = {
     val lhs = visitExpr(e.value)
     implicit val tailstrictMode: TailstrictMode =
       if (e.tailstrict) TailstrictModeEnabled else TailstrictModeDisabled
     if (e.tailstrict) {
-      lhs.cast[Val.Func].apply1(visitExpr(e.a1), e.pos)
+      TailCall.resolve(lhs.cast[Val.Func].apply1(visitExpr(e.a1), e.pos))
     } else {
       val l1 = visitAsLazy(e.a1)
       lhs.cast[Val.Func].apply1(l1, e.pos)
@@ -239,7 +252,7 @@ class Evaluator(
       if (e.tailstrict) TailstrictModeEnabled else TailstrictModeDisabled
 
     if (e.tailstrict) {
-      lhs.cast[Val.Func].apply2(visitExpr(e.a1), visitExpr(e.a2), e.pos)
+      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)
@@ -253,7 +266,9 @@ class Evaluator(
       if (e.tailstrict) TailstrictModeEnabled else TailstrictModeDisabled
 
     if (e.tailstrict) {
-      lhs.cast[Val.Func].apply3(visitExpr(e.a1), visitExpr(e.a2), visitExpr(e.a3), e.pos)
+      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)
@@ -262,41 +277,48 @@ class Evaluator(
     }
   }
 
-  protected def visitApplyBuiltin0(e: ApplyBuiltin0): Val = e.func.evalRhs(this, e.pos)
+  protected def visitApplyBuiltin0(e: ApplyBuiltin0): Val = {
+    val result = e.func.evalRhs(this, e.pos)
+    if (e.tailstrict) TailCall.resolve(result) else result
+  }
 
   protected def visitApplyBuiltin1(e: ApplyBuiltin1)(implicit scope: ValScope): Val = {
     if (e.tailstrict) {
-      e.func.evalRhs(visitExpr(e.a1), this, e.pos)
+      TailCall.resolve(e.func.evalRhs(visitExpr(e.a1), this, e.pos))
     } else {
       e.func.evalRhs(visitAsLazy(e.a1), this, e.pos)
     }
   }
 
   protected def visitApplyBuiltin2(e: ApplyBuiltin2)(implicit scope: ValScope): Val = {
     if (e.tailstrict) {
-      e.func.evalRhs(visitExpr(e.a1), visitExpr(e.a2), this, e.pos)
+      TailCall.resolve(e.func.evalRhs(visitExpr(e.a1), visitExpr(e.a2), this, e.pos))
     } else {
       e.func.evalRhs(visitAsLazy(e.a1), visitAsLazy(e.a2), this, e.pos)
     }
   }
 
   protected def visitApplyBuiltin3(e: ApplyBuiltin3)(implicit scope: ValScope): Val = {
     if (e.tailstrict) {
-      e.func.evalRhs(visitExpr(e.a1), visitExpr(e.a2), visitExpr(e.a3), this, e.pos)
+      TailCall.resolve(
+        e.func.evalRhs(visitExpr(e.a1), visitExpr(e.a2), visitExpr(e.a3), this, e.pos)
+      )
     } else {
       e.func.evalRhs(visitAsLazy(e.a1), visitAsLazy(e.a2), visitAsLazy(e.a3), this, e.pos)
     }
   }
 
   protected def visitApplyBuiltin4(e: ApplyBuiltin4)(implicit scope: ValScope): Val = {
     if (e.tailstrict) {
-      e.func.evalRhs(
-        visitExpr(e.a1),
-        visitExpr(e.a2),
-        visitExpr(e.a3),
-        visitExpr(e.a4),
-        this,
-        e.pos
+      TailCall.resolve(
+        e.func.evalRhs(
+          visitExpr(e.a1),
+          visitExpr(e.a2),
+          visitExpr(e.a3),
+          visitExpr(e.a4),
+          this,
+          e.pos
+        )
       )
     } else {
       e.func.evalRhs(
@@ -319,7 +341,7 @@ class Evaluator(
         arr(idx) = visitExpr(e.argExprs(idx))
         idx += 1
       }
-      e.func.evalRhs(arr, this, e.pos)
+      TailCall.resolve(e.func.evalRhs(arr, this, e.pos))
     } else {
       while (idx < e.argExprs.length) {
         val boundIdx = idx
@@ -638,10 +660,107 @@ class Evaluator(
       scope: ValScope): Val.Func =
     new Val.Func(outerPos, scope, params) {
       def evalRhs(vs: ValScope, es: EvalScope, fs: FileScope, pos: Position): Val =
-        visitExpr(rhs)(vs)
+        visitExprWithTailCallSupport(rhs)(vs)
       override def evalDefault(expr: Expr, vs: ValScope, es: EvalScope): Val = visitExpr(expr)(vs)
     }
 
+  /**
+   * 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`.
+   */
+  @tailrec
+  private def visitExprWithTailCallSupport(e: Expr)(implicit scope: ValScope): Val = e match {
+    case e: IfElse =>
+      visitExpr(e.cond) match {
+        case Val.True(_)  => visitExprWithTailCallSupport(e.`then`)
+        case Val.False(_) =>
+          e.`else` match {
+            case null => Val.Null(e.pos)
+            case v    => visitExprWithTailCallSupport(v)
+          }
+        case v => Error.fail("Need boolean, found " + v.prettyName, e.pos)
+      }
+    case e: LocalExpr =>
+      val bindings = e.bindings
+      val s =
+        if (bindings == null) scope
+        else {
+          val base = scope.length
+          val newScope = scope.extendBy(bindings.length)
+          var i = 0
+          while (i < bindings.length) {
+            val b = bindings(i)
+            newScope.bindings(base + i) = b.args match {
+              case null    => visitAsLazy(b.rhs)(newScope)
+              case argSpec =>
+                new Lazy(() => visitMethod(b.rhs, argSpec, b.pos)(newScope))
+            }
+            i += 1
+          }
+          newScope
+        }
+      visitExprWithTailCallSupport(e.returned)(s)
+    case e: AssertExpr =>
+      if (!visitExpr(e.asserted.value).isInstanceOf[Val.True]) {
+        e.asserted.msg match {
+          case null => Error.fail("Assertion failed", e)
+          case msg  =>
+            Error.fail("Assertion failed: " + materializeError(visitExpr(msg)), e)
+        }
+      }
+      visitExprWithTailCallSupport(e.returned)
+    // Tail-position tailstrict calls: match TailstrictableExpr to unify the tailstrict guard,
+    // then dispatch by concrete type.
+    //
+    // - Apply* (user function calls): construct a TailCall sentinel that the caller's
+    //   TailCall.resolve loop will resolve iteratively, avoiding JVM stack growth for
+    //   tail-recursive calls.
+    // - ApplyBuiltin* (built-in function calls): fall through to visitExpr, which dispatches to
+    //   visitApplyBuiltin*. Those methods already wrap their result in TailCall.resolve() when
+    //   tailstrict=true, resolving any TailCall that a user-defined callback (e.g. the function
+    //   argument to std.makeArray or std.sort) may have returned.
+    case e: TailstrictableExpr if e.tailstrict =>
+      e match {
+        case e: Apply =>
+          try {
+            val func = visitExpr(e.value).cast[Val.Func]
+            new TailCall(func, e.args.map(visitExpr(_)).asInstanceOf[Array[Eval]], e.namedNames, e)
+          } catch Error.withStackFrame(e)
+        case e: Apply0 =>
+          try {
+            val func = visitExpr(e.value).cast[Val.Func]
+            new TailCall(func, Evaluator.emptyLazyArray, null, e)
+          } catch Error.withStackFrame(e)
+        case e: Apply1 =>
+          try {
+            val func = visitExpr(e.value).cast[Val.Func]
+            new TailCall(func, Array[Eval](visitExpr(e.a1)), null, e)
+          } catch Error.withStackFrame(e)
+        case e: Apply2 =>
+          try {
+            val func = visitExpr(e.value).cast[Val.Func]
+            new TailCall(func, Array[Eval](visitExpr(e.a1), visitExpr(e.a2)), null, e)
+          } catch Error.withStackFrame(e)
+        case e: Apply3 =>
+          try {
+            val func = visitExpr(e.value).cast[Val.Func]
+            new TailCall(
+              func,
+              Array[Eval](visitExpr(e.a1), visitExpr(e.a2), visitExpr(e.a3)),
+              null,
+              e
+            )
+          } catch Error.withStackFrame(e)
+        case _ => visitExpr(e)
+      }
+    case _ =>
+      visitExpr(e)
+  }
+
   def visitBindings(bindings: Array[Bind], scope: => ValScope): Array[Eval] = {
     val arrF = new Array[Eval](bindings.length)
     var i = 0

sjsonnet/src/sjsonnet/Expr.scala

@@ -23,6 +23,27 @@ trait Expr {
 
   override def toString: String = s"$exprErrorString@$pos"
 }
+
+/**
+ * Marker trait for [[Expr]] nodes that represent function calls eligible for tail-call
+ * optimization. All Apply* (user function calls) and ApplyBuiltin* (built-in function calls) mix in
+ * this trait, providing a uniform `tailstrict` flag. The evaluator handles the two families
+ * differently when `tailstrict` is true:
+ *
+ *   - '''User function calls''' (Apply*) in tail position: the evaluator constructs a [[TailCall]]
+ *     sentinel and returns it to the caller's [[TailCall.resolve]] trampoline loop, avoiding JVM
+ *     stack growth for tail-recursive calls.
+ *   - '''Built-in function calls''' (ApplyBuiltin*): the evaluator wraps the result in
+ *     [[TailCall.resolve]] at the call site, resolving any [[TailCall]] that a user-defined
+ *     callback (e.g. the function argument to `std.makeArray` or `std.sort`) may have returned.
+ *
+ * @see
+ *   [[TailCall]] for the sentinel value used in the TCO protocol
+ */
+trait TailstrictableExpr extends Expr {
+  def tailstrict: Boolean
+}
+
 object Expr {
   private final def arrStr(a: Array[?]): String = {
     if (a == null) "null" else a.mkString("[", ", ", "]")
@@ -189,17 +210,19 @@ object Expr {
       args: Array[Expr],
       namedNames: Array[String],
       tailstrict: Boolean)
-      extends Expr {
+      extends TailstrictableExpr {
     final override private[sjsonnet] def tag = ExprTags.Apply
   }
-  final case class Apply0(pos: Position, value: Expr, tailstrict: Boolean) extends Expr {
+  final case class Apply0(pos: Position, value: Expr, tailstrict: Boolean)
+      extends TailstrictableExpr {
     final override private[sjsonnet] def tag = ExprTags.Apply0
   }
-  final case class Apply1(pos: Position, value: Expr, a1: Expr, tailstrict: Boolean) extends Expr {
+  final case class Apply1(pos: Position, value: Expr, a1: Expr, tailstrict: Boolean)
+      extends TailstrictableExpr {
     final override private[sjsonnet] def tag = ExprTags.Apply1
   }
   final case class Apply2(pos: Position, value: Expr, a1: Expr, a2: Expr, tailstrict: Boolean)
-      extends Expr {
+      extends TailstrictableExpr {
     final override private[sjsonnet] def tag = ExprTags.Apply2
   }
   final case class Apply3(
@@ -209,25 +232,25 @@ object Expr {
       a2: Expr,
       a3: Expr,
       tailstrict: Boolean)
-      extends Expr {
+      extends TailstrictableExpr {
     final override private[sjsonnet] def tag = ExprTags.Apply3
   }
   final case class ApplyBuiltin(
       pos: Position,
       func: Val.Builtin,
       argExprs: Array[Expr],
       tailstrict: Boolean)
-      extends Expr {
+      extends TailstrictableExpr {
     final override private[sjsonnet] def tag = ExprTags.ApplyBuiltin
     override def exprErrorString: String = s"std.${func.functionName}"
   }
   final case class ApplyBuiltin0(pos: Position, func: Val.Builtin0, tailstrict: Boolean)
-      extends Expr {
+      extends TailstrictableExpr {
     final override private[sjsonnet] def tag = ExprTags.ApplyBuiltin0
     override def exprErrorString: String = s"std.${func.functionName}"
   }
   final case class ApplyBuiltin1(pos: Position, func: Val.Builtin1, a1: Expr, tailstrict: Boolean)
-      extends Expr {
+      extends TailstrictableExpr {
     final override private[sjsonnet] def tag = ExprTags.ApplyBuiltin1
     override def exprErrorString: String = s"std.${func.functionName}"
   }
@@ -237,7 +260,7 @@ object Expr {
       a1: Expr,
       a2: Expr,
       tailstrict: Boolean)
-      extends Expr {
+      extends TailstrictableExpr {
     final override private[sjsonnet] def tag = ExprTags.ApplyBuiltin2
     override def exprErrorString: String = s"std.${func.functionName}"
   }
@@ -248,7 +271,7 @@ object Expr {
       a2: Expr,
       a3: Expr,
       tailstrict: Boolean)
-      extends Expr {
+      extends TailstrictableExpr {
     final override private[sjsonnet] def tag = ExprTags.ApplyBuiltin3
     override def exprErrorString: String = s"std.${func.functionName}"
   }
@@ -260,7 +283,7 @@ object Expr {
       a3: Expr,
       a4: Expr,
       tailstrict: Boolean)
-      extends Expr {
+      extends TailstrictableExpr {
     override private[sjsonnet] def tag = ExprTags.ApplyBuiltin4
     override def exprErrorString: String = s"std.${func.functionName}"
   }