/* Skullfuck -- A dead simple compiler for Brainfuck.

   Copyright (C) 2017 Jakob Kreuze, All Rights Reserved.

   Skullfuck is free software: you can redistribute it and/or modify it
   under the terms of the GNU General Public License as published by the
   Free Software Foundation, either version 3 of the License, or (at
   your option) any later version.

   Skullfuck is distributed in the hope that it will be useful, but
   WITHOUT ANY WARRANTY; without even the implied warranty of
   MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the GNU
   General Public License for more details.

   You should have received a copy of the GNU General Public License
   along with Skullfuck. If not, see <http://www.gnu.org/licenses/>. */

#include <sys/types.h>
#include <sys/wait.h>
#include <stdbool.h>
#include <stdio.h>
#include <stdlib.h>
#include <unistd.h>

#include <getopt.h>


#if defined(__x86_64__)
#define ARCH         "x86_64"
#define OS           "Linux"
#define INC_DATA_PTR "\tincq %rsi\n"
#define ADD_DATA_PTR "\taddq $0x%x, %%rsi\n"
#define DEC_DATA_PTR "\tdecq %rsi\n"
#define SUB_DATA_PTR "\tsubq $0x%x, %%rsi\n"
#define INC_DATA     "\tincb (%rsi)\n"
#define ADD_DATA     "\taddb $0x%x, (%%rsi)\n"
#define DEC_DATA     "\tdecb (%rsi)\n"
#define SUB_DATA     "\tsubb $0x%x, (%%rsi)\n"
#define ZERO_DATA    "\tmovb $0x00, (%rsi)\n"
#define WRITE_DATA   "\tmovq $0x01, %rax\n" \
                     "\tmovq $0x01, %rdi\n" \
                     "\tsyscall\n"
#define READ_DATA    "\tmovq $0x00, %rax\n" \
                     "\tmovq $0x00, %rdi\n" \
                     "\tsyscall\n"
#define DATA_IS_ZERO "\tmovb (%rsi), %cl\n" \
                     "\ttestb %cl, %cl\n"
#define JMP_END      "\tjz E%d\n"
#define JMP_START    "\tjnz S%d\n"

#define PRELUDE      "\t.section .bss\n" \
                     "\t.comm mem, %d\n" \
                     "\t.section .text\n" \
                     "\t.globl _start\n" \
                     "_start:\n" \
                     "\tmovq $mem, %%rsi\n" \
                     "\tmovq $0x01, %%rdx\n"

#define EXIT         "\tmovq $0x3c, %rax\n" \
                     "\tmovq $0x00, %rdi\n" \
                     "\tsyscall\n"
#elif defined(__i686__)
#define ARCH         "i686"
#define OS           "Linux"
#define INC_DATA_PTR "\tincl %ecx\n"
#define ADD_DATA_PTR "\taddl $0x%x, %%ecx\n"
#define DEC_DATA_PTR "\tdecl %ecx\n"
#define SUB_DATA_PTR "\tsubl $0x%x, %%ecx\n"
#define INC_DATA     "\tincb (%ecx)\n"
#define ADD_DATA     "\taddb $0x%x, (%%ecx)\n"
#define DEC_DATA     "\tdecb (%ecx)\n"
#define SUB_DATA     "\tsubb $0x%x, (%%ecx)\n"
#define ZERO_DATA    "\tmovb $0x00, (%ecx)\n"
#define WRITE_DATA   "\tmovl $0x04, %eax\n" \
                     "\tmovl $0x01, %ebx\n" \
                     "\tint $0x80\n"
#define READ_DATA    "\tmovl $0x03, %eax\n" \
                     "\tmovl $0x00, %ebx\n" \
                     "\tint $0x80\n"
#define DATA_IS_ZERO "\tmovb (%ecx), %al\n" \
                     "\ttestb %al, %al\n"
#define JMP_END      "\tjz E%d\n"
#define JMP_START    "\tjnz S%d\n"

#define PRELUDE      "\t.section .bss\n" \
                     "\t.comm mem, %d\n" \
                     "\t.section .text\n" \
                     "\t.globl _start\n" \
                     "_start:\n" \
                     "\tmovl $mem, %%ecx\n" \
                     "\tmovl $0x01, %%edx\n"

