Skip to content

Instantly share code, notes, and snippets.

@Yangff

Yangff/trainJIT.cpp

Last active December 20, 2015 00:19
Show Gist options
  • Save Yangff/6040484 to your computer and use it in GitHub Desktop.
Save Yangff/6040484 to your computer and use it in GitHub Desktop.
Download ZIP
NOI2013 Day1 T3 JIT
Raw
trainJIT.cpp
#include <cstring>
#include <cstdio>
#include <cstdlib>
#include <cassert>
#include <map>
#include <cctype>
#ifdef _WIN32
#include <Windows.h>
#endif
using namespace std;
struct _asmcode{
unsigned char * head, * tail , *p;
int size;
_asmcode(){
size = 1024;
head = p = new unsigned char[size];
tail = head+size;
assert(p);
};
~_asmcode(){
size = 0;
delete [] head;
}
void large(){
unsigned char * newarea = new unsigned char[size*2];memset(newarea,0,sizeof(char) * size * 2);
assert(newarea);
memcpy(newarea,head,sizeof(unsigned char) * size);
p = (p-head) + newarea;
delete [] newarea;
head = newarea;
tail = head+size*2;
size *= 2;
}
int push8(unsigned char c){
if (p+1>tail) large();
int r = p-head;
*(p) = c;p++;
return r;
};
int push16(short s){
push8(s&0x00ff);
return push8(s>>8);
};
int push32(long i){
push8(i&0x000000ff);
push8((i&0x0000ff00) >> 8);
push8((i&0x00ff0000) >> 16);
return push8((i&0xff000000) >> 24);
};
void write8(int p,unsigned char c){
*(head+p) = c;
}
void write16(int p,short s){
write8(p,s&0x00ff);
write8(p+1,s>>8);
};
void write32(int p,int i){
write8(p,i&0x000000ff);
write8(p+1,(i&0x0000ff00) >> 8);
write8(p+2,(i&0x00ff0000) >> 16);
write8(p+3,(i&0xff000000) >> 24);
};
};
char getact(FILE *f){
int x;
while (isspace(x = getc(f)));
return x;
}
unsigned char* compile(FILE* f){
_asmcode code;
int * pos = NULL;
map<int,int> remark; // op code pos X should fill with pos[Y] - X + 4
map<int,int> varlst; // opcode pos X should fill content + Y
// Prepare done.
int n,m;
fscanf_s(f,"%d%d",&n,&m);
pos = new int[n+2];assert(!!pos); // oi code line to op code line
int length=n;
for (int i = 1 ; i <= n; i++){
char act = getact(f);
pos[i] = code.p-code.head;
if (act=='v'){
/*
A1 [var] mov eax,dword ptr ds:[var]
05 [IMM] add eax,IMM
or
2D [IMM] sub eax,IMM
A3 [var] mov dword ptr ds:[var],eax
A1 [var1] mov eax,dword ptr ds:[var1]
r32,r/m32
03[2B] 05->disp32 [var2] add eax,dword ptr ds:[var2]
or
A3 [var1] mov dword ptr ds:[var1],eax
*/
int controlnum = 0,just = 0;bool flag; // 0==v,1==c
bool action = 0;//0=+,1=-
fscanf_s(f,"%d",&controlnum); // vf0
action = getact(f)=='+'?0:1;
flag = getact(f)=='v'?0:1;
fscanf_s(f,"%d",&just);
code.push8(0xA1);//mov ...
assert(controlnum<=m);
varlst[code.p - code.head] = (controlnum) ; // offest
code.push32(0);
if (!action && flag) {
// + const
code.push8(0x05);
code.push32(just);
}
if (action && flag) {
// - const
code.push8(0x2D);
code.push32(just);
}
if (!action && !flag){
// + var
code.push8(0x03);
code.push8(0x05);
varlst[code.p - code.head] = (just) ; // offest
code.push32(0);
}
if (action && !flag){
// + var
code.push8(0x2B);
code.push8(0x05);
varlst[code.p - code.head] = (just) ; // offest
code.push32(0);
}
code.push8(0xA3);
varlst[code.p - code.head] = (controlnum) ; // offest
code.push32(0);
}
if (act=='s'){
int s1,s2;
fscanf_s(f,"%d%d",&s1,&s2);
if (s1 > n) s1 = n + 1;
if (s2 > n) s2 = n + 1;
/*C[i+1].a = 2;
fscanf(f,"%d%d",&C[i+1].f0,&C[i+1].f1);*/
/*
A1 98 33 14 01 mov eax,dword ptr ds:[01143398h]
40 inc eax
8A 48 FF mov cl,byte ptr [eax-1]
A3 98 33 14 01 mov dword ptr ds:[01143398h],eax
80 F9 31 cmp cl,31h
0F 84 B9 FF FF FF je apb (01141000h)
else bpa();
E9 D4 FF FF FF jmp bpa (01141020h)
*/
code.push8(0xA1);
varlst[code.p - code.head] = -1;
code.push32(0); // mov eax,dword ptr ds:[01143398h]
code.push8(0x40); // inc eax
code.push8(0x8A);
code.push8(0x48);
code.push8(0xFF); // mov cl,byte ptr [eax-1]
code.push8(0xA3);
varlst[code.p - code.head] = -1;
code.push32(0); // mov dword ptr ds:[01143398h],eax
code.push8(0x80);
code.push8(0xF9);
code.push8(0x31); // cmp cl,31h
code.push8(0x0F);
code.push8(0x84);
remark[code.p - code.head] = s1;
code.push32(0); // je s1
code.push8(0xE9);
remark[code.p - code.head] = s2;
code.push32(0); // jmp s2
}
if (act=='i'){
/*C[i+1].a = 3;
C[i+1].f0 = getact(f)=='v'?1:2;
fscanf(f,"%d",&C[i+1].arg0);
C[i+1].f1 = getact(f)=='v'?1:2;
fscanf(f,"%d",&C[i+1].arg1);
fscanf(f,"%d%d",&C[i+1].s1,&C[i+1].s2);*/
bool flag1,flag2;
int num1,num2,tar1,tar2;
flag1 = getact(f)=='v'?0:1;
fscanf_s(f,"%d",&num1);
flag2 = getact(f)=='v'?0:1;
fscanf_s(f,"%d",&num2);
assert(num1<=m || (flag1));
assert(num2<=m || (flag2));
fscanf_s(f,"%d%d",&tar1,&tar2);
if (flag1&&flag2){
int target = (num1<num2)?tar1:tar2;
code.push8(0xE9);
remark[code.p - code.head] = target;
code.push32(0);
}
if (( flag1 && !flag2 ) || (!flag1 && flag2)){
// c v v c
/*
if ((1<<30)<*x) apb();
010D1050 81 3D 70 33 0D 01 00 00 00 40 cmp dword ptr ds:[10D3370h],40000000h
010D105A 0F 8F A0 FF FF FF jg apb (010D1000h)
else bpa();
010D1060 E9 BB FF FF FF jmp bpa (010D1020h)
*/
/*
if (*x<(1<<30)) apb();
010D1030 81 3D 70 33 0D 01 00 00 00 40 cmp dword ptr ds:[10D3370h],40000000h
010D103A 0F 8C C0 FF FF FF jl apb (010D1000h)
else bpa();
010D1040 E9 DB FF FF FF jmp bpa (010D1020h)
*/
code.push8(0x81);code.push8(0x3D);
if (flag1) {
varlst[code.p - code.head] = (num2); // offest
code.push32(0);
code.push32(num1);
code.push8(0x0F); code.push8(0x8F);
} else {
varlst[code.p - code.head] = (num1); // offest
code.push32(0);
code.push32(num2);
code.push8(0x0F);code.push8(0x8C);
}
remark[code.p - code.head] = tar1;
code.push32(0);
code.push8(0xE9);
remark[code.p - code.head] = tar2;
code.push32(0);
}
if (!flag1 && !flag2){
// v v
/*
if (*x<*y) apb();
010D1070 A1 70 33 0D 01 mov eax,dword ptr ds:[010D3370h]
010D1075 3B 05 74 33 0D 01 cmp eax,dword ptr ds:[010D3374h]
010D107B 0F 8C 7F FF FF FF jl apb (010D1000h)
else bpa();
010D1081 E9 9A FF FF FF jmp bpa (010D1020h)
*/
code.push8(0xA1);
varlst[code.p - code.head] = (num1) ; // offest
code.push32(0);
code.push8(0x3B);
code.push8(0x05);
varlst[code.p - code.head] = (num2) ; // offest
code.push32(0);
code.push8(0x0F);
code.push8(0x8C);
remark[code.p - code.head] = tar1;
code.push32(0);
code.push8(0xE9);
remark[code.p - code.head] = tar2;
code.push32(0);
}
};
}
int stop = code.push8(0x90);
code.push8(0x90);
code.push8(0x90);// protect
// make pos
for (auto it = remark.begin();it!=remark.end();it++){
int X = it->first;
int Y = it->second;
if (Y>=n + 1)
code.write32(X,stop - X - 4);
else
code.write32(X,pos[Y] - X - 4);
};
int codelength = code.p - code.head;
int header = (varlst.size() * 2 + 2) * sizeof(int);
unsigned char *crst = new unsigned char[header + codelength];
unsigned char *p = crst;
int t = varlst.size();
memcpy(p,&t,sizeof(int));p+=sizeof(int);
for (auto it = varlst.begin();it!=varlst.end();it++)
{
t = it->first;
memcpy(p,&t,sizeof(int));p+=sizeof(int);
t = it->second;
memcpy(p,&t,sizeof(int));p+=sizeof(int);
}
t = codelength;
memcpy(p,&t,sizeof(int));p+=sizeof(int);
memcpy(p,code.head,codelength);
delete [] pos;
return crst;
};
int v[13];
int readint(unsigned char *p){
int x = 0;
memcpy(&x,p,sizeof(int));
return x;
}
void write8(unsigned char *head,int p,unsigned char c){
*(head+p) = c;
}
void write16(unsigned char *head,int p,short s){
write8(head,p,s&0x00ff);
write8(head,p+1,s>>8);
};
void write32(unsigned char *head,int p,int i){
write8(head,p,i&0x000000ff);
write8(head,p+1,(i&0x0000ff00) >> 8);
write8(head,p+2,(i&0x00ff0000) >> 16);
write8(head,p+3,(i&0xff000000) >> 24);
};
void run(unsigned char *code,int*var,char*vc){
/*
00B9C2E0 55 push ebp
00B9C2E1 8B EC mov ebp,esp
00B9C2E3 81 EC C0 00 00 00 sub esp,0C0h
00B9C2E9 53 push ebx
00B9C2EA 56 push esi
00B9C2EB 57 push edi
00B9C2EC 8D BD 40 FF FF FF lea edi,[ebp-0C0h]
00B9C2F2 B9 30 00 00 00 mov ecx,30h
00B9C2F7 B8 CC CC CC CC mov eax,0CCCCCCCCh
00B9C2FC F3 AB rep stos dword ptr es:[edi]
}
00B9C2FE 5F pop edi
00B9C2FF 5E pop esi
00B9C300 5B pop ebx
00B9C301 8B E5 mov esp,ebp
00B9C303 5D pop ebp
00B9C304 C3 ret
*/
map<int,int> varlst;
int size = readint(code);
code += sizeof(int);
for (int i = 0; i < size; i++){
int x = readint(code) + 28;code+=sizeof(int);
int y = readint(code);code+=sizeof(int);
varlst.insert(make_pair(x,y));
}
int codelength = readint(code);code += sizeof(int);
#ifndef _WIN32
unsigned char * space = (unsigned char*)malloc(28 + codelength + 8);
#else
//unsigned char * space = (unsigned char*)malloc(28 + codelength + 8);
unsigned char * space = (unsigned char*)VirtualAlloc(NULL,28 + codelength + 8,MEM_COMMIT,PAGE_EXECUTE_READWRITE) ;
#endif
unsigned char * p = space;
unsigned char st[28] = { 0x55,0x8b,0xec,0x81,0xec,0xc0,0x00,0x00,0x00,0x53,0x56,0x57,0x8d,0xbd,0x40,0xff,0xff,0xff,0xb8,0x30,0x00,0x00,0x00,0xb8,0xcc,0xcc,0xcc,0xcc};
memcpy(p,&st,sizeof(st));p+=28;
memcpy(p,code,codelength);p+=codelength;
unsigned char ed[7] = {0x5f,0x5e,0x5b,0x8b,0xe5,0x5d,0xc3};
memcpy(p,&ed,sizeof(ed));p+=7;*p = 0x90;
#ifndef _WIN32
mprotect(space, 28 + codelength + 8, PROT_READ|PROT_WRITE|PROT_EXEC);
#endif
for (auto it = varlst.begin();it!=varlst.end();it++){
int p = it->first;
if (it->second == -1){
write32(space,p,(int)(&vc));
} else {
write32(space,p,int(var + it->second));
}
}
typedef void*(eval_code)();
eval_code* E = (eval_code*)space;
E();
#ifndef _WIN32
free(space);
#else
VirtualFree(space,28 + codelength + 8,MEM_DECOMMIT);
#endif
};
void apb(){
int *x = v+0;
int *y = v+1;
// v[a]+c
(*x)+=1<<30;
// v[a]+v[b]
*x+=*y;
}
void bpa(){
v[1]=0;
}
char * p;
int *x = v+0;
int *y = v+1;
void jmp_test(){
if (*(p++)=='1') apb();
else bpa();
};
void jmp_test2_v_c(){
if (*x<(1<<30)) apb();
else bpa();
}
void jmp_test2_c_v(){
if ((1<<30)<*x) apb();
else bpa();
}
void jmp_test2_v_v(){
if (*x<*y) apb();
else bpa();
}
int main(){
//apb();
//p = new char[25];
//jmp_test();
//jmp_test2_v_c();
//jmp_test2_c_v();
//jmp_test2_v_v();
FILE * f = NULL;
fopen_s(&f,"train.in","r");
unsigned char* code = compile(f);
char* vcode = "111211";
run(code,v,vcode);
printf("%d\n",v[1]);
system("pause");
};
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment