/r/dailyprogrammer challenge 290 "Blinking LEDs" JIT solution

leds.c

#include <stdio.h>
#include <string.h>
#include <stdint.h>
#include <stdlib.h>
#include <time.h>
#include <errno.h>

#include <sys/mman.h>
#include <unistd.h>

#define ARRAY_SIZE 8
typedef uint8_t led_array;

// led output functions

typedef led_array (*print_func_ptr)(const led_array);
typedef void (*comp_func_ptr)();
print_func_ptr print_func;

led_array print_leds(const led_array leds)
{
    static const char states[] = ".*";

    led_array mask = (led_array)1 << (8*sizeof(led_array) - 1);
    while (mask) {
        putchar(states[!!(leds & mask)]);
        mask >>= 1;
    }
    putchar('\n');
    return leds;
}


struct timespec last_frame;
uint64_t frame_time_ns;

led_array print_leds_and_wait(const led_array leds)
{
    static const char states[] = ".*";

    putchar('\r');
    led_array mask = (led_array)1 << (8*sizeof(led_array) - 1);
    while (mask) {
        putchar(states[!!(leds & mask)]);
        mask >>= 1;
    }
    fflush(stdout);

    uint64_t nsec = last_frame.tv_nsec;
    nsec += frame_time_ns;
    if (nsec > 1000000000ull) {
        last_frame.tv_sec += nsec / 1000000000ull;
        nsec %= 1000000000ull;
    }
    last_frame.tv_nsec = nsec;

    struct timespec wait_time;
    clock_gettime(CLOCK_REALTIME, &wait_time);

    int64_t next_ns = last_frame.tv_sec * 1000000000ll + last_frame.tv_nsec;
    int64_t wait_ns = wait_time.tv_sec * 1000000000ll + wait_time.tv_nsec;
    wait_ns = next_ns - wait_ns;

    wait_time.tv_nsec = wait_ns % 1000000000ll;
    wait_time.tv_sec = wait_ns / 1000000000ll;

    if(wait_ns > 0)
        nanosleep(&wait_time, NULL);

    return leds;
}








// compilation context

#define LIST_MIN_CAP 16
#define LABEL_MAX_LENGTH 32

struct label
{
    char name[LABEL_MAX_LENGTH];
    uint64_t target;
    uint64_t * refs;
    size_t ref_size;
    size_t ref_cap;
    int def_line;
};

struct asm_ctx
{
    uint8_t * buffer;
    size_t buffer_size;
    size_t write_pos;

    struct label * labels;
    size_t lbl_size;
    size_t lbl_cap;
};

size_t page_size;


struct asm_ctx* asm_ctx_create()
{
    struct asm_ctx * ctx = (struct asm_ctx*) malloc(sizeof(struct asm_ctx));
    if (!ctx) {
        fprintf(stderr, "Failed to allocate compilation context.\n");
        exit(1);
    }
    ctx->buffer = NULL;
    ctx->buffer_size = 0;
    ctx->write_pos = 0;
    ctx->labels = NULL;
    ctx->lbl_size = 0;
    ctx->lbl_cap = 0;

    page_size = sysconf(_SC_PAGESIZE);

    int prot = PROT_READ | PROT_WRITE;
    int flags = MAP_ANONYMOUS | MAP_PRIVATE;
    ctx->buffer = (uint8_t*) mmap(NULL, page_size, prot, flags, -1, 0);
    if (ctx->buffer == MAP_FAILED) {
        fprintf(stderr, "Failed to map code buffer.\n");
        exit(1);
    }
    ctx->buffer_size = page_size;

    ctx->labels = (struct label*) malloc(LIST_MIN_CAP * sizeof(struct label));
    if(!ctx->labels) {
        fprintf(stderr, "Failed to allocate label list.\n");
        exit(1);
    }
    ctx->lbl_cap = LIST_MIN_CAP;
    return ctx;
}

