Fix mutable callback state propagation

Author: fffonion

PdVm.Runtime/PdVmProgramBase.cs

@@ -7,6 +7,8 @@ public abstract class PdVmProgramBase : IPdVmCallableProgram
     private readonly List<PdVmValue> _stack = new();
     private readonly List<PdVmValue> _locals = new();
     private readonly Dictionary<int, PdVmCaptureCell> _captureCells = new();
+    private readonly Dictionary<int, PdVmCaptureCell> _mutableBorrowAliases = new();
+    private readonly HashSet<PdVmCaptureCell> _mutableBorrowCells = new(ReferenceEqualityComparer.Instance);
     private readonly List<PdVmExecutionFrame> _executionFrames = new();
     private readonly SemaphoreSlim _managedExecutionGate = new(1, 1);
     private readonly object _callbackQueueLock = new();
@@ -24,6 +26,7 @@ public abstract class PdVmProgramBase : IPdVmCallableProgram
     private bool _shutdown;
     private bool _callbackRunnerActive;
     private CancellationTokenSource _callbackRunnerCancellation = new();
+    private (PdVmValue Value, PdVmCaptureCell Cell)? _lastBorrowedCapture;
 
     private sealed record QueuedCallbackWork(
         object Generation,
@@ -363,13 +366,49 @@ protected PdVmValue PopValue()
     protected PdVmValue LoadLocalValue(byte index)
     {
         var absolute = ResolveLocalIndex(index);
-        return _captureCells.TryGetValue(absolute, out var cell) ? cell.Value : _locals[absolute];
+        if (_captureCells.TryGetValue(absolute, out var cell))
+        {
+            _lastBorrowedCapture = IsInsideScriptCallable() && _mutableBorrowCells.Contains(cell)
+                ? (cell.Value, cell)
+                : null;
+            return cell.Value;
+        }
+
+        _lastBorrowedCapture = null;
+        return _locals[absolute];
     }
 
+    private bool IsInsideScriptCallable() =>
+        GetActiveFrameOrDefault() is { PrototypeId: not null };
+
     protected void StoreLocalValue(byte index, PdVmValue value)
     {
         ArgumentNullException.ThrowIfNull(value);
         var absolute = ResolveLocalIndex(index);
+        if (_mutableBorrowAliases.TryGetValue(absolute, out var aliasCell))
+        {
+            // The CLR compiler reuses the alias local as a temporary while
+            // lowering Set. Only the container returned by Set can publish a
+            // borrowed map/array update; staging scalars (usually null) must
+            // remain local temporaries.
+            if (value.Kind is not (PdVmValueKind.Map or PdVmValueKind.Array))
+            {
+                _locals[absolute] = value;
+                _lastBorrowedCapture = null;
+                return;
+            }
+
+            if (ReferencesCaptureCell(value, aliasCell, new HashSet<PdVmCaptureCell>(ReferenceEqualityComparer.Instance)))
+            {
+                throw new InvalidOperationException("callable capture ownership cycle is unsupported");
+            }
+
+            aliasCell.Value = value;
+            _locals[absolute] = value;
+            _lastBorrowedCapture = null;
+            return;
+        }
+
         if (_captureCells.TryGetValue(absolute, out var cell))
         {
             if (ReferencesCaptureCell(value, cell, new HashSet<PdVmCaptureCell>(ReferenceEqualityComparer.Instance)))
@@ -379,8 +418,15 @@ protected void StoreLocalValue(byte index, PdVmValue value)
 
             cell.Value = value;
         }
+        else if (_lastBorrowedCapture is { } borrowed &&
+                 ReferenceEquals(value, borrowed.Value))
+        {
+            _mutableBorrowAliases[absolute] = borrowed.Cell;
+            borrowed.Cell.Value = value;
+        }
 
         _locals[absolute] = value;
+        _lastBorrowedCapture = null;
     }
 
     protected PdVmValue[] GetLocalValues() => _locals.ToArray();
@@ -534,6 +580,12 @@ protected bool CompleteActiveFrame()
             {
                 _captureCells.Remove(absolute);
             }
+            foreach (var absolute in _mutableBorrowAliases.Keys
+                         .Where(index => index >= frame.LocalBase && index < frameEnd)
+                         .ToArray())
+            {
+                _mutableBorrowAliases.Remove(absolute);
+            }
 
             if (frameEnd != _locals.Count)
             {
@@ -570,6 +622,9 @@ protected void ResetRuntimeForReuse()
         _stack.Clear();
         _locals.Clear();
         _captureCells.Clear();
+        _mutableBorrowAliases.Clear();
+        _mutableBorrowCells.Clear();
+        _lastBorrowedCapture = null;
         _executionFrames.Clear();
         _managedCallableResult = null;
         _mapIterators.Clear();
@@ -584,6 +639,9 @@ protected void ShutdownRuntime()
         _stack.Clear();
         _locals.Clear();
         _captureCells.Clear();
+        _mutableBorrowAliases.Clear();
+        _mutableBorrowCells.Clear();
+        _lastBorrowedCapture = null;
         _executionFrames.Clear();
         _managedCallableResult = null;
         _mapIterators.Clear();
@@ -760,6 +818,10 @@ private void EnterScriptCallable(
                 if (prototype.SelfSlot != slot)
                 {
                     _captureCells[absolute] = cell;
+                    if (prototype.CaptureModes[index] == PdVmRuntimeCaptureBindingMode.BorrowMut)
+                    {
+                        _mutableBorrowCells.Add(cell);
+                    }
                 }
             }
         }
@@ -841,6 +903,11 @@ private PdVmValue BindCallable(IReadOnlyList<PdVmValue> args)
                     _captureCells[absolute] = shared;
                 }
 
