V-Gears/V-Gears-Installer/src/decompiler/control_flow.cpp

/* ScummVM Tools
 *
 * ScummVM Tools is the legal property of its developers, whose
 * names are too numerous to list here. Please refer to the
 * COPYRIGHT file distributed with this source distribution.
 *
 * This program is free software; you can redistribute it and/or
 * modify it under the terms of the GNU General Public License
 * as published by the Free Software Foundation; either version 2
 * of the License, or (at your option) any later version.
 *
 * This program is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU General Public License for more details.
 *
 * You should have received a copy of the GNU General Public License
 * along with this program; if not, write to the Free Software
 * Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301, USA.
 */

#include "decompiler/control_flow.h"
#include "decompiler/stack.h"

#include <algorithm>
#include <iostream>
#include <set>

#include <boost/format.hpp>

#define PUT(vertex, group) boost::put(boost::vertex_name, _g, vertex, group);
#define PUT_EDGE(edge, isJump) boost::put(boost::edge_attribute, _g, edge, isJump);
#define PUT_ID(vertex, id) boost::put(boost::vertex_index, _g, vertex, id);
#define GET(vertex) (boost::get(boost::vertex_name, _g, vertex))
#define GET_EDGE(edge) (boost::get(boost::edge_attribute, _g, edge))

ControlFlow::ControlFlow(InstVec& insts, Engine& engine)
  : mInsts(insts), 
    mEngine(engine)
{

	// Automatically add a function if we're not supposed to look for more functions and no functions are defined
	// This avoids a special case for when no real functions exist in the script
    if (mEngine._functions.empty())
    {
        mEngine._functions[(*insts.begin())->_address] = Function((*insts.begin())->_address, (insts.back())->_address);
    }

	GroupPtr prev = NULL;
	int id = 0;
	// Create vertices
	for (InstIterator it = insts.begin(); it != insts.end(); ++it) {
		GraphVertex cur = boost::add_vertex(_g);
		_addrMap[(*it)->_address] = cur;
		PUT(cur, new Group(cur, it, it, prev));
		PUT_ID(cur, id);
		id++;

		// Add reference to vertex if function starts here
		if (mEngine._functions.find((*it)->_address) != mEngine._functions.end())
			mEngine._functions[(*it)->_address]._v = cur;

		prev = GET(cur);
	}

	// Add regular edges
	FuncMap::iterator fn;
    GraphVertex last = {};
	bool addEdge = false;
	prev = NULL;
	for (InstIterator it = insts.begin(); it != insts.end(); ++it) {
		if (mEngine._functions.find((*it)->_address) != mEngine._functions.end()) {
			addEdge = false;
		}

		GraphVertex cur = find(it);
		if (addEdge) {
			GraphEdge e = boost::add_edge(last, cur, _g).first;
			PUT_EDGE(e, false);
		}

		last = cur;
		addEdge = !((*it)->IsUncondJump() || (*it)->isReturn());
		prev = GET(cur);

	}

	// Add jump edges
	for (InstIterator it = insts.begin(); it != insts.end(); ++it) {
		if ((*it)->isJump()) {
			GraphEdge e = boost::add_edge(find(it), find((*it)->GetDestAddress()), _g).first;
			PUT_EDGE(e, true);
		}
	}
}

GraphVertex ControlFlow::find(const InstPtr inst) {
	return _addrMap[inst->_address];
}

GraphVertex ControlFlow::find(ConstInstIterator it) {
	return _addrMap[(*it)->_address];
}

GraphVertex ControlFlow::find(uint32 address) {
	std::map<uint32, GraphVertex>::iterator it = _addrMap.find(address);
	if (it == _addrMap.end())
		std::cerr << "Request for instruction at unknown address " << boost::format("0x%08x") % address << std::endl;
	return it->second;
}

