From f9ab6c58c1904a5a3b737b8d5a8c6ad7f2c0e4d2 Mon Sep 17 00:00:00 2001 From: Yoshi Askharoun Date: Wed, 24 Sep 2025 17:46:17 -0500 Subject: [PATCH] Add brute-force dominance computation --- libs/UIX.DecompXml/ControlFlowAnalyzer.cs | 137 ++++++++++++---------- libs/UIX.DecompXml/Decompiler.Script.cs | 26 ++-- 2 files changed, 94 insertions(+), 69 deletions(-) diff --git a/libs/UIX.DecompXml/ControlFlowAnalyzer.cs b/libs/UIX.DecompXml/ControlFlowAnalyzer.cs index d09085d..9ab0856 100644 --- a/libs/UIX.DecompXml/ControlFlowAnalyzer.cs +++ b/libs/UIX.DecompXml/ControlFlowAnalyzer.cs @@ -75,47 +75,63 @@ public static class ControlFlowAnalyzer blocks.Add(block); } - // Pass III: Resolve next and branch target blocks - for (var b = 0; b < blocks.Count; b++) - { - var block = (BasicControlFlowBlock)blocks[b]; - - if (block.NextOffset is not uint.MaxValue) - block = block with { Next = blocks[IndexOfBlockFromStartOffset(block.NextOffset)] }; - - if (block.BranchTargetOffset is not uint.MaxValue) - block = block with { BranchTarget = blocks[IndexOfBlockFromStartOffset(block.BranchTargetOffset)] }; - - blocks[b] = block; - } - return blocks; - - int IndexOfBlockFromStartOffset(uint startOffset) - { - for (int b = 0; b <= blocks.Count; b++) - { - var block = blocks[b]; - if (block.StartOffset == startOffset) - return b; - } - - return -1; - } } - public static List CollapseBlocks(this List blocks) + public static IProgramBlock GetByOffset(this List blocks, uint offset) { - for (int b = 0; b < blocks.Count; b++) + return blocks.First(b => offset >= b.StartOffset && offset <= b.EndOffset); + } + + public static IProgramBlock GetByStartOffset(this List blocks, uint offset) => blocks.FirstOrDefault(b => offset == b.StartOffset); + + public static IProgramBlock GetByInstruction(this List blocks, Instruction instruction) => blocks.GetByOffset(instruction.Offset); + + public static bool IsAlwaysExecuted(this List blocks, uint offset) + { + if (blocks.Count == 1) + return true; + + // Checks if the entry node is post-dominated by the block containing this offset. + // Essentially, do all paths through this method execute code at this offset? + + var blockOfInterest = blocks.GetByStartOffset(offset); + if (blockOfInterest is null) + return false; + + HashSet visitedStartOffsets = []; + Stack stack = []; + + var filteredBlocks = blocks + .OrderBy(b => b.StartOffset) + .Where(b => b.StartOffset != offset); + stack.Push(filteredBlocks.First()); + + while (stack.Count > 0) { - var block = blocks[b]; - if (block.Body[^1].OpCode is OpCode.JumpIfFalse) + var current = stack.Pop(); + + if (visitedStartOffsets.Contains(current.StartOffset)) + continue; + + visitedStartOffsets.Add(current.StartOffset); + + foreach (var childOffset in current.GetChildrenStartOffsets()) { - // If conditions always end with JMPF + if (childOffset == blockOfInterest.StartOffset) + continue; + + var child = blocks.GetByStartOffset(childOffset); + stack.Push(child); } } - return []; + HashSet exitBlockOffsets = new(blocks + .Where(b => b.Body[^1].OpCode is OpCode.ReturnValue or OpCode.ReturnVoid) + .Select(b => b.StartOffset)); + + visitedStartOffsets.IntersectWith(exitBlockOffsets); + return visitedStartOffsets.Count == 0; } public static string SerializeToGraphviz(IEnumerable blocks) @@ -160,45 +176,48 @@ public interface IProgramBlock { uint StartOffset { get; } uint EndOffset { get; } - uint NextOffset { get; } List Body { get; } - IProgramBlock Next { get; } + uint NextOffset { get; } - bool HasEdgeTo(IProgramBlock block); + IEnumerable GetChildrenStartOffsets(); } -public abstract record ProgramBlock(uint StartOffset, uint EndOffset, List Body, - uint NextOffset = uint.MaxValue, IProgramBlock Next = null) +public abstract record ProgramBlock(uint StartOffset, uint EndOffset, List Body, uint NextOffset = uint.MaxValue) : IProgramBlock { - public virtual bool HasEdgeTo(IProgramBlock block) => NextOffset == block.StartOffset; + public virtual IEnumerable GetChildrenStartOffsets() + { + if (NextOffset is not uint.MaxValue) + yield return NextOffset; + } } public record BasicControlFlowBlock(uint StartOffset, uint EndOffset, List Body, - uint NextOffset = uint.MaxValue, uint BranchTargetOffset = uint.MaxValue, - IProgramBlock Next = null, IProgramBlock BranchTarget = null) - : ProgramBlock(StartOffset, EndOffset, Body, NextOffset, Next) + uint NextOffset = uint.MaxValue, uint BranchTargetOffset = uint.MaxValue) + : ProgramBlock(StartOffset, EndOffset, Body, NextOffset) { - public override bool HasEdgeTo(IProgramBlock block) => base.HasEdgeTo(block) || BranchTargetOffset == block.StartOffset; -} - -public record ConditionalControlFlowBlock(uint StartOffset, uint EndOffset, List Body, - uint NextOffset = uint.MaxValue, IProgramBlock Next = null, - IProgramBlock IfBlock = null, List ElseIfBlocks = null, IProgramBlock ElseBlock = null) - : ProgramBlock(StartOffset, EndOffset, Body, NextOffset, Next) -{ - public ConditionalControlFlowBlock(IProgramBlock programBlock, - IProgramBlock ifBlock = null, List elseIfBlocks = null, IProgramBlock elseBlock = null) - : this(programBlock.StartOffset, programBlock.EndOffset, programBlock.Body, - programBlock.NextOffset, programBlock.Next, ifBlock, elseIfBlocks ?? [], elseBlock) + public override IEnumerable GetChildrenStartOffsets() { - } + foreach (var offset in base.GetChildrenStartOffsets()) + yield return offset; - public override bool HasEdgeTo(IProgramBlock block) + if (BranchTargetOffset is not uint.MaxValue) + yield return BranchTargetOffset; + } +} + +public static class IProgramBlockExtensions +{ + public static bool HasEdgeTo(this IProgramBlock block, IProgramBlock targetBlock, List blocks) { - return base.HasEdgeTo(block) - || IfBlock.StartOffset == block.StartOffset - || ElseIfBlocks.Any(b => b.StartOffset == block.StartOffset) - || ElseBlock.StartOffset == block.StartOffset; + return block + .GetChildren(blocks) + .Contains(targetBlock); + } + + public static IEnumerable GetChildren(this IProgramBlock block, List blocks) + { + foreach (var startOffset in block.GetChildrenStartOffsets()) + yield return blocks.First(b => b.StartOffset == startOffset); } } diff --git a/libs/UIX.DecompXml/Decompiler.Script.cs b/libs/UIX.DecompXml/Decompiler.Script.cs index ff01074..d05ef8f 100644 --- a/libs/UIX.DecompXml/Decompiler.Script.cs +++ b/libs/UIX.DecompXml/Decompiler.Script.cs @@ -54,6 +54,11 @@ partial class Decompiler if (jumpFalseToOffsets.Contains(instruction.Offset)) { + if (blockStack.Count < 2) + { + throw new InvalidOperationException("Expected two blocks left on the stack"); + } + var currentBlock = blockStack.Pop() with { EndOffset = instruction.Offset }; currentBlock.FinalizeBlock(blockStack.Peek()); @@ -61,24 +66,25 @@ partial class Decompiler // 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 = controlBlocks.First(b => instruction.Offset >= b.StartOffset && instruction.Offset <= b.EndOffset); - if (controlBlocks.Count(b => b.HasEdgeTo(currentControlBlock)) <= 1) + var currentControlBlock = controlBlocks.GetByInstruction(instruction); + if (controlBlocks.Count(b => b.HasEdgeTo(currentControlBlock, controlBlocks)) <= 1) { blockStack.Push(new(instruction.Offset, uint.MaxValue, SyntaxKind.ElseClause, null)); } } - if (jumpToOffsets.Contains(instruction.Offset)) + if (controlBlocks.IsAlwaysExecuted(instruction.Offset)) { - // This address is code that will be unconditionally executed. - // For now, we'll assume that this is the end of IF/ELSE clauses. while (blockStack.Count > 1) { - if (blockStack.Peek().Kind is SyntaxKind.ElseClause) - { - var currentBlock = blockStack.Pop() with { EndOffset = instruction.Offset }; - currentBlock.FinalizeBlock(blockStack.Peek()); - } + var currentBlock = blockStack.Pop(); + if (currentBlock.EndOffset is uint.MaxValue) + currentBlock = currentBlock with { EndOffset = instruction.Offset }; + + if (currentBlock.EndOffset != instruction.Offset) + break; + + currentBlock.FinalizeBlock(blockStack.Peek()); } }