From 8c01cf7fb7943e5af17272888b57c44854b86bd0 Mon Sep 17 00:00:00 2001
From: Brian McKenna <brian@brianmckenna.org>
Date: Thu, 17 Sep 2026 09:30:38 +1000
Subject: [PATCH] Fix defer instance of Speedy.Control

This previously just eagerly evaluated the deferred computation. We
just need to have a constructor which suspends the deferred
computation and then unwraps it in the existing tail recursive loop.
---
 .../daml/lf/speedy/PrettyLightweight.scala    |  1 +
 .../digitalasset/daml/lf/speedy/Speedy.scala  |  6 ++-
 .../daml/lf/speedy/ContTTest.scala            | 41 +++++++++++++++++++
 3 files changed, 47 insertions(+), 1 deletion(-)
 create mode 100644 community/daml-lf/interpreter/src/test/scala/com/digitalasset/daml/lf/speedy/ContTTest.scala

diff --git a/community/daml-lf/interpreter/src/main/scala/com/digitalasset/daml/lf/speedy/PrettyLightweight.scala b/community/daml-lf/interpreter/src/main/scala/com/digitalasset/daml/lf/speedy/PrettyLightweight.scala
index f77afe02a76..08a8cd8229f 100644
--- a/community/daml-lf/interpreter/src/main/scala/com/digitalasset/daml/lf/speedy/PrettyLightweight.scala
+++ b/community/daml-lf/interpreter/src/main/scala/com/digitalasset/daml/lf/speedy/PrettyLightweight.scala
@@ -20,6 +20,7 @@ private[speedy] object PrettyLightweight { // lightweight pretty printer for CEK
       case Control.Question(_) => "question"
       case Control.Complete(_) => "complete"
       case Control.Error(_) => "error"
+      case Control.Suspend(_) => "suspend"
       case Control.WeAreUnset => "unset"
     }
 
diff --git a/community/daml-lf/interpreter/src/main/scala/com/digitalasset/daml/lf/speedy/Speedy.scala b/community/daml-lf/interpreter/src/main/scala/com/digitalasset/daml/lf/speedy/Speedy.scala
index 4b9a6519483..26d4db6057a 100644
--- a/community/daml-lf/interpreter/src/main/scala/com/digitalasset/daml/lf/speedy/Speedy.scala
+++ b/community/daml-lf/interpreter/src/main/scala/com/digitalasset/daml/lf/speedy/Speedy.scala
@@ -940,6 +940,9 @@ private[lf] object Speedy {
             interruptionCountDown -= 1
             metrics.incrCount[StepCount](StepCount.StepCtx)
             thisControl match {
+              case Control.Suspend(thunk) =>
+                control = thunk()
+                loop()
               case Control.Value(value) =>
                 popTempStackToBase()
                 control = popKont().execute(this, value)
@@ -1296,9 +1299,10 @@ private[lf] object Speedy {
     final case class Complete(res: SValue) extends Control[Nothing]
     final case class Error(err: interpretation.Error) extends Control[Nothing]
     final case object WeAreUnset extends Control[Nothing]
+    final case class Suspend[Q](thunk: () => Control[Q]) extends Control[Q]
 
     implicit object `Defer Control` extends cats.Defer[Control] {
-      override def defer[A](x: => Control[A]): Control[A] = x
+      override def defer[A](x: => Control[A]): Control[A] = Suspend(() => x)
     }
   }
 
diff --git a/community/daml-lf/interpreter/src/test/scala/com/digitalasset/daml/lf/speedy/ContTTest.scala b/community/daml-lf/interpreter/src/test/scala/com/digitalasset/daml/lf/speedy/ContTTest.scala
new file mode 100644
index 00000000000..b9bf173eaa4
--- /dev/null
+++ b/community/daml-lf/interpreter/src/test/scala/com/digitalasset/daml/lf/speedy/ContTTest.scala
@@ -0,0 +1,41 @@
+// Copyright (c) 2026 Digital Asset (Switzerland) GmbH and/or its affiliates. All rights reserved.
+// SPDX-License-Identifier: Apache-2.0
+
+package com.digitalasset.daml.lf
+package speedy
+
+import cats.data.ContT
+import com.digitalasset.canton.logging.SuppressingLogging
+import com.digitalasset.daml.lf.speedy.SResult.SResultFinal
+import org.scalatest.matchers.should.Matchers
+import org.scalatest.wordspec.AnyWordSpec
+
+import Speedy.Control
+
+class ContTTest
+    extends AnyWordSpec
+    with Matchers
+    with SuppressingLogging {
+
+  "ContT[Control, _, _]" should {
+
+    "drive a deep flatMap chain without overflowing stack" in {
+      val N = 100000
+      val chain: ContT[Control, Nothing, Int] =
+        (1 to N).foldLeft(ContT.pure[Control, Nothing, Int](0)) { (acc, _) =>
+          acc.flatMap(i => ContT.pure[Control, Nothing, Int](i + 1))
+        }
+
+      val control: Control[Nothing] = chain.run(n => Control.Value(SValue.SInt64(n.toLong)))
+
+      val machine = Speedy.Machine.fromPureSExpr(
+        PureCompiledPackages.Empty(Compiler.Config.Default),
+        SExpr.SEValue(SValue.SUnit),
+        MachineLogger(),
+      )
+      machine.setControl(control)
+
+      machine.run() shouldBe SResultFinal(SValue.SInt64(N.toLong))
+    }
+  }
+}