ILSpy/ICSharpCode.Decompiler/Ast/Transforms/AddCheckedBlocks.cs

// Copyright (c) AlphaSierraPapa for the SharpDevelop Team (for details please see \doc\copyright.txt)
// This code is distributed under MIT X11 license (for details please see \doc\license.txt)

using System;
using System.Linq;
using ICSharpCode.Decompiler.ILAst;
using ICSharpCode.NRefactory.CSharp;

namespace ICSharpCode.Decompiler.Ast.Transforms
{
	/// <summary>
	/// Add checked/unchecked blocks.
	/// </summary>
	public class AddCheckedBlocks : IAstTransform
	{
		#region Annotation
		sealed class CheckedUncheckedAnnotation {
			/// <summary>
			/// true=checked, false=unchecked
			/// </summary>
			public bool IsChecked;
		}
		
		public static readonly object CheckedAnnotation = new CheckedUncheckedAnnotation { IsChecked = true };
		public static readonly object UncheckedAnnotation = new CheckedUncheckedAnnotation { IsChecked = false };
		#endregion
		
		/* 
			We treat placing checked/unchecked blocks as an optimization problem, with the following goals:
				1. Use minimum number of checked blocks+expressions
				2. Prefer checked expressions over checked blocks
				3. Make the scope of checked expressions as small as possible
				4. Make the scope of checked blocks as large as possible
				(where goal 1 has the highest priority)
		 */
		
		#region struct Cost
		struct Cost
		{
			// highest possible cost so that the Blocks+Expressions addition doesn't overflow
			public static readonly Cost Infinite = new Cost(0x3fffffff, 0x3fffffff);
			
			public readonly int Blocks;
			public readonly int Expressions;
			
			public Cost(int blocks, int expressions)
			{
				this.Blocks = blocks;
				this.Expressions = expressions;
			}
			
			public static bool operator <(Cost a, Cost b)
			{
				return a.Blocks + a.Expressions < b.Blocks + b.Expressions
					|| a.Blocks + a.Expressions == b.Blocks + b.Expressions && a.Blocks < b.Blocks;
			}
			
			public static bool operator >(Cost a, Cost b)
			{
				return a.Blocks + a.Expressions > b.Blocks + b.Expressions
					|| a.Blocks + a.Expressions == b.Blocks + b.Expressions && a.Blocks > b.Blocks;
			}
			
			public static bool operator <=(Cost a, Cost b)
			{
				return a.Blocks + a.Expressions < b.Blocks + b.Expressions
					|| a.Blocks + a.Expressions == b.Blocks + b.Expressions && a.Blocks <= b.Blocks;
			}
			
			public static bool operator >=(Cost a, Cost b)
			{
				return a.Blocks + a.Expressions > b.Blocks + b.Expressions
					|| a.Blocks + a.Expressions == b.Blocks + b.Expressions && a.Blocks >= b.Blocks;
			}
			
			public static Cost operator +(Cost a, Cost b)
			{
				return new Cost(a.Blocks + b.Blocks, a.Expressions + b.Expressions);
			}
			
			public override string ToString()
			{
				return string.Format("[{0} + {1}]", Blocks, Expressions);
			}
		}
		#endregion
		
		#region class InsertedNode
		/// <summary>
		/// Holds the blocks and expressions that should be inserted
		/// </summary>
		abstract class InsertedNode
		{
			public static InsertedNode operator +(InsertedNode a, InsertedNode b)
			{
				if (a == null)
					return b;
				if (b == null)
					return a;
				return new InsertedNodeList(a, b);
			}
			
			public abstract void Insert();
		}
		
		class InsertedNodeList : InsertedNode
		{
			readonly InsertedNode child1, child2;
			
			public InsertedNodeList(AddCheckedBlocks.InsertedNode child1, AddCheckedBlocks.InsertedNode child2)
			{
				this.child1 = child1;
				this.child2 = child2;
			}
			
			public override void Insert()
			{
				child1.Insert();
				child2.Insert();
			}
		}
		
		class InsertedExpression : InsertedNode
		{
			readonly Expression expression;
			readonly bool isChecked;
			
			public InsertedExpression(Expression expression, bool isChecked)
			{
				this.expression = expression;
				this.isChecked = isChecked;
			}
			
			public override void Insert()
			{
				if (isChecked)
					expression.ReplaceWith(e => new CheckedExpression { Expression = e });
				else
					expression.ReplaceWith(e => new UncheckedExpression { Expression = e });
			}
		}
		
		class ConvertCompoundAssignment : InsertedNode
		{
			readonly Expression expression;
			readonly bool isChecked;
			