#define EXIT         "\tmovl $0x01, %eax\n" \
                     "\tmovl $0x00, %ebx\n" \
                     "\tint $0x80\n"
#endif

#define MAX_CELLS      30000
#define MAX_LOOP_DEPTH 64

#define SKULLFUCK_VERSION "0.3.0"

#define EXIT_SUCCESS 0
#define EXIT_FAILURE 1



/* Structure for storing runtime parameters. */
struct params {
    bool compile_only;
    char *in, *out, *as;
};


/* A stack data structure for storing loop depth. */
struct loop_stack {
    int _mem[MAX_LOOP_DEPTH];
    int _stack_index;
};


/* Pushes `n` onto the stack specified by `s`. */
static void push_loop_index(struct loop_stack *s, int n) {
    s->_mem[s->_stack_index] = n;
    s->_stack_index++;
}


/* Pops the top of `s` off and returns it. */
static int pop_loop_index(struct loop_stack *s) {
    s->_stack_index--;
    return s->_mem[s->_stack_index];
}


/* Loop optimizations. */
enum {
    NONE,
    LOOP_EMPTY,
    LOOP_ZERO,
};


/* Optimization subroutine which detects common loop patterns, such as
   `[-]`, increments the `in` pointer to beyond them, and returns the
   identity of the pattern. */
static int optimize_loops(char **in) {
    if (*(*in + 1) == '-' && *(*in + 2) == ']') {
        *in += 2;
        return LOOP_ZERO;
    }
    if (*(*in + 1) == ']') {
        *in += 1;
        return LOOP_EMPTY;
    }
    return NONE;
}


/* Minor optimization subroutine. Increments the `in` pointer to beyond
   the collection of adjacent `op` characters, and returns the number
   of operations counted. */
static int reduce(char **in, char op) {
    int i;
    for (i = 0; (*in)[i] == op; i++);
    *in += i - 1;
    return i;
}


/* Converts the brainfuck instructions read from `in` to assembly
   instructions and writes them to `out`. `in` is expected to be
   null-terminated. */
static void write_instructions(FILE *out, char *in) {
    int               next_loop = 0, cur_loop = 0, op_count, optimization;
    struct loop_stack *loops = calloc(sizeof(struct loop_stack), 1);

    for (char *cur = in; *cur != '\0'; cur++) {
        switch(*cur) {
            case '>':
                op_count = reduce(&cur, '>');
                if (op_count > 1) {
                    fprintf(out, ADD_DATA_PTR, op_count);
                } else {
                    fputs(INC_DATA_PTR, out);
                }
                break;
            case '<':
                op_count = reduce(&cur, '<');
                if (op_count > 1) {
                    fprintf(out, SUB_DATA_PTR, op_count);
                } else {
                    fputs(DEC_DATA_PTR, out);
                }
                break;
            case '+':
                op_count = reduce(&cur, '+');
                if (op_count > 1) {
                    fprintf(out, ADD_DATA, op_count);
                } else {
                    fputs(INC_DATA, out);
                }
                break;
            case '-':
                op_count = reduce(&cur, '-');
                if (op_count > 1) {
                    fprintf(out, SUB_DATA, op_count);
                } else {
                    fputs(DEC_DATA, out);
                }
                break;
            case '.':
                fputs(WRITE_DATA, out);
                break;
            case ',':
                fputs(READ_DATA, out);
                break;
            case '[':
                optimization = optimize_loops(&cur);
                if (optimization == LOOP_ZERO) {
                    fputs(ZERO_DATA, out);
                } else if (optimization == NONE) {
                    cur_loop = next_loop;
                    push_loop_index(loops, next_loop++);
                    fprintf(out, "S%d:\n", cur_loop);
                    fputs(DATA_IS_ZERO, out);
                    fprintf(out, JMP_END, cur_loop);
                }
                break;
            case ']':
                cur_loop = pop_loop_index(loops);
                fprintf(out, "E%d:\n", cur_loop);
                fputs(DATA_IS_ZERO, out);
                fprintf(out, JMP_START, cur_loop);
                break;
        }
    }

    fputs(EXIT, out);
    free(loops);
}