void ControlFlow::merge(GraphVertex g1, GraphVertex g2) {
	// Update property
	GroupPtr gr1 = GET(g1);
	GroupPtr gr2 = GET(g2);
	gr1->end_ = gr2->end_;
	PUT(g1, gr1);

	// Update address map
	ConstInstIterator it = gr2->start_;
	do {
		_addrMap[(*it)->_address] = g1;
		++it;
	} while (gr2->start_ != gr2->end_ && it != gr2->end_);

	// Add outgoing edges from g2
	OutEdgeRange r = boost::out_edges(g2, _g);
	for (OutEdgeIterator e = r.first; e != r.second; ++e) {
		GraphEdge newE = boost::add_edge(g1, boost::target(*e, _g), _g).first;
		PUT_EDGE(newE, GET_EDGE(*e));
	}

	// Update _next pointer
	gr1->_next = gr2->_next;
	if (gr2->_next != NULL)
		gr2->_next->_prev = gr2->_prev;

	// Remove edges to/from g2
	boost::clear_vertex(g2, _g);
	// Remove vertex
	boost::remove_vertex(g2, _g);
}

typedef std::pair<GraphVertex, int> LevelEntry;

void ControlFlow::setStackLevel(GraphVertex g, int level) {
	Stack<LevelEntry> levelStack;
	std::set<GraphVertex> seen;
	levelStack.push(LevelEntry(g, level));
	seen.insert(g);
	while (!levelStack.empty()) {
		LevelEntry e = levelStack.pop();
		GroupPtr gr = GET(e.first);
		if (gr->_stackLevel != -1) {
			if (gr->_stackLevel != e.second)
				std::cerr << boost::format("WARNING: Inconsistency in expected stack level for instruction at address 0x%08x (current: %d, requested: %d)\n") % (*gr->start_)->_address % gr->_stackLevel % e.second;
			continue;
		}
		gr->_stackLevel = e.second;

		OutEdgeRange r = boost::out_edges(e.first, _g);
		for (OutEdgeIterator oe = r.first; oe != r.second; ++oe) {
			GraphVertex target = boost::target(*oe, _g);
			if (seen.find(target) == seen.end()) {
				levelStack.push(LevelEntry(target, e.second + (*gr->start_)->_stackChange));
				seen.insert(target);
			}
		}
	}
}

void ControlFlow::createGroups() 
{
    if (!mEngine._functions.empty() && GET(mEngine._functions.begin()->second._v)->_stackLevel != -1)
    {
        return;
    }

	for (FuncMap::iterator fn = mEngine._functions.begin(); fn != mEngine._functions.end(); ++fn)
		setStackLevel(fn->second._v, 0);
	ConstInstIterator curInst, nextInst;
	nextInst = mInsts.begin();
	nextInst++;
	int stackLevel = 0;
	int expectedStackLevel = 0;
	for (curInst = mInsts.begin(); nextInst != mInsts.end(); ++curInst, ++nextInst) {
		GraphVertex cur = find(curInst);
		GraphVertex next = find(nextInst);

		GroupPtr grCur = GET(cur);
		GroupPtr grNext = GET(next);

		// Don't process unreachable code
		if (grCur->_stackLevel < 0) {
			stackLevel = grNext->_stackLevel;
			continue;
		}

		expectedStackLevel = grCur->_stackLevel;
		// If expected stack level decreases in next vertex, then use next vertex level as expected level
		if (expectedStackLevel > grNext->_stackLevel && grNext->_stackLevel >= 0) {
			expectedStackLevel = grNext->_stackLevel;
			// Also set the stack level of the current group to remember that we expect it to be lower
			grCur->_stackLevel = expectedStackLevel;
		}

		stackLevel += (*curInst)->_stackChange;

		// For stack operations, the new stack level becomes the expected stack level starting from the next group
		if ((*curInst)->isStackOp()) {
			expectedStackLevel = stackLevel;
			grNext->_stackLevel = stackLevel;
		}

		// Group ends after a jump
		if ((*curInst)->isJump()) {
			stackLevel = grNext->_stackLevel;
			continue;
		}

		// Group ends with a return
		if ((*curInst)->isReturn()) {
			stackLevel = grNext->_stackLevel;
			continue;
		}

		// Group ends before target of a jump
		if (in_degree(next, _g) != 1) {
			stackLevel = grNext->_stackLevel;
			continue;
		}

		// This part is only relevant if we use the stack level.
		if (!mEngine.UsePureGrouping()) {
			// If group has no instructions with stack effect >= 0, don't merge on balanced stack
			bool forceMerge = true;
			ConstInstIterator it = grCur->start_;
			do {
				if ((*it)->_stackChange >= 0)
					forceMerge = false;
				++it;
			} while (grCur->start_ != grCur->end_ && it != grCur->end_);

			// Group ends when stack is balanced, unless just before conditional jump
			if (stackLevel == expectedStackLevel && !forceMerge && !(*nextInst)->isCondJump()) {
				continue;
			}
		}

		// All checks passed, merge groups
		merge(cur, next);
	}

	// FIXME: The short-circuit detection is disabled because short-circuited groups require some special handling
	// in the code generation. It's not entirely clear how to handle it properly, though: you need to deduce which
	// effect is created by the conditional jumps in the middle of a block, which seems to get fairly complex when
	// you have multiple groups that are merged by the short-circuit detection.
	//detectShortCircuit();
}

