Move stack analysis to the ByteCode class

Decompiler.csproj

@@ -43,10 +43,11 @@
     <Compile Include="src\ByteCode.StackBehaviour.cs" />
     <Compile Include="src\ByteCode.Type.cs" />
     <Compile Include="src\ByteCodeCollection.cs" />
+    <Compile Include="src\CilStack.cs" />
     <Compile Include="src\MainForm.cs" />
     <Compile Include="src\MainForm.Designer.cs" />
     <Compile Include="src\Program.cs" />
-    <Compile Include="src\StackAnalysis.cs" />
+    <Compile Include="src\ByteCode.StackAnalysis.cs" />
     <Compile Include="src\Util.cs" />
     <EmbeddedResource Include="src\MainForm.resx">
       <DependentUpon>MainForm.cs</DependentUpon>

src/AstMetodBodyBuilder.cs

@@ -19,7 +19,6 @@ namespace Decompiler
 			methodDef.Body.Simplify();
 			
 			ByteCodeCollection body = new ByteCodeCollection(methodDef);
-			StackAnalysis stackAnalysis = new StackAnalysis(methodDef, body);
 			
 			foreach(VariableDefinition varDef in methodDef.Body.Variables) {
 				Ast.VariableDeclaration astVar = new Ast.VariableDeclaration(varDef.Name);
@@ -44,7 +43,7 @@ namespace Decompiler
 					int argCount = byteCode.PopCount;
 					Ast.Expression[] args = new Ast.Expression[argCount];
 					for(int i = 0; i < argCount; i++) {
-						ByteCode allocBy = stackAnalysis.StackBefore[byteCode].Peek(argCount - i).AllocadedBy;
+						ByteCode allocBy = byteCode.StackBefore.Peek(argCount - i).AllocadedBy;
 						string name = string.Format("expr{0:X2}", allocBy.Offset);
 						args[i] = new Ast.IdentifierExpression(name);
 					}
@@ -54,7 +53,7 @@ namespace Decompiler
 						args);
 					if (codeExpr is Ast.Expression) {
 						if (byteCode.PushCount == 1) {
-							string type = stackAnalysis.GetTypeOf(byteCode).FullName;
+							string type = byteCode.Type.FullName;
 							string name = string.Format("expr{0:X2}", byteCode.Offset);
 							Ast.LocalVariableDeclaration astLocal = new Ast.LocalVariableDeclaration(new Ast.TypeReference(type.ToString()));
 							astLocal.Variables.Add(new Ast.VariableDeclaration(name, (Ast.Expression)codeExpr));
@@ -69,7 +68,7 @@ namespace Decompiler
 					astStatement = MakeComment(description);
 				}
 				//astBlock.Children.Add(MakeComment(description));
-				if (stackAnalysis.BranchTargetOf[byteCode].Count > 0) {
+				if (byteCode.BranchesHere.Count > 0) {
 					astBlock.Children.Add(new Ast.LabelStatement(string.Format("IL_{0:X2}", byteCode.Offset)));
 				}
 				astBlock.Children.Add(astStatement);

src/ByteCode.StackAnalysis.cs

@@ -0,0 +1,71 @@
+using System;
+using System.Collections.Generic;
+
+using Cecil = Mono.Cecil;
+using Mono.Cecil;
+using Mono.Cecil.Cil;
+
+namespace Decompiler
+{
+	public partial class ByteCode
+	{
+		CilStack stackBefore;
+		CilStack stackAfter;
+		List<ByteCode> branchesHere = new List<ByteCode>();
+		
+		public CilStack StackBefore {
+			get { return stackBefore; }
+		}
+		
+		public CilStack StackAfter {
+			get { return stackAfter; }
+		}
+		
+		public List<ByteCode> BranchesHere {
+			get { return branchesHere; }
+		}
+		
+		public void MergeStackBeforeWith(CilStack stack)
+		{
+			CilStack mergedStack;
+			if (CilStack.Merge(this.stackBefore, stack, out mergedStack)) {
+				// Stacks are identical
+				return;
+			}
+			this.stackBefore = mergedStack;
+			
+			stackAfter = SimulateEffectOnStack(this.StackBefore);
+			
+			switch(this.OpCode.FlowControl) {
+				case FlowControl.Branch:
+					this.BranchTarget.MergeStackBeforeWith(this.StackAfter);
+					break;
+				case FlowControl.Cond_Branch:
+					this.Next.MergeStackBeforeWith(this.StackAfter);
+					this.BranchTarget.MergeStackBeforeWith(this.StackAfter);
+					break;
+				case FlowControl.Next:
+				case FlowControl.Call:
+					this.Next.MergeStackBeforeWith(this.StackAfter);
+					break;
+				case FlowControl.Return:
+					if (this.StackAfter.Count > 0) throw new Exception("Non-empty stack at return instruction");
+					break;
+				default: throw new NotImplementedException();
+			}
+		}
+		
+		public CilStack SimulateEffectOnStack(CilStack oldStack)
+		{
+			CilStack newStack = oldStack.Clone();
+			List<Cecil.TypeReference> typeArgs = new List<Cecil.TypeReference>();
+			foreach(CilStackSlot slot in newStack.PopCount(this.PopCount)) {
+				typeArgs.Add(slot.Type);
+			}
+			for (int i = 0; i < this.PushCount; i++) {
+				newStack.Push(new CilStackSlot(this, this.GetType(typeArgs.ToArray())));
+			}
+			return newStack;
+		}
+	}
+}

src/ByteCode.Type.cs

@@ -23,6 +23,16 @@ namespace Decompiler
 			return new Cecil.TypeReference(type.Name, type.Namespace, null, type.IsValueType);
 		}
 		
+		public Cecil.TypeReference Type {
+			get {
+				if (this.PushCount == 0) {
+					return TypeVoid;
+				} else {
+					return this.StackAfter.Peek(1).Type;
+				}
+			}
+		}
+		
 		public Cecil.TypeReference GetType(params Cecil.TypeReference[] args)
 		{
 			try {

src/ByteCodeCollection.cs

@@ -49,15 +49,6 @@ namespace Decompiler
 			return list.IndexOf(byteCode);
 		}
 		
-		public void Remove(ByteCode byteCode)
-		{
-			int index = IndexOf(byteCode);
-			list.RemoveAt(index);
-			byteCode.Next = null;
-			byteCode.Previous = null;
-			UpdateNextPrevious();
-		}
-		
 		public ByteCodeCollection(MethodDefinition methodDef)
 		{
 			foreach(Instruction inst in methodDef.Body.Instructions) {
@@ -68,7 +59,13 @@ namespace Decompiler
 					byteCode.Operand = GetByOffset(((Instruction)byteCode.Operand).Offset);
 				}
 			}
+			foreach(ByteCode byteCode in this) {
+				if (byteCode.CanBranch) {
+					byteCode.BranchTarget.BranchesHere.Add(byteCode);
+				}
+			}
 			UpdateNextPrevious();
+			UpdateStackAnalysis();
 		}
 		
 		void UpdateNextPrevious()
@@ -78,5 +75,12 @@ namespace Decompiler
 				this[i + 1].Previous = this[i];
 			}
 		}
