/************************************************************************
Function length calculation algorithm - by Darawk:

1.  Scan the function's code for branches, and record each branch.  Stop
	upon reaching an end-point*.  This group of instructions constitutes
	the current "block".
2.  QSort the branch list
3.  Recursively repeat steps 1 & 2 with each branch, skipping duplicates
	and intra-block branches.

	*end-point: A ret instruction or an unconditional backwards jump,
				that jumps to a previous block.
************************************************************************/
u32 GetFunctionLength(void *begin)
{
	void *end = GetFunctionEnd(begin);
	u32 delta = (u32)((DWORD_PTR)end - (DWORD_PTR)begin);
	delta    += mlde32(end);
	return delta;
}

void *GetFunctionEnd(void *func)
{
	void *block = func;
	vector<void *> branchList;
	//	ptr now points to the end of this block
	void *blockend = GetBranchListFromBlock(block, branchList);

	//	If there are no branches, then return
	//	the empty list.  If we don't have this
	//	here the loop will crash on an empty
	//	branch list.
	if(branchList.size() == 0) return blockend;

	//	Sort the list so that we can identify and
	//	discard, intra-block branches.  And optimize
	//	the removal of duplicates.
	std::sort(branchList.begin(), branchList.end());

	void *prev = NULL;
	vector<void *>::iterator branch;
	for(branch = branchList.begin(); branch != branchList.end(); branch++)
	{
		//	Skip branches that jump into a block we've already
		//	processed.  
		if(*branch < blockend || *branch == prev)
			continue;

		blockend = GetFunctionEnd(*branch);
		prev     = *branch;
	}

	return blockend;
}

void *GetBranchListFromBlock(void *block, vector<void *> &branchList)
{
	u8 *ptr = (u8 *)block;

	//	If we reach an end-point, then this block is complete
	while(!IsEndPoint(ptr, block))
	{
		//	Record all branching instructions that we encounter
		void *address = GetBranchAddress(ptr);
		if(address)
		{
			branchList.push_back(address);
		}

		//	Next instruction
		ptr += mlde32(ptr);
	}
	
	return ptr;
}


void *GetBranchAddress(u8 *instr)
{
	s32 offset = 0;
	//	This code will determine what type of branch it is, and
	//	determine the address it will branch to.
	switch(*instr)
	{
	case INSTR_SHORTJMP:
	case INSTR_RELJCX:
		offset  = (s32)(*(s8 *)(instr + 1));
		offset += 2;	
		break;
	case INSTR_RELJMP:
		offset  = *(s32 *)(instr + 1);
		offset += 5;
		break;
	case INSTR_NEAR_PREFIX:
		if(*(instr + 1) >= INSTR_NEARJCC_BEGIN && *(instr + 1) <= INSTR_NEARJCC_END)
		{
			offset  = *(s32 *)(instr + 2);
			offset += 5;
		}
		break;
	default:
		//	Check to see if it's in the valid range of JCC values.
		//	e.g. ja, je, jne, jb, etc..
		if(*instr >= INSTR_SHORTJCC_BEGIN && *instr <= INSTR_SHORTJCC_END)
		{
			offset  = (s32)*((s8 *)(instr + 1));
			offset += 2;
		}
		break;
	}

	if(offset == 0) return NULL;
	return instr + offset;
}

bool IsEndPoint(u8 *instr, void *curblock)
{
	void *address;
	s32 offset;
	switch(*instr)
	{
	case INSTR_RET:
	case INSTR_RETN:
	case INSTR_RETFN:
	case INSTR_RETF:
		return true;
		break;

	//	The following two checks, look for an instance in which
	//	an unconditional jump returns us to a previous block,
	//	thus creating a pseudo-endpoint.
	case INSTR_SHORTJMP:
		offset  = (s32)(*(s8 *)(instr + 1));
		address = instr + offset;
		if(address <= curblock) return true;
		break;
	case INSTR_RELJMP:
		offset  = *(s32 *)(instr + 1);
		address = instr + offset;
        if(address <= curblock) return true;
		break;
	default:
		return false;
		break;
	}

	return false;
}