void label_init(struct label * lbl, const char * name, uint64_t target, int line)
{
    strncpy(lbl->name, name, LABEL_MAX_LENGTH - 1);
    lbl->target = target;
    lbl->def_line = line;
    lbl->refs = (uint64_t*) malloc(LIST_MIN_CAP * sizeof(uint64_t));
    if (!lbl->refs) {
        fprintf(stderr, "Failed to allocate reference list.\n");
        exit(1);
    }
    lbl->ref_size = 0;
    lbl->ref_cap = LIST_MIN_CAP;
}

void label_insert_reference(struct label * lbl, uint64_t ref)
{
    if (lbl->ref_size == lbl->ref_cap) {
        lbl->ref_cap += lbl->ref_cap >> 1;
        lbl->refs = realloc(lbl->refs, lbl->ref_cap);
        if (!lbl->ref_cap) {
            fprintf(stderr, "Failed to increase reference list size.\n");
            exit(1);
        }
    }
    lbl->refs[lbl->ref_size++] = ref;
}

void asm_ctx_insert_label(struct asm_ctx * ctx, const char * lbl_name, uint64_t lbl_addr, int line)
{
    if (ctx->lbl_size == ctx->lbl_cap) {
        ctx->lbl_cap += ctx->lbl_cap >> 1;
        ctx->labels = realloc(ctx->labels, ctx->lbl_cap);
        if (!ctx->labels) {
            fprintf(stderr, "Failed to increase label list size.\n");
            exit(1);
        }
    }
    label_init(&ctx->labels[ctx->lbl_size++], lbl_name, lbl_addr, line);
}

void label_release(struct label * lbl)
{
    free(lbl->refs);
}

void asm_ctx_release(struct asm_ctx * ctx)
{
    munmap(ctx->buffer, ctx->buffer_size);
    if (ctx->labels) {
        for(size_t i = 0; i < ctx->lbl_size; ++i) {
            label_release(&ctx->labels[i]);
        }
        free(ctx->labels);
        ctx->labels = NULL;
    }
    free(ctx);
}

void asm_ctx_release_labels(struct asm_ctx * ctx)
{
    if (ctx->labels) {
        for(size_t i = 0; i < ctx->lbl_size; ++i) {
            label_release(&ctx->labels[i]);
        }
        free(ctx->labels);
        ctx->labels = NULL;
    }
}

comp_func_ptr asm_ctx_finalize(struct asm_ctx * ctx)
{
    mprotect(ctx->buffer, ctx->buffer_size, PROT_READ | PROT_EXEC);
    return (comp_func_ptr) ctx->buffer+8;
}



void asm_ctx_ensure_buffer_size(struct asm_ctx * ctx, size_t buf_size)
{
    if (ctx->buffer_size < buf_size) {
        int prot = PROT_READ | PROT_WRITE;
        int flags = MAP_ANONYMOUS | MAP_PRIVATE;
        uint8_t * buf = mmap(NULL, ctx->buffer_size + page_size, prot, flags, -1, 0);
        if (buf == MAP_FAILED) {
            fprintf(stderr, "Failed to reallocate code buffer.\n");
            exit(1);
        }
        memcpy(buf, ctx->buffer, ctx->buffer_size);
        munmap(ctx->buffer, ctx->buffer_size);
        ctx->buffer_size += page_size;
        ctx->buffer = buf;
    }
}

#define MAX_LINE_LENGTH 1024
#define WHITESPACE " \t\n"