			public ConvertCompoundAssignment(Expression expression, bool isChecked)
			{
				this.expression = expression;
				this.isChecked = isChecked;
			}
			
			public override void Insert()
			{
				AssignmentExpression assign = expression.Annotation<ReplaceMethodCallsWithOperators.RestoreOriginalAssignOperatorAnnotation>().Restore(expression);
				expression.ReplaceWith(assign);
				if (isChecked)
					assign.Right = new CheckedExpression { Expression = assign.Right.Detach() };
				else
					assign.Right = new UncheckedExpression { Expression = assign.Right.Detach() };
			}
		}
		
		class InsertedBlock : InsertedNode
		{
			readonly Statement firstStatement; // inclusive
			readonly Statement lastStatement; // exclusive
			readonly bool isChecked;
			
			public InsertedBlock(Statement firstStatement, Statement lastStatement, bool isChecked)
			{
				this.firstStatement = firstStatement;
				this.lastStatement = lastStatement;
				this.isChecked = isChecked;
			}
			
			public override void Insert()
			{
				BlockStatement newBlock = new BlockStatement();
				// Move all statements except for the first
				Statement next;
				for (Statement stmt = firstStatement.GetNextStatement(); stmt != lastStatement; stmt = next) {
					next = stmt.GetNextStatement();
					newBlock.Add(stmt.Detach());
				}
				// Replace the first statement with the new (un)checked block
				if (isChecked)
					firstStatement.ReplaceWith(new CheckedStatement { Body = newBlock });
				else
					firstStatement.ReplaceWith(new UncheckedStatement { Body = newBlock });
				// now also move the first node into the new block
				newBlock.Statements.InsertAfter(null, firstStatement);
			}
		}
		#endregion
		
		#region class Result
		/// <summary>
		/// Holds the result of an insertion operation.
		/// </summary>
		class Result
		{
			public Cost CostInCheckedContext;
			public InsertedNode NodesToInsertInCheckedContext;
			public Cost CostInUncheckedContext;
			public InsertedNode NodesToInsertInUncheckedContext;
		}
		#endregion
		
		public void Run(AstNode node)
		{
			BlockStatement block = node as BlockStatement;
			if (block == null) {
				for (AstNode child = node.FirstChild; child != null; child = child.NextSibling) {
					Run(child);
				}
			} else {
				Result r = GetResultFromBlock(block);
				if (r.NodesToInsertInUncheckedContext != null)
					r.NodesToInsertInUncheckedContext.Insert();
			}
		}
		
		Result GetResultFromBlock(BlockStatement block)
		{
			// For a block, we are tracking 4 possibilities:
			// a) context is checked, no unchecked block open
			Cost costCheckedContext = new Cost(0, 0);
			InsertedNode nodesCheckedContext = null;
			// b) context is checked, an unchecked block is open
			Cost costCheckedContextUncheckedBlockOpen = Cost.Infinite;
			InsertedNode nodesCheckedContextUncheckedBlockOpen = null;
			Statement uncheckedBlockStart = null;
			// c) context is unchecked, no checked block open
			Cost costUncheckedContext = new Cost(0, 0);
			InsertedNode nodesUncheckedContext = null;
			// d) context is unchecked, a checked block is open
			Cost costUncheckedContextCheckedBlockOpen = Cost.Infinite;
			InsertedNode nodesUncheckedContextCheckedBlockOpen = null;
			Statement checkedBlockStart = null;
			
			Statement statement = block.Statements.FirstOrDefault();
			while (true) {
				// Blocks can be closed 'for free'. We use '<=' so that blocks are closed as late as possible (goal 4)
				if (costCheckedContextUncheckedBlockOpen <= costCheckedContext) {
					costCheckedContext = costCheckedContextUncheckedBlockOpen;
					nodesCheckedContext = nodesCheckedContextUncheckedBlockOpen + new InsertedBlock(uncheckedBlockStart, statement, false);
				}
				if (costUncheckedContextCheckedBlockOpen <= costUncheckedContext) {
					costUncheckedContext = costUncheckedContextCheckedBlockOpen;
					nodesUncheckedContext = nodesUncheckedContextCheckedBlockOpen + new InsertedBlock(checkedBlockStart, statement, true);
				}
				if (statement == null)
					break;
				// Now try opening blocks. We use '<' so that blocks are opened as early as possible. (goal 4)
				if (costCheckedContext + new Cost(1, 0) < costCheckedContextUncheckedBlockOpen) {
					costCheckedContextUncheckedBlockOpen = costCheckedContext + new Cost(1, 0);
					nodesCheckedContextUncheckedBlockOpen = nodesCheckedContext;
					uncheckedBlockStart = statement;
				}
				if (costUncheckedContext + new Cost(1, 0) < costUncheckedContextCheckedBlockOpen) {
					costUncheckedContextCheckedBlockOpen = costUncheckedContext + new Cost(1, 0);
					nodesUncheckedContextCheckedBlockOpen = nodesUncheckedContext;
					checkedBlockStart = statement;
				}
				// Now handle the statement
				Result stmtResult = GetResult(statement);
				
				costCheckedContext += stmtResult.CostInCheckedContext;
				nodesCheckedContext += stmtResult.NodesToInsertInCheckedContext;
				costCheckedContextUncheckedBlockOpen += stmtResult.CostInUncheckedContext;
				nodesCheckedContextUncheckedBlockOpen += stmtResult.NodesToInsertInUncheckedContext;
				costUncheckedContext += stmtResult.CostInUncheckedContext;
				nodesUncheckedContext += stmtResult.NodesToInsertInUncheckedContext;
				costUncheckedContextCheckedBlockOpen += stmtResult.CostInCheckedContext;
				nodesUncheckedContextCheckedBlockOpen += stmtResult.NodesToInsertInCheckedContext;
				
				statement = statement.GetNextStatement();
			}
			
			return new Result {
				CostInCheckedContext = costCheckedContext, NodesToInsertInCheckedContext = nodesCheckedContext,
				CostInUncheckedContext = costUncheckedContext, NodesToInsertInUncheckedContext = nodesUncheckedContext
			};
		}
		
