Add explicit successor edges to finally handler exits

This fixes the regalloc problem described in #33.

One downside of this approach is that it will require all ResumeInsts
to be in sync with region exit targets, but there are currently no
passes messing with those.

src/DistIL/Analysis/ProtectedRegionAnalysis.cs

@@ -33,6 +33,9 @@ public ProtectedRegionAnalysis(MethodBody method)
             }
             region.Blocks.Add(block);
 
+            // Don't need to visit resume edges because they only exist only for consistency.
+            if (block.Last is ResumeInst) continue;
+
             // Add succs to worklist
             var succRegion = block.Last is LeaveInst ? region.Parent! : region;
             foreach (var succ in block.Succs) {

src/DistIL/Frontend/BlockState.cs

@@ -4,7 +4,7 @@ internal class BlockState
 {
     readonly ILImporter _importer;
     readonly MethodBody _body;
-    readonly int _startOffset;
+    public readonly int StartOffset;
 
     private ModuleDef _mod => _body.Definition.Module;
 
@@ -24,7 +24,7 @@ public BlockState(ILImporter importer, int startOffset)
     {
         _importer = importer;
         _body = importer._body;
-        _startOffset = startOffset;
+        this.StartOffset = startOffset;
 
         Block = _body.CreateBlock();
         EntryBlock = Block;
@@ -800,25 +800,36 @@ private void ImportRet()
 
     private void ImportLeave(int targetOffset)
     {
-        var chainBlock = _importer.GetBlock(targetOffset).Block;
-        var currRegion = _importer._regionTree!.FindEnclosing(_startOffset).Parent!;
-        var parentRegion = _importer._regionTree!.FindEnclosing(targetOffset);
+        var currRegion = _importer._regionTree!.FindEnclosing(StartOffset);
+        var destRegion = _importer._regionTree!.FindEnclosing(targetOffset);
+        var targetBlock = _importer.GetBlock(targetOffset).Block;
+
+        currRegion.ExitTargets.Add(targetBlock);
 
         // Create a chain of blocks leaving all nested regions until target (in reverse order)
-        while (currRegion != parentRegion) {
-            // TODO: avoid creating new blocks (can't use preds as they may be in another region)
+        while (currRegion.Parent! != destRegion) {
+            currRegion = currRegion.Parent!;
+
+            currRegion.LeavingBlocks ??= new();
+
+            if (currRegion.LeavingBlocks.TryGetValue(targetBlock, out var exitBlock)) {
+                targetBlock = exitBlock;
+                break;
+            }
             var nextBlock = _body.CreateBlock(insertAfter: Block);
-            nextBlock.InsertLast(new LeaveInst(chainBlock!));
+            nextBlock.InsertLast(new LeaveInst(targetBlock!));
 
-            chainBlock = nextBlock;
-            currRegion = currRegion.Parent!;
+            currRegion.LeavingBlocks.Add(targetBlock, nextBlock);
+            targetBlock = nextBlock;
         }
-        TerminateBlock(new LeaveInst(chainBlock), clearStack: true);
+        TerminateBlock(new LeaveInst(targetBlock), clearStack: true);
     }
     private void ImportResume(bool isFromFilter)
     {
         var filterResult = isFromFilter ? Pop() : null;
-        TerminateBlock(new ResumeInst(filterResult), clearStack: true);
+        TerminateBlock(new ResumeInst([], filterResult), clearStack: true);
+
+        _importer._blocksEndingWithEhResume.Add(this);
     }
 
     private void ImportThrow(bool isRethrow)

src/DistIL/Frontend/ILImporter.cs

@@ -13,6 +13,7 @@ public class ILImporter
     internal readonly Value?[] _varSlots;
 
     readonly Dictionary<int, BlockState> _blocks = new();
+    internal List<BlockState> _blocksEndingWithEhResume = new();
 
     private ILImporter(MethodDef method)
     {
@@ -44,6 +45,9 @@ private MethodBody ImportCode()
         CreateBlocks(leaders);
         CreateGuards(ehRegions);
         ImportBlocks(code, leaders);
+
+        AddFinallyResumeEdges(ehRegions);
+
         return _body;
     }
 
@@ -154,6 +158,21 @@ private void ImportBlocks(Span<ILInstruction> code, BitSet leaders)
         }
     }
 
+    private void AddFinallyResumeEdges(ExceptionClause[] clauses)
+    {
+        foreach (var block in _blocksEndingWithEhResume) {
+            var parentClause = clauses.First(c => {
+                return c.Kind == ExceptionRegionKind.Filter 
+                    ? block.StartOffset >= c.FilterStart && block.StartOffset < c.FilterEnd 
+                    : block.StartOffset >= c.HandlerStart && block.StartOffset < c.HandlerEnd;
+            });
+            var protectedNode = _regionTree!.FindEnclosing(parentClause.TryStart);
+
+            var resume = (ResumeInst)block.Block.Last;
+            resume.SetExitTargets(protectedNode.ExitTargets);
+        }
+    }
+
     private void AnalyseVars(Span<ILInstruction> code, BitSet leaders)
     {
         int blockStartIdx = 0;

src/DistIL/Frontend/RegionNode.cs

@@ -10,6 +10,9 @@ internal class RegionNode
     public RegionNode? Parent;
     public List<RegionNode> Children = new();
 
+    public HashSet<BasicBlock> ExitTargets = new();
+    public Dictionary<BasicBlock, BasicBlock>? LeavingBlocks; // Target -> Leave. May form a chain through multiple regions at a time.
+
     public RegionNode FindEnclosing(int start, int end)
     {
         foreach (var child in Children) {

src/DistIL/IR/BasicBlock.cs

@@ -317,7 +317,7 @@ public IEnumerable<Instruction> NonPhis()
     }
 
     private static bool IsBranchWithSuccEdges(Instruction? inst)
-        => inst is BranchInst or SwitchInst or LeaveInst;
+        => inst is BranchInst or SwitchInst or LeaveInst or ResumeInst;
 
     // Enumerating block users (ignoring phis) will lead directly to predecessors.
     // GuardInst`s will not yield duplicates because handler blocks can only have one predecessor guard.
@@ -362,8 +362,9 @@ internal SuccIterator(BasicBlock block)
                 //  Switch: [value, targetBlock0, targetBlock1, ...]  (targets are never duplicated)
                 //  Guard:  [handlerBlock, filterBlock?]
                 //  Leave:  [targetBlock]
+                //  Resume: [filterResult?, targetBlock0, targetBlock1, ...]
                 var opers = _currInst.Operands;
-                _operIdx = opers.Length >= 2 && _currInst is not GuardInst ? 1 : 0;
+                _operIdx = opers.Length >= 2 && _currInst is not (GuardInst or ResumeInst { IsFromFilter: false }) ? 1 : 0;
             } else {
                 _currInst = block.First as GuardInst;
                 Ensure.That(block.Last is not GuardInst); // prevents an infinite loop in MoveNext()

src/DistIL/IR/Instructions/EHInsts.cs

@@ -77,16 +77,41 @@ public class ResumeInst : Instruction
             ReplaceOperand(0, value);
         }
     }
+
     [MemberNotNullWhen(true, nameof(FilterResult))]
-    public bool IsFromFilter => Operands.Length > 0;
+    public bool IsFromFilter => Operands is [Value and not BasicBlock, ..];
 
     public override string InstName => "resume";
     public override bool IsBranch => true;
 
-    public ResumeInst(Value? filterResult = null)
-        : base(filterResult == null ? [] : [filterResult]) { }
+    public ResumeInst(IEnumerable<BasicBlock> exitTargets, Value? filterResult = null)
+        : base(filterResult == null ? exitTargets.Distinct().ToArray() : [filterResult, ..exitTargets]) { }
+
+    /// <summary> Unchecked cloning constructor. </summary>
+    public ResumeInst(int _, Value[] operands)
+        : base(operands) { }
 
     public override void Accept(InstVisitor visitor) => visitor.Visit(this);
+
+    public void SetExitTargets(IReadOnlySet<BasicBlock> targets)
+    {
+        int startIdx = IsFromFilter ? 1 : 0;
+        int newCount = targets.Count + startIdx;
+
+        if (newCount > _operands.Length) {
+            startIdx = GrowOperands(newCount - _operands.Length);
+        } else if (newCount < _operands.Length) {
+            RemoveOperands(newCount, _operands.Length - newCount);
+        }
+
+        foreach (var target in targets) {
+            _operands[startIdx] = target;
+            target.AddUse(this, startIdx);
+            startIdx++;
+        }
+    }
+
+    public IEnumerable<BasicBlock> GetExitTargets() => _operands.Skip(IsFromFilter ? 1 : 0).Cast<BasicBlock>();
 }
 
 /// <summary> Throws or rethrows an exception. </summary>

src/DistIL/IR/Utils/IRCloner.cs

@@ -277,6 +277,6 @@ public void Visit(SelectInst inst)
         public void Visit(GuardInst inst) => Out(new GuardInst(inst.Kind, Remap(inst.HandlerBlock), inst.CatchType == null ? null : Remap(inst.CatchType), inst.HasFilter ? Remap(inst.FilterBlock) : null));
         public void Visit(ThrowInst inst) => Out(new ThrowInst(inst.IsRethrow ? null : Remap(inst.Exception)));
         public void Visit(LeaveInst inst) => Out(new LeaveInst(Remap(inst.Target)));
-        public void Visit(ResumeInst inst) => Out(new ResumeInst(inst.IsFromFilter ? Remap(inst.FilterResult) : null));
+        public void Visit(ResumeInst inst) => Out(new ResumeInst(0, RemapArgs(inst.Operands)));
     }
 }