void asm_ctx_compile_stdin(struct asm_ctx * ctx)
{
    mprotect(ctx->buffer, ctx->buffer_size, PROT_READ | PROT_WRITE);
    ctx->write_pos = 0;

    char line[MAX_LINE_LENGTH];
    int line_nr = 0;

    // address of print_func
    *(uint64_t*) &ctx->buffer[0] = (uint64_t) print_func;
    ctx->write_pos += 8;
    // xor eax, eax
    ctx->buffer[ctx->write_pos++] = 0x31;
    ctx->buffer[ctx->write_pos++] = 0xc0;
    // push rbx
    ctx->buffer[ctx->write_pos++] = 0x53;
    // xor ebx, ebx
    ctx->buffer[ctx->write_pos++] = 0x31;
    ctx->buffer[ctx->write_pos++] = 0xdb;
    // push r12
    ctx->buffer[ctx->write_pos++] = 0x41;
    ctx->buffer[ctx->write_pos++] = 0x54;
    // mov r12, $print_func
    ctx->buffer[ctx->write_pos++] = 0x4c;
    ctx->buffer[ctx->write_pos++] = 0x8b;
    ctx->buffer[ctx->write_pos++] = 0x25;
    ctx->buffer[ctx->write_pos++] = 0xea;
    ctx->buffer[ctx->write_pos++] = 0xff;
    ctx->buffer[ctx->write_pos++] = 0xff;
    ctx->buffer[ctx->write_pos++] = 0xff;

    while (fgets(line, MAX_LINE_LENGTH, stdin)) {
        char * token = strtok(line, WHITESPACE);
        line_nr++;
        if (!token) {
            continue;
        }


        if (strcmp(token, "ld") == 0) {
            int reg = 0;
            uint64_t val = 0;
            token = strtok(NULL, WHITESPACE);
            if (!token) {
                fprintf(stderr, "%d: ld: missing operands\n", line_nr);
                exit(1);
            }

            char * op = token;
            token = strtok(NULL, WHITESPACE);
            if (token) {
                fprintf(stderr, "%d: ld: too many operands\n", line_nr);
                exit(1);
            }

            token = strtok(op, ",");
            if (!token) {
                fprintf(stderr, "%d: ld: cannot parse operands\n", line_nr);
                exit(1);
            }
            if (strcmp(token, "a") == 0) {
                reg = 1;
            } else if (strcmp(token, "b") == 0) {
                reg = 2;
            } else {
                fprintf(stderr, "%d: ld: invalid destination\n", line_nr);
                exit(1);
            }

            token = strtok(NULL, ",");
            if (!token) {
                fprintf(stderr, "%d: ld: missing second operand\n", line_nr);
                exit(1);
            }
            errno = 0;
            val = strtoull(token, NULL, 0);
            if (errno == ERANGE) {
                fprintf(stderr, "%d: ld: cannot parse second operand\n", line_nr);
                exit(1);
            }


#if ARRAY_SIZE == 8


            if (val >= (1 << 8)) {
                fprintf(stderr, "%d: ld: second operand out of range\n", line_nr);
                exit(1);
            }

            asm_ctx_ensure_buffer_size(ctx, ctx->write_pos + 4);

            ctx->buffer[ctx->write_pos++] = 0x66;
            ctx->buffer[ctx->write_pos++] = reg == 1 ? 0xb8 : 0xbb;
            ctx->buffer[ctx->write_pos++] = val;
            ctx->buffer[ctx->write_pos++] = 0;


#elif ARRAY_SIZE == 16


            if (val >= (1 << 16)) {
                fprintf(stderr, "%d: ld: second operand out of range\n", line_nr);
                exit(1);
            }

            asm_ctx_ensure_buffer_size(ctx, ctx->write_pos + 4);

            ctx->buffer[ctx->write_pos++] = 0x66;
            ctx->buffer[ctx->write_pos++] = reg == 1 ? 0xb8 : 0xbb;
            *(uint16_t*) &ctx->buffer[ctx->write_pos] = val;
            ctx->write_pos += 2;


#elif ARRAY_SIZE == 32


            if (val >= (1ull << 32)) {
                fprintf(stderr, "%d: ld: second operand out of range\n", line_nr);
                exit(1);
            }

            asm_ctx_ensure_buffer_size(ctx, ctx->write_pos + 5);

            ctx->buffer[ctx->write_pos++] = reg == 1 ? 0xb8 : 0xbb;
            *(uint32_t*) &ctx->buffer[ctx->write_pos] = val;
            ctx->write_pos += 4;


#endif


        } else if (strcmp(token, "out") == 0) {
            token = strtok(NULL, WHITESPACE);
            if (!token) {
                fprintf(stderr, "%d: out: missing operands\n", line_nr);
                exit(1);
            }
            if (strcmp(token, "(0),a") != 0) {
                fprintf(stderr, "%d: out: invalid operands\n", line_nr);
                exit(1);
            }
            token = strtok(NULL, WHITESPACE);
            if (token) {
                fprintf(stderr, "%d: out: too many operands\n", line_nr);
                exit(1);
            }

            // mov edi, eax
            ctx->buffer[ctx->write_pos++] = 0x89;
            ctx->buffer[ctx->write_pos++] = 0xc7;
            // call r12
            ctx->buffer[ctx->write_pos++] = 0x41;
            ctx->buffer[ctx->write_pos++] = 0xff;
            ctx->buffer[ctx->write_pos++] = 0xd4;


        } else if (strcmp(token, "rlca") == 0) {
            token = strtok(NULL, WHITESPACE);
            if (!token) {
#if ARRAY_SIZE == 8
            ctx->buffer[ctx->write_pos++] = 0xd0;
            ctx->buffer[ctx->write_pos++] = 0xc0;
#elif ARRAY_SIZE == 16
            ctx->buffer[ctx->write_pos++] = 0x66;
            ctx->buffer[ctx->write_pos++] = 0xd1;
            ctx->buffer[ctx->write_pos++] = 0xc0;
#elif ARRAY_SIZE == 32
            ctx->buffer[ctx->write_pos++] = 0xd1;
            ctx->buffer[ctx->write_pos++] = 0xc0;
#endif
            } else {
                uint64_t val;
                errno = 0;
                val = strtoull(token, NULL, 0);
                if (errno == ERANGE) {
                    fprintf(stderr, "%d: rlca: cannot parse operand\n", line_nr);
                    exit(1);
                }
                if (val >= sizeof(led_array) * 8) {
                    fprintf(stderr, "%d: rlca: operand too large\n", line_nr);
                    exit(1);
                }
#if ARRAY_SIZE == 8
                ctx->buffer[ctx->write_pos++] = 0xc0;
#elif ARRAY_SIZE == 16
                ctx->buffer[ctx->write_pos++] = 0x66;
                ctx->buffer[ctx->write_pos++] = 0xc1;
#elif ARRAY_SIZE == 32
                ctx->buffer[ctx->write_pos++] = 0xc1;
#endif
                ctx->buffer[ctx->write_pos++] = 0xc0;
                ctx->buffer[ctx->write_pos++] = val;
            }
        } else if (strcmp(token, "rrca") == 0) {
            token = strtok(NULL, WHITESPACE);
            if (!token) {
#if ARRAY_SIZE == 8
            ctx->buffer[ctx->write_pos++] = 0xd0;
#elif ARRAY_SIZE == 16
            ctx->buffer[ctx->write_pos++] = 0x66;
            ctx->buffer[ctx->write_pos++] = 0xd1;
#elif ARRAY_SIZE == 32
            ctx->buffer[ctx->write_pos++] = 0xd1;
#endif
            ctx->buffer[ctx->write_pos++] = 0xc8;
            } else {
                uint64_t val;
                errno = 0;
                val = strtoull(token, NULL, 0);
                if (errno == ERANGE) {
                    fprintf(stderr, "%d: rrca: cannot parse operand\n", line_nr);
                    exit(1);
                }
                if (val >= sizeof(led_array) * 8) {
                    fprintf(stderr, "%d: rrca: operand too large\n", line_nr);
                    exit(1);
                }
#if ARRAY_SIZE == 8
                ctx->buffer[ctx->write_pos++] = 0xc0;
#elif ARRAY_SIZE == 16
                ctx->buffer[ctx->write_pos++] = 0x66;
                ctx->buffer[ctx->write_pos++] = 0xc1;
#elif ARRAY_SIZE == 32
                ctx->buffer[ctx->write_pos++] = 0xc1;
#endif
                ctx->buffer[ctx->write_pos++] = 0xc8;
                ctx->buffer[ctx->write_pos++] = val;
            }
        } else if (strcmp(token, "djnz") == 0) {
            token = strtok(NULL, WHITESPACE);
            if (!token) {
                fprintf(stderr, "%d: djnz: missing operand\n", line_nr);
                exit(1);
            }

            // dec rbx
            ctx->buffer[ctx->write_pos++] = 0x48;
            ctx->buffer[ctx->write_pos++] = 0xff;
            ctx->buffer[ctx->write_pos++] = 0xcb;

            // jnz $label
            ctx->buffer[ctx->write_pos++] = 0x0f;
            ctx->buffer[ctx->write_pos++] = 0x85;
            ctx->buffer[ctx->write_pos++] = 0x00;
            ctx->buffer[ctx->write_pos++] = 0x00;
            ctx->buffer[ctx->write_pos++] = 0x00;
            ctx->buffer[ctx->write_pos++] = 0x00;


            size_t i;
            for (i = 0; i < ctx->lbl_size; ++i) {
                if (strcmp(ctx->labels[i].name, token) == 0) break;
            }
            if (i == ctx->lbl_size) {
                asm_ctx_insert_label(ctx, token, 0, -1);
            }
            label_insert_reference(&ctx->labels[i], (uint64_t)&ctx->buffer[ctx->write_pos]);
        } else {
            char * lbl = token;
            int len = strlen(lbl);
            if (lbl[len-1] != ':') {
                fprintf(stderr, "%d: unknown opcode\n", line_nr);
                exit(1);
            }
            lbl[len-1] = '\0';
            token = strtok(NULL, WHITESPACE);
            if (token) {
                fprintf(stderr, "%d: token after label\n", line_nr);
                exit(1);
            }

            size_t i;
            for (i = 0; i < ctx->lbl_size; ++i) {
                if (strcmp(ctx->labels[i].name, lbl) == 0) break;
            }
            if (i != ctx->lbl_size) {
                if (ctx->labels[i].def_line != -1) {
                    fprintf(stderr, "%d: duplicate label first declared on line %d\n", line_nr, ctx->labels[i].def_line);
                    exit(1);
                }
                ctx->labels[i].def_line = line_nr;
                ctx->labels[i].target = (uint64_t) &ctx->buffer[ctx->write_pos];
            } else {
                asm_ctx_insert_label(ctx, lbl, (uint64_t)&ctx->buffer[ctx->write_pos], line_nr);
            }
        }
    }
    // pop r12
    ctx->buffer[ctx->write_pos++] = 0x41;
    ctx->buffer[ctx->write_pos++] = 0x5c;
    // pop rbx
    ctx->buffer[ctx->write_pos++] = 0x5b;
    // ret
    ctx->buffer[ctx->write_pos++] = 0xc3;


    // resolve jumps
    for (size_t i = 0; i < ctx->lbl_size; ++i) {
        for (size_t j = 0; j < ctx->labels[i].ref_size; ++j) {
            uint32_t rel = ctx->labels[i].target - ctx->labels[i].refs[j];
            *(uint32_t*) (ctx->labels[i].refs[j] - 4) = rel;
        }
    }
}


int main()
{
    print_func = print_leds_and_wait;
    clock_gettime(CLOCK_REALTIME, &last_frame);
    frame_time_ns = 100 * 1000000ull;

    struct asm_ctx * ctx = asm_ctx_create();
    asm_ctx_compile_stdin(ctx);
    asm_ctx_release_labels(ctx);

    comp_func_ptr f = asm_ctx_finalize(ctx);
    f();

    asm_ctx_release(ctx);

    if (print_func == print_leds_and_wait) {
        putchar('\n');
    }

    return 0;
}