bfjit.cpp

#define _CRT_SECURE_NO_WARNINGS

#include <Windows.h>
#include <stdio.h>
#include <string.h>

typedef unsigned char uchar;

#define JMPSTACK_SIZE 1024
#define ADDRBUF_SIZE 1024
#define TRI(a, b, c) ((a)<<6|(b)<<3|(c))

uchar* __fastcall write32(uchar *dst, unsigned int src) {
  *dst++ = src >> 0;
	*dst++ = src >> 8;
	*dst++ = src >> 16;
	*dst++ = src >> 24;
	return dst;
}

unsigned int __fastcall read32(uchar *src) {
	return src[0] << 0
		|  src[1] << 8
		|  src[2] << 16
		|  src[3] << 24;
}

void __fastcall mput(char c)
{
	putchar(c);
}

char __fastcall mget()
{
	return getchar();
}

int main(int argc, char** argv) {

	if (argc == 1) {
		fprintf(stderr, "Usage: bfjit [in.b] <out>\n");
		fprintf(stderr, "\t-h: Dump the generated x86 hex to out\n");
		fprintf(stderr, "\t-e: Treat the argument as a brainfuck program\n");
		fprintf(stderr, "\t    instead of a filename\n");
		exit(-1);
	}

	// Allocate executable memory for the code
	uchar *native = (uchar*)VirtualAlloc(NULL, 2048*4, MEM_COMMIT | MEM_RESERVE, PAGE_EXECUTE_READWRITE);
	uchar *out = native;

	memset(native, 0, 2048*4);

	char *infile = nullptr, *outfile = nullptr;

	int hexdump = 0;
	int argcode = 0;

	// Parse arguments
	for (char **argp = argv + 1; argp != argv + argc; argp++) {
		char *arg = *argp;
		if (arg[0] == '-') {
			for (char *f = &arg[1]; *f; f++) {
				switch (*f) {
				case 'h':
					hexdump = 1;
					break;
				case 'e':
					argcode = 1;
					break;
				default:
					fprintf(stderr, "Unknown flag '%c'\n", *f); 
					break;
				}
			}
		} else {
			if (infile == nullptr)
				infile = arg;
			else if (outfile == nullptr)
				outfile = arg;
			else {
				fprintf(stderr, "Unexpected argument\n");
			}
		}
	}

	if (!infile) {
		fprintf(stderr, "Input file not specified\n");
		exit(-1);
	}

	// Source reading
	FILE* in = 0;
	if (!argcode) {
		in = fopen(infile, "r");
		if (!in) {
			fprintf(stderr, "Failed to open input file\n");
			exit(-1);
		}
	}
	char c;

	// Create a stack to store the loop begin addresses
	uchar **loopbuf = (uchar**)malloc(sizeof(uchar*) * JMPSTACK_SIZE), **loopstack = loopbuf + JMPSTACK_SIZE;

	uchar **addrbuf = (uchar**)malloc(sizeof(uchar*) * ADDRBUF_SIZE), **addrhead = addrbuf;

	memset(addrbuf, 0, sizeof(uchar*) * ADDRBUF_SIZE);

	// Save the state
	*out++ = 0x60; // PUSHAD

	// Reset EAX
	*out++ = 0x31; // XOR
	*out++ = TRI(0x3, 0x0, 0x0); // EAX, EAX

	// For +-<>
	int counter = 0;
	int mem;

	// Peeking
	int hasPeek = 0;
	char peek;

	char *acptr = infile;

	unsigned int addr;

	while (argcode ? *acptr : !feof(in)) {
		c = hasPeek ? hasPeek = 0, peek : (argcode ? *acptr++ : fgetc(in));
		mem = 0;
		switch (c) {
			
		case '.':
			*out++ = 0x60; // PUSHAD
			*out++ = 0x8A; // MOV
			*out++ = TRI(0x0, 0x1, 0x4); // ECX, (SIB no disp)
			*out++ = TRI(0x0, 0x0, 0x1); // [ECX + 1*EAX]
			*out++ = 0xE8; // CALL
			addr = (unsigned int)(&mput) - ((unsigned int)(out) + 4);
			*addrhead++ = out;
			out = write32(out, addr);
			*out++ = 0x61; // POPAD
			break;
		case ',':
			*out++ = 0x8D; // LEA
			*out++ = TRI(0x0, 0x3, 0x4); // EBX, (SIB no disp)
			*out++ = TRI(0x0, 0x0, 0x1); // [ECX + 1*EAX]
			*out++ = 0x60; // PUSHAD
			*out++ = 0xE8; // CALL
			addr = (unsigned int)(&mget) - ((unsigned int)(out) + 4);
			*addrhead++ = out;
			out = write32(out, addr);
			*out++ = 0x88; // MOV
			*out++ = TRI(0x0, 0x0, 0x3); // [EBX], AL
			*out++ = 0x61; // POPAD
			break;
			
		case '+':
		case '-':
			mem = 1;
		case '<':
		case '>':
			// Optimize by combining multiple operations
			counter++;

			hasPeek = 0;
			
			// Look ahead the next meaningful character
			while ((argcode ? *acptr : !feof(in)) && !hasPeek)
			{
				switch (peek = argcode ? *acptr++ : fgetc(in)) {
				case '-':
				case '+':
				case ',':
				case '.':
				case '<':
				case '>':
				case '[':
				case ']':
					hasPeek = 1;
					break;
				default:
					break;
				}
			}

			// Emit the combined instruction
			if (c != peek) {
				uchar rrm = 0;
				if (mem) counter &= 0xFF;
				int islong = counter <= 0x7F ? 0 : 2;

				// Increment or decrement
				if (counter == 1) {
					*out++ = mem ? 0xfe : 0xff; // INC (8 bit mem else 32 bit)
					rrm = (mem ? c == '-' : c == '<') ? 0x01 : 0x00;
				} else {
					*out++ = 0x80 | (mem ? 0 : islong) | mem; // IMM
					rrm = (mem ? c == '-' : c == '<') ? 0x05 : 0x00;
				}

				if (mem) {
					*out++ = TRI(0x0, rrm, 0x4); // SIB
					*out++ = TRI(0x0, 0x0, 0x1); // [ECX + 1*EAX]
				} else {
					*out++ = TRI(0x3, rrm, 0x0); // EAX
				}
				if (counter != 1) {
					if (islong || mem) {
						out = write32(out, counter);
					} else {
						*out++ = (uchar)counter;
					}
				}
				if (!mem) {
					*out++ = 0x21; // AND
					*out++ = TRI(0x3, 0x2, 0x0); // EAX EDX
				}
				counter = 0;
			}
			break;
		case '[':
			// Compare the memory cell with zero
			*out++ = 0x80;               // IMM-instruction
			*out++ = TRI(0x0, 0x7, 0x4); // CMP (SIB)
			*out++ = TRI(0x0, 0x0, 0x1); // [ECX + 1*EAX]
			*out++ = 0x00;               // 0

			// Push the location to write the jump address
			*--loopstack = out;

			// Create space for the opcode
			out += 2;
			break;
		case ']': {
			// Pop the jump slot and calculate the distance
			uchar *loopstart = *loopstack++;
			int dist = out - loopstart;

			// If the jump is short (fits in a byte) optimize
			if (dist < 0x7F - 2) {

				// [: Conditional short jump over the body
				*loopstart++ = 0x74; // JZ
				*loopstart++ = (uchar)dist;

				// ]: Uncondtitional short jump to the test
				*out++ = 0xEB; // JMP (short)
				*out++ = (uchar)-(signed char)(dist + 2 + 4);

			} else {
				// Reserve memory for 32-bit jump offset
				memmove(loopstart + (4), loopstart, dist);
				out += (4);

				// Move the call offsets
				for (uchar** a = addrbuf; *a; a++) {
					if (*a > loopstart) {
						*a += 4;
						unsigned int val = read32(*a);
						write32(*a, val - 4);
					}
				}

				// [: Conditional long jump over the body
				*loopstart++ = 0x0F; // EXT opcode
				*loopstart++ = 0x84; // JZ
				loopstart = write32(loopstart, dist - 2 + 1 + 4);

				// ]: Unconditional long jump to the test
				*out++ = 0xE9; // JMP (long)
				out = write32(out, -(dist - 2 + 1 + 4 + 2 + 4 + 4));
			}
			} break;
		default:
			break;
		}
	}

	// Load the state
	*out++ = 0x61; // POPAD

	// Return
	*out++ = 0xC3; // RET

	if (hexdump) {
		if (!outfile) {
			fprintf(stderr, "Output file not specified\n");
			exit(-1);
		}

		FILE *hexout = fopen(outfile, "w");

		if (!hexout) {
			fprintf(stderr, "Failed to open output file\n");
			exit(-1);
		}

		int accum = 0;
		for (uchar *n = native; n != out; n++) {
			fprintf(hexout, "%02x ", *n);
			if (accum++ % 16 == 15)
				fprintf(hexout, "\n");
		}

		fclose(hexout);
	}
	else {
		char *memory = (char*)malloc(0x100);
		memset(memory, 0x0, 0x100);

		((void (__fastcall*)(char*, unsigned int))native)(memory, 0xFF);

		free(memory);
	}
	if (!argcode)
		fclose(in);

	free(loopbuf);
	free(addrbuf);
	VirtualFree(native, 0, MEM_RELEASE);
}