+		
+		void UpdateStackAnalysis()
+		{
+			if (this.Count > 0) {
+				this[0].MergeStackBeforeWith(CilStack.Empty);
+			}
+		}
 	}
 }

src/CilStack.cs

@@ -0,0 +1,135 @@
+using System;
+using System.Collections.Generic;
+
+using Ast = ICSharpCode.NRefactory.Ast;
+using ICSharpCode.NRefactory.Ast;
+
+using Cecil = Mono.Cecil;
+using Mono.Cecil;
+using Mono.Cecil.Cil;
+
+namespace Decompiler
+{
+	// Imutable
+	public struct CilStackSlot {
+		ByteCode allocadedBy;
+		Cecil.TypeReference type;
+		
+		public ByteCode AllocadedBy {
+			get { return allocadedBy; }
+		}
+		
+		public Cecil.TypeReference Type {
+			get { return type; }
+		}
+		
+		public CilStackSlot(ByteCode allocadedBy, Cecil.TypeReference type)
+		{
+			this.allocadedBy = allocadedBy;
+			this.type = type;
+		}
+		
+		public override int GetHashCode()
+		{
+			int hashCode = 0;
+			if (allocadedBy != null) hashCode ^= allocadedBy.GetHashCode(); 
+			if (type != null) hashCode ^= type.GetHashCode(); 
+			return hashCode;
+		}
+		
+		public override bool Equals(object obj)
+		{
+			if (!(obj is CilStackSlot)) return false; 
+			CilStackSlot myCilStackSlot = (CilStackSlot)obj;
+			return object.Equals(this.allocadedBy, myCilStackSlot.allocadedBy) && object.Equals(this.type, myCilStackSlot.type);
+		}
+		
+		public override string ToString()
+		{
+			if (allocadedBy == null) {
+				return "<??>";
+			} else {
+				return string.Format("expr{0:X2}", this.allocadedBy.Offset);
+			}
+		}
+	}
+	
+	/// <remarks> The tail of the list is the top of the stack </remarks>
+	public class CilStack: List<CilStackSlot> {
+		public static CilStack Empty = new CilStack();
+		
+		public CilStack Clone()
+		{
+			return new CilStack(this);
+		}
+		
+		public CilStackSlot[] PopCount(int count)
+		{
+			CilStackSlot[] poped = this.GetRange(this.Count - count, count).ToArray();
+			this.RemoveRange(this.Count - count, count);
+			return poped;
+		}
+		
+		public void Push(CilStackSlot slot)
+		{
+			this.Add(slot);
+		}
+		
+		public CilStackSlot Peek(int depth)
+		{
+			return this[this.Count - depth];
+		}
+		
+		public CilStack(): base()
+		{
+			
+		}
+		
+		public CilStack(IEnumerable<CilStackSlot> slotEnum): base(slotEnum)
+		{
+			
+		}
+		
+		// Return true if stacks are same
+		public static bool Merge(CilStack stack1, CilStack stack2, out CilStack merged)
+		{
+			// Both null
+			if (stack1 == null && stack2 == null) {
+				throw new Exception("Both stacks are null");
+			}
+			// One of stacks null, one is not
+			if (stack1 == null || stack2 == null) {
+				merged = stack1 ?? stack2;
+				return false;
+			}
+			// Both are non-null
+			if (stack1.Count != stack2.Count) {
+				throw new Exception("Stack merge error: different sizes");
+			}
+			
+			bool same = true;
+			int count = stack1.Count;
+			merged = stack1.Clone();
+			for (int i = 0; i < count; i++) {
+				if (!stack1[i].Equals(stack2[i])) {
+					merged[i] = new CilStackSlot(null, null); // Merge slots
+					same = false;
+				}
+			}
+			return same;
+		}
+		
+		public override string ToString()
+		{
+			string ret = "Stack: {";
+			bool first = true;
+			foreach(CilStackSlot slot in this) {
+				if (!first) ret += ", ";
+				ret += slot.ToString();
+				first = false;
+			}
+			ret += "}";
+			return ret;
+		}
+	}
+}