void ControlFlow::detectShortCircuit() {
	ConstInstIterator lastInst = mInsts.end();
	--lastInst;
	GraphVertex cur = find(lastInst);
	GroupPtr gr = GET(cur);
	while (gr->_prev != NULL) {
		bool doMerge = false;
		cur = find(gr->start_);
		GraphVertex prev = find(gr->_prev->start_);
		// Block is candidate for short-circuit merging if it and the preceding block both end with conditional jumps
		if (out_degree(cur, _g) == 2 && out_degree(prev, _g) == 2) {
			doMerge = true;
			OutEdgeRange rCur = boost::out_edges(cur, _g);
			std::vector<GraphVertex> succs;

			// Find possible target vertices
			for (OutEdgeIterator it = rCur.first; it != rCur.second; ++it) {
				succs.push_back(boost::target(*it, _g));
			}

			// Check if vertex would add new targets - if yes, don't merge
			OutEdgeRange rPrev = boost::out_edges(prev, _g);
			for (OutEdgeIterator it = rPrev.first; it != rPrev.second; ++it) {
				GraphVertex target = boost::target(*it, _g);
				doMerge &= (std::find(succs.begin(), succs.end(), target) != succs.end() || target == cur);
			}

			if (doMerge) {
				gr = gr->_prev;
				merge(prev, cur);
				continue;
			}
		}
		gr = gr->_prev;
	}
}

const Graph &ControlFlow::analyze() {
	detectDoWhile();
	detectWhile();
	detectBreak();
	detectContinue();
	detectIf();
	detectElse();
	return _g;
}

void ControlFlow::detectWhile() {
	VertexRange vr = boost::vertices(_g);
	for (VertexIterator v = vr.first; v != vr.second; ++v) {
		GroupPtr gr = GET(*v);
		// Undetermined block that ends with conditional jump
		if (out_degree(*v, _g) == 2 && gr->_type == kNormalGroupType) {
			InEdgeRange ier = boost::in_edges(*v, _g);
			bool isWhile = false;
			for (InEdgeIterator e = ier.first; e != ier.second; ++e) {
				GroupPtr sourceGr = GET(boost::source(*e, _g));
				// Block has ingoing edge from block later in the code that isn't a do-while condition
				if ((*sourceGr->start_)->_address > (*gr->start_)->_address && sourceGr->_type != kDoWhileCondGroupType)
					isWhile = true;
			}
			if (isWhile)
				gr->_type = kWhileCondGroupType;
		}
	}
}

void ControlFlow::detectDoWhile() {
	VertexRange vr = boost::vertices(_g);
	for (VertexIterator v = vr.first; v != vr.second; ++v) {
		GroupPtr gr = GET(*v);
		// Undetermined block that ends with conditional jump...
		if (out_degree(*v, _g) == 2 && gr->_type == kNormalGroupType) {
			OutEdgeRange oer = boost::out_edges(*v, _g);
			for (OutEdgeIterator e = oer.first; e != oer.second; ++e) {
				GroupPtr targetGr = GET(boost::target(*e, _g));
				// ...to earlier in code
				if ((*targetGr->start_)->_address < (*gr->start_)->_address)
					gr->_type = kDoWhileCondGroupType;
			}
		}
	}
}

