jit-x64-opt.dasc

#include <stdint.h>
#include <assert.h>
#include "util.h"
#include "jit-x64-opt-inc.h"

|.arch x64
|.actionlist actions
|
|// Use rbx as our cell pointer.
|// Since rbx is a callee-save register, it will be preserved
|// across our calls to getchar and putchar.
|.define PTR, rbx
|
|// Macro for calling a function.
|// In cases where our target is <=2**32 away we can use
|//   | call &addr
|// But since we don't know if it will be, we use this safe
|// sequence instead.
|.macro callp, addr
|  mov64  rax, (uintptr_t)addr
|  call   rax
|.endmacro

#define Dst &state
#define MAX_NESTING 256
#define MAX_DATA 256

dasm_State *state;
unsigned int maxpc = 0;
int pcstack[MAX_NESTING];
int *top = pcstack, *limit = pcstack + MAX_NESTING;
char *p;

int main(int argc, char *argv[])
{
	if (argc < 2) err("Usage: jit-x64 <inputfile>");
	initjit(&state, actions);

	// Function prologue.
	|  push PTR
	|  mov  PTR, rdi      // rdi store 1st argument

	p = read_file(argv[1]);
	compile();

	// Function epilogue.
	|  pop  PTR
	|  ret

	void (*fptr)(char*) = jitcode(&state);
	char *mem = calloc(30000, 1);
	fptr(mem);
	free(mem);
	free_jitcode(fptr);
	return 0;
}

DataState *create(int offset)
{
    DataState *tmp = (DataState *) malloc(sizeof(DataState));
    tmp->offset = offset;
    tmp->diff = tmp->totalDiff = 0;
    tmp->predict = 1;
    return tmp;
}

DataState *find(int offset, DataState *data[], int sizeOfData)
{
    for (int i = 0; i < sizeOfData; ++i)
    {
        if(data[i]->offset == offset)
            return data[i];
    }

    return NULL;
}

void inc(int offset, DataState *(*data)[], int *sizeOfData)
{
    DataState *tmp = find(offset, *data, *sizeOfData);
    if (tmp == NULL)
        tmp = (*data)[(*sizeOfData)++] = create(offset);

    ++(tmp->diff);
    ++(tmp->totalDiff);
}

void dec(int offset, DataState *(*data)[], int *sizeOfData)
{
    DataState *tmp = find(offset, *data, *sizeOfData);
    if (tmp == NULL)
        tmp = (*data)[(*sizeOfData)++] = create(offset);

    --(tmp->diff);
    --(tmp->totalDiff);
}

void movPtr(int from, int to) {
    int def = from - to;
    if (def < 0)
        | add  PTR, -def
    else if (def > 0)
        | sub  PTR, def
}

void genLoops(int *ptr, DataState *(*data)[], int *sizeOfData, DataState *cond)
{
    movPtr(0, cond->offset);
    int curPtr = cond->offset;

    if (cond->totalDiff != -1)
    {
        | mov  al, byte [PTR]
        | mov  ah, 0
        | mov  cl, -cond->totalDiff
        | idiv cl
        | mov  cl, al
        
    }    
    else
        | mov  cl, byte [PTR]

    for (int i = 0; i < *sizeOfData; ++i)
    {
        DataState *dataState = (*data)[i];
        if (dataState == cond)
        {
            movPtr(curPtr, dataState->offset);
            curPtr = dataState->offset;
            | mov  byte [PTR], 0

        }
        else if (dataState->diff != 0)
        {
            movPtr(curPtr, dataState->offset);
            curPtr = dataState->offset;
            | mov al, (dataState->diff)
            | imul cl
            | add  byte [PTR], al
        }

        if (dataState != NULL)
        {
            free(dataState);
            (*data)[i] = NULL;
        }
    }

    movPtr(curPtr, *ptr);
    *ptr = 0;
    *sizeOfData = 0;
}

void genAll(int *ptr, DataState *(*data)[], int *sizeOfData)
{
    int curPtr = 0;
    for (int i = 0; i < *sizeOfData; ++i)
    {
        DataState *dataState = (*data)[i];
        if (dataState->diff != 0)
        {
            movPtr(curPtr, dataState->offset);
            curPtr = dataState->offset;
                
            if (dataState->diff > 0)
                | add  byte [PTR], dataState->diff
            else
                | sub  byte [PTR], -dataState->diff
        }
        free((*data)[i]);
        (*data)[i] = NULL;
    }

    movPtr(curPtr, *ptr);
    *ptr = 0;
    *sizeOfData = 0;
}

void getPtr(int ptr, DataState *(*data)[], int *sizeOfData)
{
    movPtr(0, ptr);

    DataState *tmp = find(ptr, *data, *sizeOfData);
    if (tmp == NULL)
        tmp = (*data)[(*sizeOfData)++] = create(ptr);

    tmp->diff = tmp->totalDiff = tmp->predict = 0;

    |  callp getchar
	|  mov   byte [PTR], al
    movPtr(ptr, 0);
}

void putPtr(int ptr, DataState *(*data)[], int *sizeOfData)
{
    movPtr(0, ptr);

    DataState *tmp = find(ptr, *data, *sizeOfData);
    if (tmp != NULL && tmp->diff != 0)
    {
        if (tmp->diff > 0)
            | add  byte [PTR], tmp->diff
        else
            | sub  byte [PTR], -tmp->diff

        tmp->diff = 0;
    }

    |  movzx edi, byte [PTR]
    |  callp putchar
    movPtr(ptr, 0);
}

void compile()
{
    int predict = 1;
    int ptr = 0;
    DataState *data[MAX_DATA], *tmp;
    int sizeOfData = 0;

	for (; *p; p++) 
    {

		switch (*p) {
		case '>':
            ++ptr;
            break;

		case '<':
            --ptr;
            break;

		case '+':
            inc(ptr, &data, &sizeOfData);
            break;

		case '-':
            dec(ptr, &data, &sizeOfData);
            break;

		case '.':
            putPtr(ptr, &data, &sizeOfData);
            predict = 0;
			break;

		case ',':
            getPtr(ptr, &data, &sizeOfData);
            predict = 0;
			break;

		case '[':
            genAll(&ptr, &data, &sizeOfData);
            predict = 0;

            maxpc += 2;
            *top = maxpc;
            top++;
            dasm_growpc(&state, maxpc);
            |  cmp  byte [PTR], 0
            |  je   =>(maxpc-2)
            |=>(maxpc-1):

            (*p)++;
            compile();
			break;

		case ']':
            (*p)++;

            if (predict && ptr == 0)
            {
                tmp = find(ptr, data, sizeOfData);
                if (tmp != NULL && tmp->predict && tmp->totalDiff < 0)
                {
                    genLoops(&ptr, &data, &sizeOfData, tmp);
                    top--;
                    |=>(*top-2):
                    return;
                }
            }

            genAll(&ptr, &data, &sizeOfData);
            top--;
            |  cmp  byte [PTR], 0
            |  jne  =>(*top-1)
            |=>(*top-2):
			return ;
		}
	}

    genAll(&ptr, &data, &sizeOfData);

    return ;
}