src/DistIL/IR/Utils/IRPrinter.cs

@@ -57,6 +57,9 @@ public static void ExportDot(MethodBody method, TextWriter tw, IReadOnlyList<IPr
                     port = succ == br.Then ? "sw" : "se";
                 } else if (block.Last is LeaveInst) {
                     style.Color = "red";
+                } else if (block.Last is ResumeInst) {
+                    style.Color = "gray";
+                    style.Dashed = true;
                 }
                 if (block.Guards().Any(g => g.HandlerBlock == succ || g.FilterBlock == succ)) {
                     port = "e";

src/DistIL/IR/Utils/IRVerifier.cs

@@ -107,11 +107,29 @@ private void CheckTerminator(BasicBlock block, ProtectedRegionAnalysis regionAna
                 if (region == regionAnalysis.Root) {
                     Error(block, "Cannot leave or resume from root region");
                 } else if (block.Last is ResumeInst resume) {
-                    bool isValid = region.StartBlock.Users().OfType<GuardInst>().FirstOrDefault() is { } guard &&
-                        (resume.IsFromFilter
-                            ? (guard.FilterBlock == region.StartBlock)
-                            : guard.Kind is GuardKind.Finally or GuardKind.Fault);
+                    var guard = region.StartBlock.Users().OfType<GuardInst>().FirstOrDefault();
+
+                    if (guard == null) {
+                        Error(block, "Missing guard for filter or finally region");
+                        break;
+                    }
+
+                    bool isValid = 
+                        resume.IsFromFilter
+                            ? guard.FilterBlock == region.StartBlock
+                            : guard.Kind is GuardKind.Finally or GuardKind.Fault;
                     Check(isValid, block, "ResumeInst should not be used outside filter, finally, or fault handlers");
+
+                    var protectedRegion = regionAnalysis.GetBlockRegion(guard.Block);
+                    var succBlocks = new HashSet<BasicBlock>();
+
+                    foreach (var exitBlock in protectedRegion.GetExitBlocks()) {
+                        Debug.Assert(exitBlock.Last is LeaveInst);
+                        succBlocks.Add(exitBlock.Succs.First());
+                    }
+                    succBlocks.SymmetricExceptWith(resume.GetExitTargets());
+
+                    Check(succBlocks.Count == 0, resume, "ResumeInst targets must be consistent with exits in sibling protected region");
                 }
                 break;
             }

src/DistIL/Passes/Linq/LinqSources.cs

@@ -127,7 +127,7 @@ protected override void EmitEnd(LoopBuilder loop)
             var isNotDisposable = builder.CreateEq(enumer, ConstNull.Create());
             builder.Fork(isNotDisposable, (elseBuilder, newBlock) => elseBuilder.CreateCallVirt(disposeFn, enumer));
         }
