From 8b1250d52bdc2ffc23d060acf4e7a8ab1cbfb0d2 Mon Sep 17 00:00:00 2001 From: PhilippNaused Date: Fri, 13 Feb 2026 08:45:48 +0100 Subject: [PATCH] Implement improved MaxStackCalculator --- src/DotNet/Writer/MaxStackCalculator.cs | 336 +++++++++++------------- 1 file changed, 159 insertions(+), 177 deletions(-) diff --git a/src/DotNet/Writer/MaxStackCalculator.cs b/src/DotNet/Writer/MaxStackCalculator.cs index ea745e2d6..30dbe0fdc 100644 --- a/src/DotNet/Writer/MaxStackCalculator.cs +++ b/src/DotNet/Writer/MaxStackCalculator.cs @@ -1,177 +1,159 @@ -// dnlib: See LICENSE.txt for more info - -using System.Collections.Generic; -using dnlib.DotNet.Emit; - -namespace dnlib.DotNet.Writer { - /// - /// Calculates max stack usage by using a simple pass over all instructions. This value - /// can be placed in the fat method header's MaxStack field. - /// - public struct MaxStackCalculator { - IList instructions; - IList exceptionHandlers; - readonly Dictionary stackHeights; - bool hasError; - int currentMaxStack; - - /// - /// Gets max stack value - /// - /// All instructions - /// All exception handlers - /// Max stack value - public static uint GetMaxStack(IList instructions, IList exceptionHandlers) { - new MaxStackCalculator(instructions, exceptionHandlers).Calculate(out uint maxStack); - return maxStack; - } - - /// - /// Gets max stack value - /// - /// All instructions - /// All exception handlers - /// Updated with max stack value - /// true if no errors were detected, false otherwise - public static bool GetMaxStack(IList instructions, IList exceptionHandlers, out uint maxStack) => - new MaxStackCalculator(instructions, exceptionHandlers).Calculate(out maxStack); - - internal static MaxStackCalculator Create() => new MaxStackCalculator(true); - - MaxStackCalculator(bool dummy) { - instructions = null; - exceptionHandlers = null; - stackHeights = new Dictionary(); - hasError = false; - currentMaxStack = 0; - } - - MaxStackCalculator(IList instructions, IList exceptionHandlers) { - this.instructions = instructions; - this.exceptionHandlers = exceptionHandlers; - stackHeights = new Dictionary(); - hasError = false; - currentMaxStack = 0; - } - - internal void Reset(IList instructions, IList exceptionHandlers) { - this.instructions = instructions; - this.exceptionHandlers = exceptionHandlers; - stackHeights.Clear(); - hasError = false; - currentMaxStack = 0; - } - - internal bool Calculate(out uint maxStack) { - var exceptionHandlers = this.exceptionHandlers; - var stackHeights = this.stackHeights; - for (int i = 0; i < exceptionHandlers.Count; i++) { - var eh = exceptionHandlers[i]; - if (eh is null) - continue; - Instruction instr; - if ((instr = eh.TryStart) is not null) - stackHeights[instr] = 0; - if ((instr = eh.FilterStart) is not null) { - stackHeights[instr] = 1; - currentMaxStack = 1; - } - if ((instr = eh.HandlerStart) is not null) { - bool pushed = eh.IsCatch || eh.IsFilter; - if (pushed) { - stackHeights[instr] = 1; - currentMaxStack = 1; - } - else - stackHeights[instr] = 0; - } - } - - int stack = 0; - bool resetStack = false; - var instructions = this.instructions; - for (int i = 0; i < instructions.Count; i++) { - var instr = instructions[i]; - if (instr is null) - continue; - - if (resetStack) { - stackHeights.TryGetValue(instr, out stack); - resetStack = false; - } - stack = WriteStack(instr, stack); - var opCode = instr.OpCode; - var code = opCode.Code; - if (code == Code.Jmp) { - if (stack != 0) - hasError = true; - } - else { - instr.CalculateStackUsage(out int pushes, out int pops); - if (pops == -1) - stack = 0; - else { - stack -= pops; - if (stack < 0) { - hasError = true; - stack = 0; - } - stack += pushes; - } - } - if (stack < 0) { - hasError = true; - stack = 0; - } - - switch (opCode.FlowControl) { - case FlowControl.Branch: - WriteStack(instr.Operand as Instruction, stack); - resetStack = true; - break; - - case FlowControl.Call: - if (code == Code.Jmp) - resetStack = true; - break; - - case FlowControl.Cond_Branch: - if (code == Code.Switch) { - if (instr.Operand is IList targets) { - for (int j = 0; j < targets.Count; j++) - WriteStack(targets[j], stack); - } - } - else - WriteStack(instr.Operand as Instruction, stack); - break; - - case FlowControl.Return: - case FlowControl.Throw: - resetStack = true; - break; - } - } - - maxStack = (uint)currentMaxStack; - return !hasError; - } - - int WriteStack(Instruction instr, int stack) { - if (instr is null) { - hasError = true; - return stack; - } - var stackHeights = this.stackHeights; - if (stackHeights.TryGetValue(instr, out int stack2)) { - if (stack != stack2) - hasError = true; - return stack2; - } - stackHeights[instr] = stack; - if (stack > currentMaxStack) - currentMaxStack = stack; - return stack; - } - } -} +// dnlib: See LICENSE.txt for more info + +using System.Collections.Generic; +using dnlib.DotNet.Emit; + +namespace dnlib.DotNet.Writer { + /// + /// Calculates max stack usage by using a simple pass over all instructions. This value + /// can be placed in the fat method header's MaxStack field. + /// + public struct MaxStackCalculator { + IList instructions; + IList exceptionHandlers; + ushort?[] stackHeights; + uint maxStack; + bool hasError; + + /// + /// Gets max stack value + /// + /// All instructions + /// All exception handlers + /// Max stack value + public static uint GetMaxStack(IList instructions, IList exceptionHandlers) { + new MaxStackCalculator(instructions, exceptionHandlers).Calculate(out uint maxStack); + return maxStack; + } + + /// + /// Gets max stack value + /// + /// All instructions + /// All exception handlers + /// Updated with max stack value + /// true if no errors were detected, false otherwise + public static bool GetMaxStack(IList instructions, IList exceptionHandlers, out uint maxStack) => + new MaxStackCalculator(instructions, exceptionHandlers).Calculate(out maxStack); + + /// + /// Gets the stack height for each instruction + /// + /// All instructions + /// All exception handlers + /// The stack height for each instruction + public static ushort?[] GetStackHeights(IList instructions, IList exceptionHandlers) { + var helper = new MaxStackCalculator(instructions, exceptionHandlers); + helper.ExploreAll(); + return helper.stackHeights; + } + + internal static MaxStackCalculator Create() => new MaxStackCalculator(); + + MaxStackCalculator(IList instructions, IList exceptionHandlers) => Reset(instructions, exceptionHandlers); + + internal void Reset(IList instructions, IList exceptionHandlers) { + this.instructions = instructions; + this.exceptionHandlers = exceptionHandlers; + stackHeights = new ushort?[instructions.Count]; + maxStack = 0; + hasError = false; + } + + internal bool Calculate(out uint maxStack) { + ExploreAll(); + maxStack = this.maxStack; + return !hasError; + } + + void ExploreAll() { + Explore(0, 0); + + foreach (var handler in exceptionHandlers) { + if (handler.FilterStart is not null) { + Explore(handler.FilterStart, 1); + } + if (handler.HandlerStart is not null) { + bool pushed = handler.IsCatch || handler.IsFilter; + Explore(handler.HandlerStart, (ushort)(pushed ? 1 : 0)); + } + } + } + + void Explore(Instruction instr, ushort stackHeight) => Explore(instructions.IndexOf(instr), stackHeight); + + void Explore(int index, ushort stackHeight) { +start: + var previous = stackHeights[index]; + if (previous is not null) { + if (previous != stackHeight) { + hasError = true; + } + return; // already visited this instruction + } + stackHeights[index] = stackHeight; + if (stackHeight > maxStack) + maxStack = stackHeight; + + var instr = instructions[index]; + instr.CalculateStackUsage(out int pushes, out int pops); + if (pops == -1) { + stackHeight = 0; + } + else { + if (stackHeight < pops) { + hasError = true; // stack underflow + return; + } + stackHeight -= (ushort)pops; + stackHeight += (ushort)pushes; + } + switch (instr.OpCode.FlowControl) { + case FlowControl.Break: + case FlowControl.Call: + case FlowControl.Meta: + case FlowControl.Next: { + if (instr.OpCode.Code == Code.Jmp) { + return; // method terminates here + } + else { + index++; + goto start; // just continue to the next instruction + } + } + case FlowControl.Return: { + if (stackHeight > 1) + hasError = true; + return; // method terminates here + } + case FlowControl.Throw: { + return; // method terminates here + } + case FlowControl.Branch: // unconditional branch + { + var target = (Instruction)instr.Operand; + index = instructions.IndexOf(target); + goto start; // tail recursion + } + case FlowControl.Cond_Branch: { + if (instr.OpCode.Code == Code.Switch) { + foreach (var target in (IList)instr.Operand) { + Explore(target, stackHeight); // explore the branch target + } + index++; + goto start; // explore the next instruction + } + else { + var target = (Instruction)instr.Operand; + Explore(target, stackHeight); // explore the branch target + index++; + goto start; // explore the next instruction + } + } + default: + hasError = true; + return; + } + } + } +}