void ControlFlow::detectBreak() {
	VertexRange vr = boost::vertices(_g);
	for (VertexIterator v = vr.first; v != vr.second; ++v) {
		GroupPtr gr = GET(*v);
		// Undetermined block with unconditional jump...
		if (gr->_type == kNormalGroupType && ((*gr->end_)->IsUncondJump()) && out_degree(*v, _g) == 1) {
			OutEdgeIterator oe = boost::out_edges(*v, _g).first;
			GraphVertex target = boost::target(*oe, _g);
			GroupPtr targetGr = GET(target);
			// ...to somewhere later in the code...
			if ((*gr->start_)->_address >= (*targetGr->start_)->_address)
				continue;
			InEdgeRange ier = boost::in_edges(target, _g);
			for (InEdgeIterator ie = ier.first; ie != ier.second; ++ie) {
				GroupPtr sourceGr = GET(boost::source(*ie, _g));
				// ...to block immediately after a do-while condition, or to jump target of a while condition
				if ((targetGr->_prev == sourceGr && sourceGr->_type == kDoWhileCondGroupType) || sourceGr->_type == kWhileCondGroupType) {
					if (validateBreakOrContinue(gr, sourceGr))
						gr->_type = kBreakGroupType;
				}
			}
		}
	}
}

void ControlFlow::detectContinue() {
	VertexRange vr = boost::vertices(_g);
	for (VertexIterator v = vr.first; v != vr.second; ++v) {
		GroupPtr gr = GET(*v);
		// Undetermined block with unconditional jump...
		if (gr->_type == kNormalGroupType && ((*gr->end_)->IsUncondJump()) && out_degree(*v, _g) == 1) {
			OutEdgeIterator oe = boost::out_edges(*v, _g).first;
			GraphVertex target = boost::target(*oe, _g);
			GroupPtr targetGr = GET(target);
			// ...to a while or do-while condition...
			if (targetGr->_type == kWhileCondGroupType || targetGr->_type == kDoWhileCondGroupType) {
				bool isContinue = true;
				// ...unless...
				OutEdgeRange toer = boost::out_edges(target, _g);
				bool afterJumpTargets = true;
				for (OutEdgeIterator toe = toer.first; toe != toer.second; ++toe) {
					// ...it is targeting a while condition which jumps to the next sequential group
					if (targetGr->_type == kWhileCondGroupType && GET(boost::target(*toe, _g)) == gr->_next)
						isContinue = false;
					// ...or the instruction is placed after all jump targets from condition
					if ((*GET(boost::target(*toe, _g))->start_)->_address > (*gr->start_)->_address)
						afterJumpTargets = false;
				}
				if (afterJumpTargets)
					isContinue = false;

				if (isContinue && validateBreakOrContinue(gr, targetGr))
					gr->_type = kContinueGroupType;
			}
		}
	}
}

bool ControlFlow::validateBreakOrContinue(GroupPtr gr, GroupPtr condGr) {
	GroupPtr from, to, cursor;

	if (condGr->_type == kDoWhileCondGroupType) {
		to = condGr;
		from = gr;
	}	else {
		to = gr;
		from = condGr->_next;
	}

	GroupType ogt = (condGr->_type == kDoWhileCondGroupType ? kWhileCondGroupType : kDoWhileCondGroupType);
	// Verify that destination deals with innermost while/do-while
	for (cursor = from; cursor->_next != NULL && cursor != to; cursor = cursor->_next) {
		if (cursor->_type == condGr->_type) {
			OutEdgeRange oerValidate = boost::out_edges(find(cursor->start_), _g);
			for (OutEdgeIterator oeValidate = oerValidate.first; oeValidate != oerValidate.second; ++oeValidate) {
				GraphVertex vValidate = boost::target(*oeValidate, _g);
				GroupPtr gValidate = GET(vValidate);
				// For all other loops of same type found in range, all targets must fall within that range
				if ((*gValidate->start_)->_address < (*from->start_)->_address || (*gValidate->start_)->_address > (*to->start_)->_address )
					return false;

				InEdgeRange ierValidate = boost::in_edges(vValidate, _g);
				for (InEdgeIterator ieValidate = ierValidate.first; ieValidate != ierValidate.second; ++ieValidate) {
					GroupPtr igValidate = GET(boost::source(*ieValidate, _g));
					// All loops of other type going into range must be placed within range
					if (igValidate->_type == ogt && ((*igValidate->start_)->_address < (*from->start_)->_address || (*igValidate->start_)->_address > (*to->start_)->_address ))
					return false;
				}
			}
		}
	}
	return true;
}