-        builder.Emit(new ResumeInst());
+        builder.Emit(new ResumeInst([succBlock]));
 
         loop.Header.Block.InsertFirst(new GuardInst(GuardKind.Finally, finallyBlock));
     }

tests/PracticeTests/EhRegionTests.cs

@@ -0,0 +1,70 @@
+namespace DistIL.PracticeTests;
+
+using DistIL.Attributes;
+
+[Optimize]
+public class EhRegionTests
+{
+    [Theory, InlineData(0), InlineData(1)]
+    [CheckCodeGenAfterRun]
+    public void MultiExit_Finally1(int path)
+    {
+        int x = path;
+
+        try {
+            try {
+                Utils.DoNotOptimize(x); // fake side effect
+
+                if (x > 0) {
+                    goto Target1;
+                } else {
+                    goto Target2;
+                }
+            } finally {
+                x += 10;
+            }
+        Target1:
+            x += 5;
+        Target2:
+            x += 2;
+        } finally {
+            Assert.Equal(path == 0 ? 12 : 18, x);
+        }
+
+        // CHECK: try finally
+        // CHECK: leave [[target1:.*]]
+        // CHECK: leave [[target2:.*]]
+        // CHECK: resume [[target1]], [[target2]]
+    }
+
+    [Fact]
+    public void CrossingVars_Catch1()
+    {
+        int x = 1;
+        string str = "1;0";
+
+        try {
+            int y = int.Parse(str.Split(';')[1]);
+            x = 2;
+            x = 100 / y;
+        } catch (DivideByZeroException) {
+            Assert.Equal(2, x);
+            x = 400;
+        }
+        Assert.Equal(400, x);
+    }
+
+    [Fact]
+    public void CrossingVars_Finally1()
+    {
+        int x = 1;
+        string str = "123";
+
+        try {
+            x = int.Parse(str) * 1000;
+        } finally {
+            x *= 2;
+        }
+        Assert.Equal(123 * 1000 * 2, x);
+    }
+}

tests/PracticeTests/GlobalUsings.cs

@@ -1 +1,11 @@
-global using Xunit;
+global using Xunit;
+
+using System.Runtime.CompilerServices;
+
+using DistIL.Attributes;
+
+public static class Utils
+{
+    [DoNotOptimize, MethodImpl(MethodImplOptions.NoInlining)]
+    public static T DoNotOptimize<T>(T value) => value;
+}

tests/PracticeTests/RegAllocTests.cs

@@ -0,0 +1,26 @@
+namespace DistIL.PracticeTests;
+
+using DistIL.Attributes;
+
+[Optimize]
+public class RegAllocTests
+{
+    [Fact]
+    public void LiveReg_Loop_Finally()
+    {
+        long a = 0, b = 0, c = 0;
+
+        for (int i = 0; i < 100 && c < 100; i++, c++) {
+            try {
+                a += i;
+            } finally {
+                int tmp = (int)(a % 5);  // buggy regalloc will reuse `i`
+                if (tmp == 1 || tmp == 3) {
+                    b += a;
+                }
+            }
+        }
+        Assert.Equal(4950, a);
+        Assert.Equal(99950, b);
+    }
+}