+                if (mode == PdVmRuntimeCaptureBindingMode.BorrowMut)
+                {
+                    _mutableBorrowCells.Add(shared);
+                }
+
                 _locals[absolute] = shared.Value;
                 cells[index] = shared;
             }
@@ -866,6 +933,7 @@ private PdVmCallOutcome DetachLocal(IReadOnlyList<PdVmValue> args)
 
         var absolute = ResolveLocalIndex((byte)slot);
         _captureCells.Remove(absolute);
+        _mutableBorrowAliases.Remove(absolute);
         _locals[absolute] = PdVmValue.Null();
         return PdVmCallOutcome.Returned(PdVmCallReturn.None);
     }
@@ -1494,6 +1562,11 @@ private void AbortManagedInvocation(int stackBase, int localBase)
         {
             _captureCells.Remove(absolute);
         }
+        foreach (var absolute in _mutableBorrowAliases.Keys.Where(index => index >= localBase).ToArray())
+        {
+            _mutableBorrowAliases.Remove(absolute);
+        }
+        _lastBorrowedCapture = null;
 
         if (_locals.Count > localBase)
         {

PdVm.Tests/PdVmTypedDotNetInteropTests.cs

@@ -175,6 +175,35 @@ public async Task ManagedCallbackCanCallAClosureKeptInRootLocals()
         Assert.Equal(1, result.AsInt());
     }
 
+    [Fact]
+    public async Task MutableBorrowAliasPublishesContainerUpdates()
+    {
+        using var fixture = new SourceFixture(
+            "let mut state: map<int> = { value: 0 };\n" +
+            "pub fn increment() -> int {\n" +
+            "    let mut current = &mut state;\n" +
+            "    current.value = current.value + 1;\n" +
+            "    state.value\n" +
+            "}\n");
+        var output = PdVmDotNetSourceCompiler.CompileFile(fixture.SourcePath, fixture.OutputPath);
+        var program = Assert.IsAssignableFrom<IPdVmCallableProgram>(
+            PdVmAssemblyLoader.CreateProgram(Assembly.Load(File.ReadAllBytes(output))));
+        var host = PdVmDefaultHost.CreateConsoleHost();
+        _ = PdVmExecution.Run(program, host);
+
+        var first = await program.InvokeCallableAsync(
+            program.ResolveCallable("increment"),
+            Array.Empty<PdVmValue>(),
+            host);
+        var second = await program.InvokeCallableAsync(
+            program.ResolveCallable("increment"),
+            Array.Empty<PdVmValue>(),
+            host);
+
+        Assert.Equal(1, first.AsInt());
+        Assert.Equal(2, second.AsInt());
+    }
+
     [Fact]
     public async Task CallbackAdapterSchemaIsValidatedBeforeScriptExecution()
     {

examples/dotnet-minesweeper.rss

@@ -423,14 +423,19 @@ let mut display = new_board_array(64);
 let mut flash_cells = new_board_array(64);
 
 fn render_board(render_rows: int, render_columns: int, render_paused: bool, render_flash_until: int, render_pressed_index: int) -> null {
+    // These arrays are resized when the difficulty changes. Borrow the root
+    // slots so this renderer observes the replacement arrays instead of the
+    // 8x8 values captured during startup.
+    let current_display = &display;
+    let current_flash_cells = &flash_cells;
     let mut render_styles = new_board_array(render_rows * render_columns);
     let mut render_index = 0;
     while render_index < render_rows * render_columns {
-        let mut render_style = display[render_index].copy();
+        let mut render_style = current_display[render_index].copy();
         if render_paused {
             render_style = 0;
         }
-        if render_flash_until > 0 && flash_cells[render_index].copy() == 1 {
+        if render_flash_until > 0 && current_flash_cells[render_index].copy() == 1 {
             render_style = 2;
         }
         if render_index == render_pressed_index && render_style == 0 {