void ControlFlow::detectIf() {
	VertexRange vr = boost::vertices(_g);
	for (VertexIterator v = vr.first; v != vr.second; ++v) {
		GroupPtr gr = GET(*v);
		// if: Undetermined block with conditional jump
		if (gr->_type == kNormalGroupType && ((*gr->end_)->isCondJump())) {
			gr->_type = kIfCondGroupType;
		}
	}
}

void ControlFlow::detectElse() {
	VertexRange vr = boost::vertices(_g);
	for (VertexIterator v = vr.first; v != vr.second; ++v) {
		GroupPtr gr = GET(*v);
		if (gr->_type == kIfCondGroupType) {
			OutEdgeRange oer = boost::out_edges(*v, _g);
			GraphVertex target;
			uint32 maxAddress = 0;
			GroupPtr targetGr;
			// Find jump target
			for (OutEdgeIterator oe = oer.first; oe != oer.second; ++oe) {
				targetGr = GET(boost::target(*oe, _g));
				if ((*targetGr->start_)->_address > maxAddress) {
					target = boost::target(*oe, _g);
					maxAddress = (*targetGr->start_)->_address;
				}
			}
			targetGr = GET(target);
			// else: Jump target of if immediately preceded by an unconditional jump...
			if (!(*targetGr->_prev->end_)->IsUncondJump())
				continue;
			// ...which is not a break or a continue...
			if (targetGr->_prev->_type == kContinueGroupType || targetGr->_prev->_type == kBreakGroupType)
				continue;
			// ...to later in the code
			OutEdgeIterator toe = boost::out_edges(find((*targetGr->_prev->start_)->_address), _g).first;
			GroupPtr targetTargetGr = GET(boost::target(*toe, _g));
			if ((*targetTargetGr->start_)->_address > (*targetGr->end_)->_address) {
				if (validateElseBlock(gr, targetGr, targetTargetGr)) {
					targetGr->_startElse = true;
					targetTargetGr->_prev->_endElse.push_back(targetGr.get());
				}
			}
		}
	}
}

bool ControlFlow::validateElseBlock(GroupPtr ifGroup, GroupPtr start, GroupPtr end) {
	for (GroupPtr cursor = start; cursor != end; cursor = cursor->_next) {
		if (cursor->_type == kIfCondGroupType || cursor->_type == kWhileCondGroupType || cursor->_type == kDoWhileCondGroupType) {
			// Validate outgoing edges of conditions
			OutEdgeRange oer = boost::out_edges(find(cursor->start_), _g);
			for (OutEdgeIterator oe = oer.first; oe != oer.second; ++oe) {
				GraphVertex target = boost::target(*oe, _g);
				GroupPtr targetGr = GET(target);
				// Each edge from condition must not leave the range [start, end]
				if ((*start->start_)->_address > (*targetGr->start_)->_address || (*targetGr->start_)->_address > (*end->start_)->_address)
					return false;
			}
		}

		// If previous group ends an else, that else must start inside the range
		for (ElseEndIterator it = cursor->_prev->_endElse.begin(); it != cursor->_prev->_endElse.end(); ++it)
		{
			if ((*(*it)->start_)->_address < (*start->start_)->_address)
				return false;
		}

		// Unless group is a simple unconditional jump...
		if ((*cursor->start_)->IsUncondJump())
			continue;

		// ...validate ingoing edges
		InEdgeRange ier = boost::in_edges(find(cursor->start_), _g);
		for (InEdgeIterator ie = ier.first; ie != ier.second; ++ie) {
			GraphVertex source = boost::source(*ie, _g);
			GroupPtr sourceGr = GET(source);

			// Edges going to conditions...
			if (sourceGr->_type == kIfCondGroupType || sourceGr->_type == kWhileCondGroupType || sourceGr->_type == kDoWhileCondGroupType) {
				// ...must not come from outside the range [start, end]...
				if ((*start->start_)->_address > (*sourceGr->start_)->_address || (*sourceGr->start_)->_address > (*end->start_)->_address) {
					// ...unless source is simple unconditional jump...
					if ((*sourceGr->start_)->IsUncondJump())
						continue;
					// ...or the edge is from the if condition associated with this else
					if (ifGroup == sourceGr)
						continue;
					return false;
				}
			}
		}
	}
	return true;
}