Skip to content

Commit dde0933

Browse files
PhilippNausedPhilippNaused
committed
Implement improved MaxStackCalculator
1 parent 207e5e1 commit dde0933

1 file changed

Lines changed: 163 additions & 177 deletions

File tree

‎src/DotNet/Writer/MaxStackCalculator.cs‎

Lines changed: 163 additions & 177 deletions
Original file line numberDiff line numberDiff line change
@@ -1,177 +1,163 @@
1-
// dnlib: See LICENSE.txt for more info
2-
3-
using System.Collections.Generic;
4-
using dnlib.DotNet.Emit;
5-
6-
namespace dnlib.DotNet.Writer {
7-
/// <summary>
8-
/// Calculates max stack usage by using a simple pass over all instructions. This value
9-
/// can be placed in the fat method header's MaxStack field.
10-
/// </summary>
11-
public struct MaxStackCalculator {
12-
IList<Instruction> instructions;
13-
IList<ExceptionHandler> exceptionHandlers;
14-
readonly Dictionary<Instruction, int> stackHeights;
15-
bool hasError;
16-
int currentMaxStack;
17-
18-
/// <summary>
19-
/// Gets max stack value
20-
/// </summary>
21-
/// <param name="instructions">All instructions</param>
22-
/// <param name="exceptionHandlers">All exception handlers</param>
23-
/// <returns>Max stack value</returns>
24-
public static uint GetMaxStack(IList<Instruction> instructions, IList<ExceptionHandler> exceptionHandlers) {
25-
new MaxStackCalculator(instructions, exceptionHandlers).Calculate(out uint maxStack);
26-
return maxStack;
27-
}
28-
29-
/// <summary>
30-
/// Gets max stack value
31-
/// </summary>
32-
/// <param name="instructions">All instructions</param>
33-
/// <param name="exceptionHandlers">All exception handlers</param>
34-
/// <param name="maxStack">Updated with max stack value</param>
35-
/// <returns><c>true</c> if no errors were detected, <c>false</c> otherwise</returns>
36-
public static bool GetMaxStack(IList<Instruction> instructions, IList<ExceptionHandler> exceptionHandlers, out uint maxStack) =>
37-
new MaxStackCalculator(instructions, exceptionHandlers).Calculate(out maxStack);
38-
39-
internal static MaxStackCalculator Create() => new MaxStackCalculator(true);
40-
41-
MaxStackCalculator(bool dummy) {
42-
instructions = null;
43-
exceptionHandlers = null;
44-
stackHeights = new Dictionary<Instruction, int>();
45-
hasError = false;
46-
currentMaxStack = 0;
47-
}
48-
49-
MaxStackCalculator(IList<Instruction> instructions, IList<ExceptionHandler> exceptionHandlers) {
50-
this.instructions = instructions;
51-
this.exceptionHandlers = exceptionHandlers;
52-
stackHeights = new Dictionary<Instruction, int>();
53-
hasError = false;
54-
currentMaxStack = 0;
55-
}
56-
57-
internal void Reset(IList<Instruction> instructions, IList<ExceptionHandler> exceptionHandlers) {
58-
this.instructions = instructions;
59-
this.exceptionHandlers = exceptionHandlers;
60-
stackHeights.Clear();
61-
hasError = false;
62-
currentMaxStack = 0;
63-
}
64-
65-
internal bool Calculate(out uint maxStack) {
66-
var exceptionHandlers = this.exceptionHandlers;
67-
var stackHeights = this.stackHeights;
68-
for (int i = 0; i < exceptionHandlers.Count; i++) {
69-
var eh = exceptionHandlers[i];
70-
if (eh is null)
71-
continue;
72-
Instruction instr;
73-
if ((instr = eh.TryStart) is not null)
74-
stackHeights[instr] = 0;
75-
if ((instr = eh.FilterStart) is not null) {
76-
stackHeights[instr] = 1;
77-
currentMaxStack = 1;
78-
}
79-
if ((instr = eh.HandlerStart) is not null) {
80-
bool pushed = eh.IsCatch || eh.IsFilter;
81-
if (pushed) {
82-
stackHeights[instr] = 1;
83-
currentMaxStack = 1;
84-
}
85-
else
86-
stackHeights[instr] = 0;
87-
}
88-
}
89-
90-
int stack = 0;
91-
bool resetStack = false;
92-
var instructions = this.instructions;
93-
for (int i = 0; i < instructions.Count; i++) {
94-
var instr = instructions[i];
95-
if (instr is null)
96-
continue;
97-
98-
if (resetStack) {
99-
stackHeights.TryGetValue(instr, out stack);
100-
resetStack = false;
101-
}
102-
stack = WriteStack(instr, stack);
103-
var opCode = instr.OpCode;
104-
var code = opCode.Code;
105-
if (code == Code.Jmp) {
106-
if (stack != 0)
107-
hasError = true;
108-
}
109-
else {
110-
instr.CalculateStackUsage(out int pushes, out int pops);
111-
if (pops == -1)
112-
stack = 0;
113-
else {
114-
stack -= pops;
115-
if (stack < 0) {
116-
hasError = true;
117-
stack = 0;
118-
}
119-
stack += pushes;
120-
}
121-
}
122-
if (stack < 0) {
123-
hasError = true;
124-
stack = 0;
125-
}
126-
127-
switch (opCode.FlowControl) {
128-
case FlowControl.Branch:
129-
WriteStack(instr.Operand as Instruction, stack);
130-
resetStack = true;
131-
break;
132-
133-
case FlowControl.Call:
134-
if (code == Code.Jmp)
135-
resetStack = true;
136-
break;
137-
138-
case FlowControl.Cond_Branch:
139-
if (code == Code.Switch) {
140-
if (instr.Operand is IList<Instruction> targets) {
141-
for (int j = 0; j < targets.Count; j++)
142-
WriteStack(targets[j], stack);
143-
}
144-
}
145-
else
146-
WriteStack(instr.Operand as Instruction, stack);
147-
break;
148-
149-
case FlowControl.Return:
150-
case FlowControl.Throw:
151-
resetStack = true;
152-
break;
153-
}
154-
}
155-
156-
maxStack = (uint)currentMaxStack;
157-
return !hasError;
158-
}
159-
160-
int WriteStack(Instruction instr, int stack) {
161-
if (instr is null) {
162-
hasError = true;
163-
return stack;
164-
}
165-
var stackHeights = this.stackHeights;
166-
if (stackHeights.TryGetValue(instr, out int stack2)) {
167-
if (stack != stack2)
168-
hasError = true;
169-
return stack2;
170-
}
171-
stackHeights[instr] = stack;
172-
if (stack > currentMaxStack)
173-
currentMaxStack = stack;
174-
return stack;
175-
}
176-
}
177-
}
1+
// dnlib: See LICENSE.txt for more info
2+
3+
using System.Collections.Generic;
4+
using System.Linq;
5+
using dnlib.DotNet.Emit;
6+
7+
namespace dnlib.DotNet.Writer {
8+
/// <summary>
9+
/// Calculates max stack usage by using a simple pass over all instructions. This value
10+
/// can be placed in the fat method header's MaxStack field.
11+
/// </summary>
12+
public struct MaxStackCalculator {
13+
IList<Instruction> instructions;
14+
IList<ExceptionHandler> exceptionHandlers;
15+
uint?[] stackHeights;
16+
bool hasError;
17+
18+
/// <summary>
19+
/// Gets max stack value
20+
/// </summary>
21+
/// <param name="instructions">All instructions</param>
22+
/// <param name="exceptionHandlers">All exception handlers</param>
23+
/// <returns>Max stack value</returns>
24+
public static uint GetMaxStack(IList<Instruction> instructions, IList<ExceptionHandler> exceptionHandlers) {
25+
new MaxStackCalculator(instructions, exceptionHandlers).Calculate(out uint maxStack);
26+
return maxStack;
27+
}
28+
29+
/// <summary>
30+
/// Gets max stack value
31+
/// </summary>
32+
/// <param name="instructions">All instructions</param>
33+
/// <param name="exceptionHandlers">All exception handlers</param>
34+
/// <param name="maxStack">Updated with max stack value</param>
35+
/// <returns><c>true</c> if no errors were detected, <c>false</c> otherwise</returns>
36+
public static bool GetMaxStack(IList<Instruction> instructions, IList<ExceptionHandler> exceptionHandlers, out uint maxStack) =>
37+
new MaxStackCalculator(instructions, exceptionHandlers).Calculate(out maxStack);
38+
39+
/// <summary>
40+
/// Gets the stack height for each instruction.
41+
/// </summary>
42+
/// <param name="instructions">All instructions</param>
43+
/// <param name="exceptionHandlers">All exception handlers</param>
44+
/// <returns>Max stack value</returns>
45+
public static uint?[] GetStackHeights(IList<Instruction> instructions, IList<ExceptionHandler> exceptionHandlers) {
46+
var helper = new MaxStackCalculator(instructions, exceptionHandlers);
47+
helper.ExploreAll();
48+
return helper.stackHeights;
49+
}
50+
51+
internal static MaxStackCalculator Create() => new MaxStackCalculator();
52+
53+
MaxStackCalculator(IList<Instruction> instructions, IList<ExceptionHandler> exceptionHandlers) => Reset(instructions, exceptionHandlers);
54+
55+
internal void Reset(IList<Instruction> instructions, IList<ExceptionHandler> exceptionHandlers) {
56+
this.instructions = instructions;
57+
this.exceptionHandlers = exceptionHandlers;
58+
stackHeights = new uint?[instructions.Count];
59+
hasError = false;
60+
}
61+
62+
internal bool Calculate(out uint maxStack) {
63+
ExploreAll();
64+
maxStack = stackHeights.Max() ?? 0;
65+
return !hasError;
66+
}
67+
68+
void ExploreAll() {
69+
Explore(0, 0);
70+
71+
foreach (var handler in exceptionHandlers) {
72+
if (handler!.FilterStart is not null) {
73+
var idx = IndexOf(handler.FilterStart);
74+
Explore(idx, 1);
75+
}
76+
if (handler!.HandlerStart is not null) {
77+
var idx = IndexOf(handler.HandlerStart);
78+
bool pushed = handler.IsCatch || handler.IsFilter;
79+
Explore(idx, pushed ? 1u : 0u);
80+
}
81+
}
82+
}
83+
84+
readonly int IndexOf(Instruction instr) {
85+
var index = instructions.IndexOf(instr);
86+
if (index < 0)
87+
throw new InvalidMethodException($"Instruction {instr} not found.");
88+
return index;
89+
}
90+
91+
void Explore(int index, uint stackHeight) {
92+
start:
93+
var previous = stackHeights[index];
94+
if (previous is not null) {
95+
if (previous != stackHeight) {
96+
hasError = true;
97+
}
98+
return; // already visited this instruction with the same stack height, no need to explore further
99+
}
100+
stackHeights[index] = stackHeight;
101+
102+
var instr = instructions[index];
103+
instr.CalculateStackUsage(out int pushes, out int pops);
104+
if (pops == -1) {
105+
stackHeight = 0;
106+
}
107+
else {
108+
if (stackHeight < pops) {
109+
hasError = true;
110+
return;
111+
}
112+
stackHeight -= (uint)pops;
113+
stackHeight += (uint)pushes;
114+
}
115+
switch (instr.OpCode.FlowControl) {
116+
case FlowControl.Break:
117+
case FlowControl.Call:
118+
case FlowControl.Meta:
119+
case FlowControl.Next: {
120+
if (instr.OpCode.Code == Code.Jmp) {
121+
return; // method terminates here, no need to explore further
122+
}
123+
else {
124+
index++;
125+
goto start; // just continue to the next instruction
126+
}
127+
}
128+
case FlowControl.Return: {
129+
if (stackHeight > 1)
130+
hasError = true;
131+
return; // method terminates here
132+
}
133+
case FlowControl.Throw: {
134+
return; // method terminates here
135+
}
136+
case FlowControl.Branch: // unconditional branch
137+
{
138+
var target = (Instruction)instr.Operand;
139+
index = IndexOf(target);
140+
goto start; // tail recursion
141+
}
142+
case FlowControl.Cond_Branch: {
143+
if (instr.OpCode.Code == Code.Switch) {
144+
foreach (var target in (IList<Instruction>)instr.Operand) {
145+
Explore(IndexOf(target), stackHeight); // explore the branch target
146+
}
147+
index++;
148+
goto start; // explore the next instruction
149+
}
150+
else {
151+
var target = (Instruction)instr.Operand;
152+
Explore(IndexOf(target), stackHeight); // explore the branch target
153+
index++;
154+
goto start; // explore the next instruction
155+
}
156+
}
157+
default:
158+
hasError = true;
159+
return;
160+
}
161+
}
162+
}
163+
}

0 commit comments

Comments
 (0)