TranformerCPS mutable variable capture bug

[dvdvgt]

Consider the following program

def motz(n: Int): Int = {
  if (n <= 1) {
    1
  } else {
    val limit = n - 2
    def product(i: Int): Int = {
      motz(i) * motz(limit - i)
    }
    motz(n - 1) + sum(0, limit) { product }
  }
}

which is translated to

function motz_0(n_23, ks_1706, k_937) {
  motz_1: while (true) {
    if ((n_23 <= (1))) {
      return () => k_937(1, ks_1706);
    } else {
      function product_2(i_45, ks_1707, k_938) {
        return motz_0(i_45, ks_1707, (v_r_846, ks_1708) =>
          motz_0((((n_23 - (2))) - i_45), ks_1708, (v_r_847, ks_1709) =>
            () => k_938((v_r_846 * v_r_847), ks_1709)));
      }
      /* prepare call */
      const tmp_n_1 = n_23;
      const tmp_k_6 = k_937;
      k_937 = (v_r_848, ks_1710) =>
        sum_4(0, (tmp_n_1 - (2)), product_2, ks_1710, (v_r_849, ks_1711) =>
          () => tmp_k_6((v_r_848 + v_r_849), ks_1711));
      n_23 = (tmp_n_1 - (1));
      continue motz_1;
    }
  }
}

if not using optimizations (no-optimize). When the continuation k_937 is overwritten, we call sum_4 and correctly rename n to tmp_n_1. However, in product_2 we still use n_23 which is mutated before the continuation is called, altering the semantics.

Doing this correctly might be difficult since by the time we apply the substitution, the binding product_2 has already been introduced and is not considered anymore.

