Make if and foreach constructs compatible

libs/UIX.DecompXml/CodeBlockInfo.cs

@@ -1,4 +1,5 @@
 using Microsoft.CodeAnalysis.CSharp.Syntax;
+using System;
 using System.Collections.Generic;
 
 using static Microsoft.CodeAnalysis.CSharp.SyntaxFactory;
@@ -35,7 +36,11 @@ internal record CodeBlockInfo
 
             case ElseBlockInfo _:
                 var elseClause = ElseClause(blockBody);
-                var ifElseBlock = (IfStatementSyntax)parentBlock.Statements[^1];
+
+                var parentLastStatement = parentBlock.Statements[^1];
+                if (parentLastStatement is not IfStatementSyntax ifElseBlock)
+                    throw new InvalidOperationException($"Else block must be preceded by an if block, got '{parentLastStatement.Kind()}'");
+
                 parentBlock.Statements[^1] = ifElseBlock.WithElse(elseClause);
                 break;
 

libs/UIX.DecompXml/Decompiler.Script.cs

@@ -56,37 +56,58 @@ partial class Decompiler
         {
             var instruction = methodBody[i];
 
-            if (jumpFalseToOffsets.Contains(instruction.Offset) && blockStack.Count >= 2)
-            {
-
-                var currentBlock = blockStack.Pop() with { EndOffset = instruction.Offset };
-                currentBlock.FinalizeBlock(blockStack.Peek());
-
-                // This address marks the end of the affirmative branch of an IF clause.
-                // If multiple blocks lead to this address, then we're outside of the IF clause entirely.
-                // Otherwise, it's probably the start of an ELSE clause.
-
-                var currentControlBlock = cfa.GetByInstruction(instruction);
-                if (controlBlocks.Count(b => b.HasEdgeTo(currentControlBlock, controlBlocks)) <= 1)
-                {
-                    blockStack.Push(new(instruction.Offset, uint.MaxValue, new ElseBlockInfo()));
-                }
-            }
-
-            if (!foreachLoopHeadOffsets.Contains(instruction.Offset) && cfa.IsAlwaysExecuted(instruction.Offset))
-            {
             while (blockStack.Count > 1)
             {
                 var currentBlock = blockStack.Pop();
-                    if (currentBlock.EndOffset is uint.MaxValue)
-                        currentBlock = currentBlock with { EndOffset = instruction.Offset };
 
                 if (currentBlock.EndOffset != instruction.Offset)
+                {
+                    blockStack.Push(currentBlock);
                     break;
+                }
 
+                System.Diagnostics.Debug.WriteLine($"END: Finalizing {currentBlock}");
                 currentBlock.FinalizeBlock(blockStack.Peek());
             }
-            }
+
+            //if (jumpToOffsets.Contains(instruction.Offset) && TryPeekBlock<ElseBlockInfo>(out _))
+            //{
+            //    // This address marks the end of the affirmative branch of an IF clause.
+            //    // If multiple blocks lead to this address, then we're outside of the IF clause entirely.
+            //    // Otherwise, it's probably the start of an ELSE clause.
+
+            //    while (blockStack.Count > 1)
+            //    {
+            //        var currentBlock = blockStack.Pop();
+
+            //        if (currentBlock.AdditionalInfo is not (ElseBlockInfo))
+            //            break;
+
+            //        if (currentBlock.EndOffset is uint.MaxValue)
+            //            currentBlock = currentBlock with { EndOffset = instruction.Offset };
+
+            //        System.Diagnostics.Debug.WriteLine($"JMP: Finalizing ELSE {currentBlock}");
+            //        currentBlock.FinalizeBlock(blockStack.Peek());
+            //    }
+            //}
+
+            //if (!foreachLoopHeadOffsets.Contains(instruction.Offset) && cfa.IsAlwaysExecuted(instruction.Offset))
+            //{
+            //    while (blockStack.Count > 1)
+            //    {
+            //        var currentBlock = blockStack.Pop();
+            //        if (currentBlock.EndOffset is uint.MaxValue)
+            //        {
+
+            //        }
+
+            //        if (currentBlock.EndOffset != instruction.Offset || currentBlock.AdditionalInfo is not (IfBlockInfo or ElseBlockInfo))
+            //            break;
+
+            //        System.Diagnostics.Debug.WriteLine($"Always exec'ed: Automatically finalizing {currentBlock}");
+            //        currentBlock.FinalizeBlock(blockStack.Peek());
+            //    }
+            //}
 
             var opCode = instruction.OpCode;
 
@@ -299,7 +320,8 @@ partial class Decompiler
                     case OpCode.JumpIfTruePeek:
                         var jumpToOffset = (uint)instruction.Operands.First().Value;
 
-                        if (opCode is OpCode.JumpIfFalse && TryPeekBlock<ForEachBlockInfo>(out _))
+                        if (opCode is OpCode.JumpIfFalse && TryPeekBlock<ForEachBlockInfo>(out var jmpfForEachBlockInfo)
+                            && jmpfForEachBlockInfo.Type is null)
                             break;
 
                         var isPeek = opCode is OpCode.JumpIfFalsePeek or OpCode.JumpIfTruePeek;
@@ -308,7 +330,13 @@ partial class Decompiler
                         if (opCode is OpCode.JumpIfFalse)
                         {
                             // JMPF is used to evaluate the branch condition
-                            var ifBlock = new CodeBlockInfo(instruction.Offset, jumpToOffset, new IfBlockInfo(jumpCondition));
+                            var ifBlockEndOffset = methodBody
+                               .Reverse()
+                               .SkipWhile(i => i.Offset >= jumpToOffset)
+                               .First()
+                               .Offset;
+
+                            var ifBlock = new CodeBlockInfo(instruction.Offset, ifBlockEndOffset, new IfBlockInfo(jumpCondition));
                             blockStack.Push(ifBlock);
                         }
                         else
@@ -329,14 +357,33 @@ partial class Decompiler
 
                             if (foreachLoopHeadOffsets.Contains(jumpOffset))
                             {
-                                var currentBlock = blockStack.Pop() with { EndOffset = instruction.Offset };
+                                //var currentBlock = blockStack.Pop() with { EndOffset = instruction.Offset };
-                                currentBlock.FinalizeBlock(blockStack.Peek());
+                                //System.Diagnostics.Debug.WriteLine($"JMP: Finalizing loop {currentBlock}");
+                                //currentBlock.FinalizeBlock(blockStack.Peek());
                             }
                             else
                             {
                                 throw new NotImplementedException("For and while loops are not supported at this time.");
                             }
                         }