src/StackAnalysis.cs

@@ -1,230 +0,0 @@
-using System;
-using System.Collections.Generic;
-
-using Ast = ICSharpCode.NRefactory.Ast;
-using ICSharpCode.NRefactory.Ast;
-
-using Cecil = Mono.Cecil;
-using Mono.Cecil;
-using Mono.Cecil.Cil;
-
-namespace Decompiler
-{
-	// Imutable
-	public struct CilStackSlot {
-		ByteCode allocadedBy;
-		Cecil.TypeReference type;
-		
-		public ByteCode AllocadedBy {
-			get { return allocadedBy; }
-		}
-		
-		public Cecil.TypeReference Type {
-			get { return type; }
-		}
-		
-		public CilStackSlot(ByteCode allocadedBy, Cecil.TypeReference type)
-		{
-			this.allocadedBy = allocadedBy;
-			this.type = type;
-		}
-		
-		public override int GetHashCode()
-		{
-			int hashCode = 0;
-			if (allocadedBy != null) hashCode ^= allocadedBy.GetHashCode(); 
-			if (type != null) hashCode ^= type.GetHashCode(); 
-			return hashCode;
-		}
-		
-		public override bool Equals(object obj)
-		{
-			if (!(obj is CilStackSlot)) return false; 
-			CilStackSlot myCilStackSlot = (CilStackSlot)obj;
-			return object.Equals(this.allocadedBy, myCilStackSlot.allocadedBy) && object.Equals(this.type, myCilStackSlot.type);
-		}
-		
-		public override string ToString()
-		{
-			if (allocadedBy == null) {
-				return "<??>";
-			} else {
-				return string.Format("expr{0:X2}", this.allocadedBy.Offset);
-			}
-		}
-	}
-	
-	/// <remarks> The tail of the list is the top of the stack </remarks>
-	public class CilStack: List<CilStackSlot> {
-		public static CilStack Empty = new CilStack();
-		
-		public CilStack Clone()
-		{
-			return new CilStack(this);
-		}
-		
-		public CilStackSlot[] PopCount(int count)
-		{
-			CilStackSlot[] poped = this.GetRange(this.Count - count, count).ToArray();
-			this.RemoveRange(this.Count - count, count);
-			return poped;
-		}
-		
-		public void Push(CilStackSlot slot)
-		{
-			this.Add(slot);
-		}
-		
-		public CilStackSlot Peek(int depth)
-		{
-			return this[this.Count - depth];
-		}
-		
-		public CilStack(): base()
-		{
-			
-		}
-		
-		public CilStack(IEnumerable<CilStackSlot> slotEnum): base(slotEnum)
-		{
-			
-		}
-		
-		public override string ToString()
-		{
-			string ret = "Stack: {";
-			bool first = true;
-			foreach(CilStackSlot slot in this) {
-				if (!first) ret += ", ";
-				ret += slot.ToString();
-				first = false;
-			}
-			ret += "}";
-			return ret;
-		}
-	}
-	
-	public partial class StackAnalysis {
-		MethodDefinition methodDef;
-		Dictionary<ByteCode, CilStack> stackBefore = new Dictionary<ByteCode, CilStack>();
-		Dictionary<ByteCode, CilStack> stackAfter = new Dictionary<ByteCode, CilStack>();
-		Dictionary<ByteCode, List<ByteCode>> branchTargetOf = new Dictionary<ByteCode, List<ByteCode>>();
-		
-		public Dictionary<ByteCode, CilStack> StackBefore {
-			get { return stackBefore; }
-		}
-		
-		public Dictionary<ByteCode, CilStack> StackAfter {
-			get { return stackAfter; }
-		}
-		
-		public Dictionary<ByteCode, List<ByteCode>> BranchTargetOf {
-			get { return branchTargetOf; }
-		}
-		
-		public Cecil.TypeReference GetTypeOf(ByteCode byteCode)
-		{
-			if (byteCode.PushCount == 0) {
-				return ByteCode.TypeVoid;
-			} else {
-				return StackAfter[byteCode].Peek(1).Type;
-			}
-		}
-		
-		public StackAnalysis(MethodDefinition methodDef, ByteCodeCollection byteCodeCol) {
-			this.methodDef = methodDef;
-			
-			foreach(ByteCode byteCode in byteCodeCol) {
-				stackBefore[byteCode] = null;
-				stackAfter[byteCode] = null;
-				branchTargetOf[byteCode] = new List<ByteCode>();
-			}
-			
-			foreach(ByteCode byteCode in byteCodeCol) {
-				if (byteCode.CanBranch) {
-					branchTargetOf[byteCode.BranchTarget].Add(byteCode);
-				}
-			}
-			
-			if (byteCodeCol.Count > 0) {
-				ByteCode firstInst = byteCodeCol[0];
-				stackBefore[firstInst] = CilStack.Empty;
-				ProcessByteCodeRec(firstInst);
-			}
-		}
-		
-		void ProcessByteCodeRec(ByteCode byteCode)
-		{
-			stackAfter[byteCode] = ChangeStack(stackBefore[byteCode], byteCode);
-			
-			switch(byteCode.OpCode.FlowControl) {
-				case FlowControl.Branch:
-					CopyStack(byteCode, byteCode.BranchTarget);
-					break;
-				case FlowControl.Cond_Branch:
-					CopyStack(byteCode, byteCode.Next);
-					CopyStack(byteCode, byteCode.BranchTarget);
-					break;
-				case FlowControl.Next:
-				case FlowControl.Call:
-					CopyStack(byteCode, byteCode.Next);
-					break;
-				case FlowControl.Return:
-					if (stackAfter[byteCode].Count > 0) throw new Exception("Non-empty stack at the end");
-					break;
-				default: throw new NotImplementedException();
-			}
-		}
-		
-		CilStack ChangeStack(CilStack oldStack, ByteCode byteCode)
-		{
-			CilStack newStack = oldStack.Clone();
-			CilStackSlot[] popedSlots = newStack.PopCount(byteCode.PopCount);
-			List<Cecil.TypeReference> typeArgs = new List<Cecil.TypeReference>();
-			foreach(CilStackSlot slot in popedSlots) {
-				typeArgs.Add(slot.Type);
-			}
-			for (int i = 0; i < byteCode.PushCount; i++) {
-				newStack.Push(new CilStackSlot(byteCode, byteCode.GetType(typeArgs.ToArray())));
-			}
-			return newStack;
-		}
-		
-		void CopyStack(ByteCode byteCodeFrom, ByteCode byteCodeTo)
-		{
-			CilStack mergedStack;
-			if (!Merge(stackAfter[byteCodeFrom], stackBefore[byteCodeTo], out mergedStack)) {
-				stackBefore[byteCodeTo] = mergedStack;
-				ProcessByteCodeRec(byteCodeTo);
-			}
-		}
-		
-		bool Merge(CilStack stack1, CilStack stack2, out CilStack merged)
-		{
-			// Both null
-			if (stack1 == null && stack2 == null) {
-				throw new Exception("Both stacks are null");
-			}
-			// One of stacks null, one is not
-			if (stack1 == null || stack2 == null) {
-				merged = stack1 ?? stack2;
-				return false;
-			}
-			// Both are non-null
-			if (stack1.Count != stack2.Count) {
-				throw new Exception("Stack merge error: different sizes");
-			}
-			
-			bool same = true;
-			int count = stack1.Count;
-			merged = stack1.Clone();
-			for (int i = 0; i < count; i++) {
-				if (!stack1[i].Equals(stack2[i])) {
-					merged[i] = new CilStackSlot(null, null); // Merge slots
-					same = false;
-				}
-			}
-			return same;
-		}
-	}
-}