		Result GetResult(AstNode node)
		{
			if (node is BlockStatement)
				return GetResultFromBlock((BlockStatement)node);
			Result result = new Result();
			for (AstNode child = node.FirstChild; child != null; child = child.NextSibling) {
				Result childResult = GetResult(child);
				result.CostInCheckedContext += childResult.CostInCheckedContext;
				result.NodesToInsertInCheckedContext += childResult.NodesToInsertInCheckedContext;
				result.CostInUncheckedContext += childResult.CostInUncheckedContext;
				result.NodesToInsertInUncheckedContext += childResult.NodesToInsertInUncheckedContext;
			}
			Expression expr = node as Expression;
			if (expr != null) {
				CheckedUncheckedAnnotation annotation = expr.Annotation<CheckedUncheckedAnnotation>();
				if (annotation != null) {
					// If the annotation requires this node to be in a specific context, add a huge cost to the other context
					// That huge cost gives us the option to ignore a required checked/unchecked expression when there wouldn't be any
					// solution otherwise. (e.g. "for (checked(M().x += 1); true; unchecked(M().x += 2)) {}")
					if (annotation.IsChecked)
						result.CostInUncheckedContext += new Cost(10000, 0);
					else
						result.CostInCheckedContext += new Cost(10000, 0);
				}
				// Embed this node in an checked/unchecked expression:
				if (expr.Parent is ExpressionStatement) {
					// We cannot use checked/unchecked for top-level-expressions.
					// However, we could try converting a compound assignment (checked(a+=b);) or unary operator (checked(a++);)
					// back to its old form.
					if (expr.Annotation<ReplaceMethodCallsWithOperators.RestoreOriginalAssignOperatorAnnotation>() != null) {
						// We use '<' so that expressions are introduced on the deepest level possible (goal 3)
						if (result.CostInCheckedContext + new Cost(1, 1) < result.CostInUncheckedContext) {
							result.CostInUncheckedContext = result.CostInCheckedContext + new Cost(1, 1);
							result.NodesToInsertInUncheckedContext = result.NodesToInsertInCheckedContext + new ConvertCompoundAssignment(expr, true);
						} else if (result.CostInUncheckedContext + new Cost(1, 1) < result.CostInCheckedContext) {
							result.CostInCheckedContext = result.CostInUncheckedContext + new Cost(1, 1);
							result.NodesToInsertInCheckedContext = result.NodesToInsertInUncheckedContext + new ConvertCompoundAssignment(expr, false);
						}
					}
				} else {
					// We use '<' so that expressions are introduced on the deepest level possible (goal 3)
					if (result.CostInCheckedContext + new Cost(0, 1) < result.CostInUncheckedContext) {
						result.CostInUncheckedContext = result.CostInCheckedContext + new Cost(0, 1);
						result.NodesToInsertInUncheckedContext = result.NodesToInsertInCheckedContext + new InsertedExpression(expr, true);
					} else if (result.CostInUncheckedContext + new Cost(0, 1) < result.CostInCheckedContext) {
						result.CostInCheckedContext = result.CostInUncheckedContext + new Cost(0, 1);
						result.NodesToInsertInCheckedContext = result.NodesToInsertInUncheckedContext + new InsertedExpression(expr, false);
					}
				}
			}
			return result;
		}
	}
}