+                        else
+                        {
+                            // End of if block, skipping else block
+
+                            // Figure out where the else block ends by searching for the last instruction we skip
+                            var elseBlockEndOffset = methodBody
+                                .Reverse()
+                                .SkipWhile(i => i.Offset >= jumpOffset)
+                                .First()
+                                .Offset;
+
+                            //var currentBlock = blockStack.Pop() with { EndOffset = instruction.Offset };
+                            //System.Diagnostics.Debug.WriteLine($"JMP: Finalizing presumed IF {currentBlock}");
+                            //currentBlock.FinalizeBlock(blockStack.Peek());
+
+                            var elseBlock = new CodeBlockInfo(instruction.Offset, elseBlockEndOffset, new ElseBlockInfo());
+                            blockStack.Push(elseBlock);
+                        }
 
                         break;
 
@@ -380,6 +427,8 @@ partial class Decompiler
         return blockStack.Pop().Statements;
 
         bool TryPeekBlock<T>([NotNullWhen(true)] out T additionalInfo) where T : ICodeBlockAdditionalInfo
+        {
+            if (blockStack.Count > 1)
             {
                 var currentBlock = blockStack.Peek();
                 if (currentBlock.AdditionalInfo is T a)
@@ -387,6 +436,7 @@ partial class Decompiler
                     additionalInfo = a;
                     return true;
                 }
+            }
 
             additionalInfo = default;
             return false;