#include "stdafx.h"

#include <Compiler.h>

#include "../BasicCompiler_Common/common.h"

#ifdef _AMD64_
#include "../BasicCompiler64/opcode.h"
#else
#include "../BasicCompiler32/opcode.h"
#endif

extern HANDLE hHeap;

/////////////////////////////////////////////
// 再配置スケジュール
//
// メモリの再配置は.textセクション、または.dataセクションに行われることが予想される。
//        .textセクション … セクションデータへのポインタ
//        .dataセクション … vtblの再配置
/////////////////////////////////////////////

CReloc *pobj_Reloc;

CReloc::CReloc(){
	buffer=0;
	length=0;
	NowPageRVAToReloc=0;
}
CReloc::~CReloc(){
	AllFree();
}
void CReloc::AllFree(void){
	if(buffer){
		HeapDefaultFree(buffer);
		buffer=0;
	}
}
void CReloc::copy(CReloc *por){
	AllFree();

	if(por->buffer){
		buffer=(char *)HeapReAlloc(hHeap,0,buffer,por->length);
		memcpy(buffer,por->buffer,por->length);
	}
	else buffer=0;

	length=por->length;

	NowPageRVAToReloc=por->NowPageRVAToReloc;
	NowCountAddrToReloc=por->NowCountAddrToReloc;

	codeSectionAddresses = por->codeSectionAddresses;

	dataSectionAddresses = por->dataSectionAddresses;
}

void CReloc::AddSchedule_CodeSection(DWORD addr){
	if( !compiler.IsDll() ) return;

	codeSectionAddresses.push_back( addr );
}
void CReloc::AddSchedule_DataSection(DWORD addr){
	if( !compiler.IsDll() ) return;

	dataSectionAddresses.push_back( addr );
}

void CReloc::__add(DWORD addr){
	if( !compiler.IsDll() ) return;

	BOOL sw;
	sw=0;
	while((addr-(addr%MEM_ALIGNMENT))>NowPageRVAToReloc){
		NowPageRVAToReloc+=MEM_ALIGNMENT;
		sw=1;
	}

	if(sw){
		//ページを増やす
		while(length%4){
			buffer=(char *)HeapReAlloc(hHeap,0,buffer,length+sizeof(WORD));
			(*(WORD *)(buffer+length))=0;
			length+=2;

			(*(DWORD *)(buffer+NowCountAddrToReloc))+=sizeof(WORD);
		}

		buffer=(char *)HeapReAlloc(hHeap,0,buffer,
			length+
			sizeof(DWORD)+
			sizeof(DWORD));

		//Page RVA
		*(DWORD *)(buffer+length)=NowPageRVAToReloc;
		length+=sizeof(DWORD);

		NowCountAddrToReloc=length;

		//Block size
		*(DWORD *)(buffer+length)=sizeof(DWORD)*2;
		length+=sizeof(DWORD);
	}

	//リロケーション情報の追加
	buffer=(char *)HeapReAlloc(hHeap,0,buffer,length+sizeof(WORD));
	buffer[length]=0;
	buffer[length+1]=0x30;	//=IMAGE_REL_BASED_HIGHLOW
	(*(WORD *)(buffer+length))|=addr&0x0FFF;
	length+=2;

	(*(DWORD *)(buffer+NowCountAddrToReloc))+=sizeof(WORD);
}
void CReloc::ResetRelocBuffer(void){
	if(buffer) HeapDefaultFree(buffer);

	buffer=(char *)HeapAlloc(hHeap,0,1);
	length=0;
	NowPageRVAToReloc=0;

	BOOST_FOREACH( DWORD addr, codeSectionAddresses )
	{
		extern int MemPos_CodeSection;
		__add(MemPos_CodeSection + addr);
	}

	BOOST_FOREACH( DWORD addr, dataSectionAddresses )
	{
		extern int MemPos_DataSection;
		__add(MemPos_DataSection + addr);
	}
}




//////////////////////////
// 一般スケジュール
//////////////////////////

CSchedule::CSchedule(){
	pObpValues=(int *)HeapAlloc(hHeap,0,1);
	num=0;

	this->flag=0;
}
CSchedule::~CSchedule(){
	HeapDefaultFree(pObpValues);
}
void CSchedule::SetFlag(int flag){
	this->flag=flag;
}
void CSchedule::add(){
	pObpValues=(int *)HeapReAlloc(hHeap,0,pObpValues,(num+1)*sizeof(int));
	extern int obp;
	pObpValues[num]=obp;
	num++;

	if(flag&SCHEDULE_FLAG_RELOC){
		//リロケーション情報を追加する
		pobj_Reloc->AddSchedule_CodeSection(obp);
	}
}
void CSchedule::move(int iStartPos,int iSize,int offset){
	int i;
	for(i=0;i<num;i++){
		if(iStartPos<=pObpValues[i] && pObpValues[i]<iStartPos+iSize){
			pObpValues[i]+=offset;
		}
	}
}

//データテーブルスケジュール
CSchedule *pobj_DataTableSchedule;

//グローバル変数アドレススケジュール
CSchedule *pobj_GlobalVarSchedule;



/////////////////////////////////////////
// インポート アドレス スケジューリング
/////////////////////////////////////////

CImportAddrSchedule::CImportAddrSchedule(){
	ppdi=(const DllProc **)HeapAlloc(hHeap,0,1);
}
CImportAddrSchedule::~CImportAddrSchedule(){
	HeapDefaultFree(ppdi);
}

void CImportAddrSchedule::add(const DllProc *pDllProc){
	ppdi=(const DllProc **)HeapReAlloc(hHeap,0,ppdi,(num+1)*sizeof(DllProc *));
	ppdi[num]=pDllProc;

	CSchedule::add();
}

CImportAddrSchedule *pobj_ImportAddrSchedule;




/////////////////////////////////////////
// プロシージャ アドレス スケジューリング
/////////////////////////////////////////

CSubAddrSchedule::CSubAddrSchedule(){
	ppsi=(const UserProc **)HeapAlloc(hHeap,0,1);
	pbCall=(BOOL *)HeapAlloc(hHeap,0,1);
}
CSubAddrSchedule::~CSubAddrSchedule(){
	HeapDefaultFree(ppsi);
	HeapDefaultFree(pbCall);
}

void CSubAddrSchedule::add(const UserProc *pUserProc,BOOL bCall){
	if(!pUserProc) return;

	ppsi=(const UserProc **)HeapReAlloc(hHeap,0,ppsi,(num+1)*sizeof(UserProc *));
	ppsi[num]=pUserProc;
	pbCall=(BOOL *)HeapReAlloc(hHeap,0,pbCall,(num+1)*sizeof(BOOL));
	pbCall[num]=bCall;

	CSchedule::add();

	if(!bCall){
		//リロケーション情報を追加する
		extern int obp;
		pobj_Reloc->AddSchedule_CodeSection(obp);
	}
}

CSubAddrSchedule *pobj_SubAddrSchedule;