/* Prints error information and terminates with minimal cleanup. */
static void panic(char *msg) {
    perror(msg);
    fprintf(stderr, "Fatal error encountered. Terminating!\n");
    exit(EXIT_FAILURE);
}


/* Attempts to create a new file at `path`, and writes a boilerplate
   prelude to initialize memory segments and load some registers with
   constant values. A FILE pointer will be returned, or NULL if the file
   could not be created. */
static FILE *out_init(char *path) {
    FILE *fp = fopen(path, "w+");
    if (fp == NULL) return fp;
    fprintf(fp, PRELUDE, MAX_CELLS);
    return fp;
}


/* Returns a buffer containing the contents of `path`. */
static char *read_infile(char *path) {
    FILE *fp = fopen(path, "r");
    if (fp == NULL) panic(path);
    fseek(fp, 0, SEEK_END);
    size_t len = ftell(fp);
    char   *in = malloc(len);
    fseek(fp, 0, SEEK_SET);
    fread(in, len, 1, fp);
    fclose(fp);
    return in;
}


/* Assembles and links the compiled source code into a binary. */
static void create_binary(struct params p) {
    pid_t        pid;
    char  *const as_argv[] = {"as", "-o", "/tmp/skullfuck_tmp.o", p.as, NULL};
    char  *const ld_argv[] = {"ld", "-o", p.out, "/tmp/skullfuck_tmp.o", NULL};

    if ((pid = fork()) == -1) {
        panic("fork");
    } else if (pid == 0) {
        execv(AS, as_argv);
    } else {
        wait(NULL);
        if ((pid = fork()) == -1) {
            panic("fork");
        } else if (pid == 0) {
            execv(LD, ld_argv);
        } else {
            wait(NULL);
        }
    }

    if (!p.compile_only)
        remove("/tmp/skullfuck_tmp.s");
    remove("/tmp/skullfuck_tmp.o");
}


/* Parses command-line arguments into a params structure. */
static struct params parse_args(int argc, char **argv) {
    struct params p = {0};

    if (argc < 2) {
        fprintf(stderr, "USAGE: %s [options] (infile)\n", argv[0]);
        exit(EXIT_FAILURE);
    }

    int cur = 0, opt_index = 0, count = 0;
    static struct option long_opts[] = {
        {"help", no_argument, NULL, 'h'},
        {"version", no_argument, NULL, 'v'},
        {NULL, 0, NULL, 0}
    };

    do {
        count++;
        cur = getopt_long(argc, argv, "hvSo:", long_opts, &opt_index);
        switch (cur) {
            case 'h':
                printf("A dead simple, non-optimizing Brainfuck compiler.\n\n"
                       "\t-h, --help\tDisplay this help page and exit.\n"
                       "\t-v, --version\tDisplay versioning information and "
                       "exit\n\n"
                       "\t-S\t\tStop after the stage of compilation proper; "
                       "do not assemble.\n"
                       "\t-o file\t\tPlace the output in file.\n");
                exit(EXIT_SUCCESS);
            case 'v':
                printf("Skullfuck compiler for Brainfuck, version %s.\n"
                       "Toolchain for %s %s.\nDeveloped by Jakob. "
                       "<http://jakob.space>\n", SKULLFUCK_VERSION, ARCH, OS);
                exit(EXIT_SUCCESS);
            case 'S':
                p.compile_only = true;
                break;
            case 'o':
                count++;
                p.out = optarg;
                break;
        }
    } while (cur >= 0);

    if (count >= argc) {
        fprintf(stderr, "Fatal error: No input files.\n");
        exit(EXIT_FAILURE);
    }

    p.in = argv[count];
    p.as = "/tmp/skullfuck_tmp.s";
    if (p.compile_only) {
        p.out = p.out == NULL ? "./out.s" : p.out;
        p.as = p.out;
    } else if (p.out == NULL) {
        p.out = "./a.out";
    }

    return p;
}


int main(int argc, char **argv) {
    struct params p = parse_args(argc, argv);
    FILE          *fp = out_init(p.as);

    if (fp == NULL) panic(p.as);
    char *in = read_infile(p.in);
    write_instructions(fp, in);
    fclose(fp);

    if (!p.compile_only)
        create_binary(p);

    free(in);
    return EXIT_SUCCESS;
}