From 7469579addc0f19ea2bd2c89e6e7f6eb10671381 Mon Sep 17 00:00:00 2001
From: Jamie Willis <j.willis19@imperial.ac.uk>
Date: Sat, 4 Jan 2025 12:30:03 +0000
Subject: [PATCH] fix: make `CalleeSave`  (partially) re-entrant

---
 .../internal/deepembedding/backend/StrictParsley.scala   | 2 ++
 .../internal/machine/instructions/PrimitiveInstrs.scala  | 9 +++++----
 2 files changed, 7 insertions(+), 4 deletions(-)

diff --git a/parsley/shared/src/main/scala/parsley/internal/deepembedding/backend/StrictParsley.scala b/parsley/shared/src/main/scala/parsley/internal/deepembedding/backend/StrictParsley.scala
index 08957ac87..6eae2bd29 100644
--- a/parsley/shared/src/main/scala/parsley/internal/deepembedding/backend/StrictParsley.scala
+++ b/parsley/shared/src/main/scala/parsley/internal/deepembedding/backend/StrictParsley.scala
@@ -161,9 +161,11 @@ private [deepembedding] object StrictParsley {
         if (calleeSaveRequired && localRegs.nonEmpty) {
             val end = state.freshLabel()
             val calleeSave = state.freshLabel()
+            instrs += new instructions.Push(false)
             instrs += new instructions.Label(calleeSave)
             instrs += new instructions.CalleeSave(end, localRegs, reqRegs, allocatedRegs, numRegsUsedByParent)
             bodyGen |> {
+                instrs += new instructions.Push(true)
                 instrs += new instructions.Jump(calleeSave)
                 instrs += new instructions.Label(end)
             }
diff --git a/parsley/shared/src/main/scala/parsley/internal/machine/instructions/PrimitiveInstrs.scala b/parsley/shared/src/main/scala/parsley/internal/machine/instructions/PrimitiveInstrs.scala
index ef11d72c0..a958bb9e1 100644
--- a/parsley/shared/src/main/scala/parsley/internal/machine/instructions/PrimitiveInstrs.scala
+++ b/parsley/shared/src/main/scala/parsley/internal/machine/instructions/PrimitiveInstrs.scala
@@ -154,7 +154,10 @@ private [internal] object Span extends Instr {
     // $COVERAGE-ON$
 }
 
-// This instruction holds mutate state, but it is safe to do so, because it's always the first instruction of a DynCall.
+// FIXME: this isn't re-entrant, the stack allows us to tell, but the persisted
+// references are shared across nested visits. Is the reference deallocation idea also screwy?
+// perhaps deallocation only occurs when we have an empty restore stack?
+// This instruction holds mutable state, but it is safe to do so, because it's always the first instruction of a DynCall.
 private [parsley] final class CalleeSave(var label: Int, localRegs: Set[Ref[_]], reqSize: Int, slots: List[(Int, Int)], saveArray: Array[AnyRef])
     extends InstrWithLabel {
     private def this(label: Int, localRegs: Set[Ref[_]], reqSize: Int, slots: List[Int]) =
@@ -162,7 +165,6 @@ private [parsley] final class CalleeSave(var label: Int, localRegs: Set[Ref[_]],
     // this filters out the slots to ensure we only do callee-save on registers that might exist in the parent
     def this(label: Int, localRefs: Set[Ref[_]], reqSize: Int, slots: List[Int], numRegsInContext: Int) =
         this(label, localRefs, reqSize, slots.takeWhile(_ < numRegsInContext))
-    private var inUse = false
     private var oldRegs: Array[AnyRef] = null
 
     private def save(ctx: Context): Unit = {
@@ -199,17 +201,16 @@ private [parsley] final class CalleeSave(var label: Int, localRegs: Set[Ref[_]],
     }
 
     override def apply(ctx: Context): Unit = {
+        val inUse = !ctx.good || ctx.stack.pop[Boolean]()
         // Second-entry, callee-restore and either jump or fail
         if (inUse) {
             restore(ctx)
-            inUse = false
             continue(ctx)
         }
         // Entry for the first time, register as a handle, callee-save and inc
         else {
             ensureRegularInstruction(ctx)
             save(ctx)
-            inUse = true
             ctx.pushHandler(ctx.pc)
             ctx.inc()
         }