(this issue has been split off from #1200)

[jiribenes]

This looks suspicious to me:

effekt/effekt/shared/src/main/scala/effekt/generator/js/TransformerCps.scala

Lines 281 to 282 in bd89bbc

	// We need to create temporaries for all free variables that appear in arguments
	val freeInArgs = (vargs.flatMap(Variables.free) ++ bargs.flatMap(Variables.free)).toSet

When I try a bit harder, I get the following output, is it better @dvdvgt?
EDIT: no

function motz_0(n_0, ks_77, k_40) {
  motz_1: while (true) {
    if ((n_0 <= (1))) {
      return () => k_40(1, ks_77);
    } else {
      function product_0(i_1, ks_78, k_41) {
        return motz_0(i_1, ks_78, (v_r_44, ks_79) =>
          motz_0((((n_0 - (2))) - i_1), ks_79, (v_r_45, ks_80) =>
            () => k_41((v_r_44 * v_r_45), ks_80)));
      }
      /* prepare call */
      const tmp_n_0 = n_0;
      const tmp_k_1 = k_40;
      k_40 = (v_r_46, ks_81) =>
        sum_0(0, (tmp_n_0 - (2)), product_0, ks_81, (v_r_47, ks_82) =>
          () => tmp_k_1((v_r_46 + v_r_47), ks_82));
      n_0 = (tmp_n_0 - (1));
      continue motz_1;
    }
  }
}

Here's the patch

diff --git i/effekt/shared/src/main/scala/effekt/generator/js/TransformerCps.scala w/effekt/shared/src/main/scala/effekt/generator/js/TransformerCps.scala
index 4e07a0a3c..85c03c7b8 100644
--- i/effekt/shared/src/main/scala/effekt/generator/js/TransformerCps.scala
+++ w/effekt/shared/src/main/scala/effekt/generator/js/TransformerCps.scala
@@ -36,6 +36,7 @@ object TransformerCps extends Transformer {
     directStyle: Option[ContinuationInfo],
     // the current direct-style metacontinuation
     metacont: Option[Id],
+    blockFreeVars: Map[Id, Set[Id]],
     // the original declaration context (used to compile pattern matching)
     declarations: DeclarationContext,
     // the usual compiler context
@@ -63,6 +64,7 @@ object TransformerCps extends Transformer {
           None,
           None,
           None,
+          Map.empty,
           D, C)
 
         val name      = JSName(jsModuleName(module.path))
@@ -81,6 +83,7 @@ object TransformerCps extends Transformer {
           None,
           None,
           None,
+          Map.empty,
           D, C)
 
     input.definitions.map(toJS)
@@ -200,7 +203,9 @@ object TransformerCps extends Transformer {
 
     case cps.Stmt.LetDef(id, block, body) =>
       Binding { k =>
-        js.Const(nameDef(id), requiringThunk { toJS(id, block) }) :: toJS(body).run(k)
+        val blockFree = Variables.free(block)
+        val updatedContext = D.copy(blockFreeVars = D.blockFreeVars + (id -> blockFree))
+        js.Const(nameDef(id), requiringThunk { toJS(id, block) }) :: toJS(body)(using updatedContext).run(k)
       }
 
     case cps.Stmt.If(cond, thn, els) =>
@@ -278,8 +283,18 @@ object TransformerCps extends Transformer {
 
         used.jumped = true
 
+        val freeInVargs = vargs.flatMap(Variables.free).toSet
+        val freeInBargs = bargs.flatMap { b =>
+          val directFree = Variables.free(b)
+          val capturedVars = b match {
+            case cps.Block.BlockVar(blockId) => D.blockFreeVars.getOrElse(blockId, Set.empty)
+            case _ => Set.empty
+          }
+          directFree ++ capturedVars
+        }.toSet
+
         // We need to create temporaries for all free variables that appear in arguments
-        val freeInArgs = (vargs.flatMap(Variables.free) ++ bargs.flatMap(Variables.free)).toSet
+        val freeInArgs = freeInVargs ++ freeInBargs
         // Only compute free vars of continuation if it's not a ContVar
         val (isTailCall, freeInK) = k match {
           case Cont.ContVar(kid) => (kid == k1 && ks == ks1, Set.empty[Id])

[jiribenes]

I think you're right, we forget to substitute for blocks.
Here's the output of the patch below, see !!!
It's still not optimal since we never use product_0 to begin with...

We can also inline product_subst_0, see a comment under the patch.

function motz_0(n_0, ks_77, k_40) {
  motz_1: while (true) {
    if ((n_0 <= (1))) {
      return () => k_40(1, ks_77);
    } else {
      function product_0(i_1, ks_78, k_41) {
        return motz_0(i_1, ks_78, (v_r_44, ks_79) =>
          motz_0((((n_0 - (2))) - i_1), ks_79, (v_r_45, ks_80) =>
            () => k_41((v_r_44 * v_r_45), ks_80)));
      }
      /* prepare call */
      const tmp_n_0 = n_0;
      const tmp_k_1 = k_40;
      function product_subst_0(i_1, ks_78, k_41) { // !!!
        return motz_0(i_1, ks_78, (v_r_44, ks_79) =>
          motz_0((((tmp_n_0 - (2))) - i_1), ks_79, (v_r_45, ks_80) =>
            () => k_41((v_r_44 * v_r_45), ks_80)));
      }
      k_40 = (v_r_46, ks_81) =>
        sum_0(0, (tmp_n_0 - (2)), product_subst_0, ks_81, (v_r_47, ks_82) =>
          () => tmp_k_1((v_r_46 + v_r_47), ks_82));
      n_0 = (tmp_n_0 - (1));
      continue motz_1;
    }
  }
}

Patch here

diff --git i/effekt/shared/src/main/scala/effekt/generator/js/TransformerCps.scala w/effekt/shared/src/main/scala/effekt/generator/js/TransformerCps.scala
index 4e07a0a3c..252e81a74 100644
--- i/effekt/shared/src/main/scala/effekt/generator/js/TransformerCps.scala
+++ w/effekt/shared/src/main/scala/effekt/generator/js/TransformerCps.scala
@@ -36,6 +36,7 @@ object TransformerCps extends Transformer {
     directStyle: Option[ContinuationInfo],
     // the current direct-style metacontinuation
     metacont: Option[Id],
+    blockDefs: Map[Id, cps.Block],
     // the original declaration context (used to compile pattern matching)
     declarations: DeclarationContext,
     // the usual compiler context
@@ -63,6 +64,7 @@ object TransformerCps extends Transformer {
           None,
           None,
           None,
+          Map.empty,
           D, C)
 
         val name      = JSName(jsModuleName(module.path))
@@ -81,6 +83,7 @@ object TransformerCps extends Transformer {
           None,
           None,
           None,
+          Map.empty,
           D, C)
 
     input.definitions.map(toJS)
@@ -200,7 +203,8 @@ object TransformerCps extends Transformer {
 
     case cps.Stmt.LetDef(id, block, body) =>
       Binding { k =>
-        js.Const(nameDef(id), requiringThunk { toJS(id, block) }) :: toJS(body).run(k)
+        val updatedContext = D.copy(blockDefs = D.blockDefs + (id -> block))
+        js.Const(nameDef(id), requiringThunk { toJS(id, block) }) :: toJS(body)(using updatedContext).run(k)
       }
 
     case cps.Stmt.If(cond, thn, els) =>
@@ -320,10 +324,20 @@ object TransformerCps extends Transformer {
           } else None
         }
 
+        // Compute which blocks need to be substituted (those that capture overlapping variables)
+        val valueSubst = paramTmps.map { case (p, t) => p -> Expr.ValueVar(t) }
+        val blockSubst: Map[Id, cps.Block] = D.blockDefs.collect {
+          case (blockId, blockDef) if Variables.free(blockDef).intersect(overlapping).nonEmpty =>
+            val substitutedBlock = substitutions.substitute(blockDef)(using Substitution(values = valueSubst))
+            val newId = Id(s"${blockId}_subst")
+            stmts.append(js.Const(nameDef(newId), argumentToJS(substitutedBlock)))
+            blockId -> cps.Block.BlockVar(newId)
+        }
+
         // Prepare the substitution
         val subst = Substitution(
-          values = paramTmps.map { case (p, t) => p -> Expr.ValueVar(t) },
-          blocks = Map.empty,
+          values = valueSubst,
+          blocks = blockSubst,
           conts = tmp_k.map(t => k1 -> Cont.ContVar(t)).toMap,
           metaconts = tmp_ks.map(t => ks1 -> MetaCont(t)).toMap
         )

We could also inline the $blockId_subst directly with:

        val blockSubst: Map[Id, cps.Block] = D.blockDefs.collect {
          case (blockId, blockDef) if Variables.free(blockDef).intersect(overlapping).nonEmpty =>
            blockId -> substitutions.substitute(blockDef)(using Substitution(values = valueSubst))
        }