#!/usr/bin/env python3
#
# (c) Mit authors 1994-2020
#
# The package is distributed under the MIT/X11 License.
#
# THIS PROGRAM IS PROVIDED AS IS, WITH NO WARRANTY. USE IS AT THE USER’S
# RISK.

import argparse, json, textwrap

from code_util import Code, copyright_banner, disable_warnings
from code_gen import run_body, run_fn
from spec import opcode_bit, Instructions as BasicInstructions
from specializer import CacheState, gen_case
from specializer_spec import Instructions
from path import Path, State


GENERATOR_PROGRAM = 'gen-specializer'
PURPOSE = 'Run the virtual machine main loop (optimizing version).'
COPYRIGHT_YEARS = '1994-2019'


# Command-line arguments.
parser = argparse.ArgumentParser(
    prog=GENERATOR_PROGRAM,
    description='Generate optimizing interpreter',
)
parser.add_argument(
    '--no-preguess',
    action='store_true',
    help='disable preguessing',
)
parser.add_argument(
    'labels_filename',
    metavar='LABELS-FILENAME',
    help='labels file to use',
)
args = parser.parse_args()


def _opcode_test(multiguess):
    '''
    Returns a C expression representing the condition "`ir` begins with
    `multiguess`".
     - multiguess - list of Instructions.
    '''
    assert type(multiguess) is list
    opcode_value = 0 # Bit pattern to match.
    opcode_bits = 0 # Number of bits to match.
    for instruction in multiguess:
        assert opcode_bits is not None
        opcode_value |= instruction.opcode << opcode_bits
        opcode_bits += opcode_bit
        if instruction.action.terminal:
            if instruction.opcode & 0x80 != 0:
                opcode_value |= -1 << opcode_bits
            opcode_bits = None # Must match all bits.
    # Generate the guess condition.
    if opcode_bits is None:
        opcode_mask = (0 if opcode_value & 80 == 0 else -1) << opcode_bit
        return f'ir == {opcode_value | opcode_mask:#x}'
    else:
        mask = (1 << opcode_bits) - 1
        return f'(ir & {mask:#x}) == {opcode_value:#x}'


class Label:
    '''
     - index - int - the index of this Label.
     - path - Path - the instruction sequence just executed.
     - preguess - list of Instructions - the longest unique probable
       continuation.
     - guess - Instructions - the most probable next instruction.
     - if_correct - int - the Label index to jump to if `guess` is correct,
       or `None` for the fallback label.
     - if_wrong - int - the Label index to jump to if `guess` is wrong,
       or `None` for the fallback label.
    '''
    def __init__(
        self,
        index,
        path,
        preguess,
        guess,
        if_correct,
        if_wrong,
    ):
        '''
         - path - str (space-separated Instructions names).
         - preguess - str (space-separated Instructions names).
         - guess - str (Instructions name).
        '''
        self.index = index
        self.path = Path(tuple(Instructions[name] for name in path.split()))
        self.preguess = [Instructions[name] for name in preguess.split()]
        self.guess = Instructions[guess]
        self.if_correct = if_correct
        self.if_wrong = if_wrong

    def name(self):
        '''Returns the C identifier of this Label.'''
        return f'A_{self.index}'

    def cache_state(self):
        '''Returns a fresh CacheState describing this Label.'''
        return CacheState(self.cached_depth(), self.checked_depth())

    def cached_depth(self):
        return self.path.state.cached_depth()

    def checked_depth(self):
        return self.path.state.checked_depth()

    def max_cached_depth(self):
        return self.path.state.step(self.guess).max_cached_depth

    def generate_code(self, profiling=False, no_preguess=False):
        '''
        Returns a Code to place at this Label.
         - profiling - bool - `True` to increment `state_guess_correct` when
           executing the instruction.
        Returns a Code.
        '''
        # Generate the Code for the branch where `self.guess` is correct.
        c_code = Code()
        if profiling:
            c_code.append(f'state_guess_correct[{self.index}]++;')
        state = State()
        cache_state = self.cache_state()
        state = state.step(self.guess)
        c_code.extend(Code('''\
            ir = ARSHIFT(ir, 8);''',
            f'{{ // {self.guess.name}',
                gen_case(self.guess, cache_state),
            '}',
        ))
        c_code.extend(goto_label_by_index(cache_state, self.if_correct))
        pops = -state.stack_min
        pushes = state.stack_max - state.stack_min
        # Generate the code for the branch where the guess is wrong.
        w_code = Code()
        if profiling:
            w_code.append(f'state_guess_wrong[{self.index}]++;')
        cache_state = self.cache_state()
        w_code.extend(goto_label_by_index(cache_state, self.if_wrong))
        # Compute the bit string of Instructions opcodes to guess.
        multiguess = [self.guess]
        if not no_preguess and self.if_correct is not None:
            multiguess.extend(labels[self.if_correct].preguess)
        opcode_test = _opcode_test(multiguess)
        # Compute the guard condition.
        cache_state = self.cache_state()
        guard_code = self.guess.action.guard.format(
            stack_0=cache_state.lvalue(0),
            stack_1=cache_state.lvalue(1),
        )
        tests = [
            f'({opcode_test})',
            f'likely({cache_state.underflow_test(pops)})',
            f'likely({cache_state.overflow_test(pops, pushes)})',
            f'({guard_code})',
        ]
        # Generate the main Code.
        code = Code(
            '// History: {}'.format(' '.join(i.name for i in self.path)),
            '// Future: {}'.format(' '.join(i.name for i in self.preguess)),
            'assert(error == MIT_ERROR_OK);',
            f'assert(cached_depth == {self.cached_depth()});',
            'if ({}) {{'.format(' && '.join(tests)),
            c_code,
            '}',
        )
        code.extend(w_code)
        return code


# Load the labels file, which tells us what control-flow graph to generate.
with open(args.labels_filename, "rb") as h:
    labels = [
        Label(
            index,
            obj['path'],
            obj['preguess'],
            obj['guess'],
            obj['if_correct'],
            obj['if_wrong'],
        )
        for index, obj
            in enumerate(json.load(h))
    ]


# The peak depth to which we cache stack slots in C locals.
max_cached_depth = max(label.max_cached_depth() for label in labels)


def goto_label_by_index(cache_state, index):
    '''
    Returns a Code representing a jump to the specified Label.
     - index - int - the index of the Label, or `None` for `FALLBACK`.
    '''
    code = Code()
    if index is None:
        # Jump to the fallback label.
        code.extend(cache_state.flush(CacheState(0, 0)))
        code.append('goto A_FALLBACK;')
    else:
        # Jump to a Label.
        label = labels[index]
        code.extend(cache_state.flush(label.cache_state()))
        code.append(f'goto {label.name()};')
    return code


def gen_labels_code(profiling=False):
    code = Code()
    code.append(Code('''\
        // The number of stack items cached in C variables.
        int cached_depth = 0;

        // Each A_XXX label has an expectation of what instruction sequence
        // probably comes next. If correct, there is a fast path for executing
        // that sequence, otherwise we fall back on a slow path with broader
        // expectations. Ultimately, we fall back to `A_FALLBACK` which
        // handles corner cases.
        //
        // The calling convention at each A_XXX label is as follows:
        //
        //  - On entry, `error` is `MIT_ERROR_OK`. `cached_depth` is a
        //    compile-time constant indicating how many top stack items are in
        //    C variables.
        //  - Before running the `Instructions.code` for an instruction,
        //    right-shift `ir`.
        //  - When an error occurs, call `THROW(code)`.
        //    It sets `error` to the return code; then it resets `pc` and
        //    `ir` to their initial values, and does "goto error".
        //  - On exit, whether by `THROW` or otherwise, all Mit
        //    registers must be set according to the Mit spec. In particular,
        //    there must be no side effects before a THROW, except on `ir` and
        //    `pc`, which will be reset as described above.
    '''))
    for label in labels:
        code.append('')
        code.append(f'{label.name()}:')
        code.append(label.generate_code(profiling, args.no_preguess or profiling))
    code.append('')
    code.append('''\
        A_FALLBACK:
            assert(error == MIT_ERROR_OK);
            assert(cached_depth == 0);
            uint8_t opcode = (uint8_t)ir;
            ir = ARSHIFT(ir, 8);
            #define run_inner run_inner_fast
    ''')
    code.append(run_body(BasicInstructions))
    code.append('''
            #undef run_inner
            continue;

        error:
            switch (cached_depth) {
    ''')
    switch_code = Code()
    cache_state = CacheState(max_cached_depth, 0)
    while cache_state.cached_depth > 0:
        switch_code.append(f'case {cache_state.cached_depth}:')
        case_code = Code()
        case_code.extend(cache_state.flush(cache_state.cached_depth - 1))
        case_code.append('// Falls through.')
        switch_code.append(case_code)
    switch_code.append('''\
        case 0:
            THROW_LONGJMP(error);
            break;
        default:
            assert(0); // Unreachable.
        }
    ''')
    code.append(switch_code)
    return disable_warnings(
        ['-Wunsafe-loop-optimizations'],
        code,
    )

def gen_body_code(profiling=False):
    code = Code()
    code.append('''\
        mit_uword_t stack_words = mit_stack_words;
        mit_word_t error = MIT_ERROR_OK;
    ''')
    if max_cached_depth > 0:
        cache_state = CacheState(max_cached_depth, 0)
        # Initialize the cache variables to avoid (incorrect) uninitialized
        # variable warning.
        code.append('mit_word_t {};'.format(
            ', '.join(
                f'{cache_state.var(pos)} = 0'
                for pos in reversed(range(max_cached_depth))
            )
        ))
    code.append('for (;;) {')
    code.extend(gen_labels_code(profiling=profiling))
    code.append('''\
        }
    ''')
    return code


# Generate the code for the entire source file.
code = copyright_banner(GENERATOR_PROGRAM, PURPOSE, COPYRIGHT_YEARS)
code.append(f'''

    #include "config.h"

    #include <assert.h>
    #include <stdio.h>
    #include <unistd.h>

    #include "mit/mit.h"
    #include "mit/features.h"

    #include "run.h"


    #define NUM_LABELS {len(labels)}
    // The number of times we guessed correctly at each specializer label.
    static MIT_THREAD_LOCAL unsigned long long state_guess_correct[NUM_LABELS];
    // The number of times we guessed wrong at each specializer label.
    static MIT_THREAD_LOCAL unsigned long long state_guess_wrong[NUM_LABELS];

    // Extra data to be included in profile files.
    static const struct label {{
        const char *path;
        const char *guess;
        int correct_label;
        int wrong_label;
    }} label_data[NUM_LABELS] = {{'''
)
def label_to_c(l):
    return -1 if l is None else l
code.append(Code(',\n'.join([
    ('{{"{}", "{}", {}, {}}}'.format(
        ' '.join(i.name for i in l.path),
        l.guess.name,
        label_to_c(l.if_correct),
        label_to_c(l.if_wrong),
    ))
    for l in labels
])))
code.append('};')
code.append('''
    void mit_profile_reset(void)
    {
        for (unsigned i = 0; i < NUM_LABELS; i++)
            state_guess_correct[i] = state_guess_wrong[i] = 0;
    }

    int mit_profile_dump(int fd)
    {
        // Open output stream (for buffering)
        int dup_fd = dup(fd);
        if (dup_fd == -1)
            return -1;
        FILE *fp = fdopen(dup_fd, "w");
        if (fp != NULL) {
            if (fprintf(fp, "[") < 0)
                goto err;
            const char *sep = "";
            for (unsigned i = 0; i < NUM_LABELS; i++) {
                struct label l = label_data[i];
                if (fprintf(fp,
                        "%s\\n  {"
                            "\\"path\\": \\"%s\\", "
                            "\\"guess\\": \\"%s\\", "
                            "\\"if_correct\\": %d, "
                            "\\"if_wrong\\": %d, "
                            "\\"correct_count\\": %llu, "
                            "\\"wrong_count\\": %llu"
                        "}",
                        sep,
                        l.path,
                        l.guess,
                        l.correct_label,
                        l.wrong_label,
                        state_guess_correct[i],
                        state_guess_wrong[i]
                    ) < 0
                )
                    goto err;
                sep = ",";
            }
            if (fprintf(fp, "\\n]\\n") < 0)
                goto err;
            if (fclose(fp) == 0)
                return 0;
            fp = NULL;
        }
    err:
        if (fp != NULL)
            fclose(fp);
        return -1;
    }
''')

for profiling, fn in [
    (False, 'fast'),
    (True, 'profile'),
]:
    code.append('')
    code.extend(disable_warnings(
        ['-Wstack-protector', '-Wvla-larger-than='], # Stack protection cannot cope with VLAs.
        Code(
            f'static void run_inner_{fn}(mit_word_t *pc, mit_word_t ir, mit_word_t * restrict stack, mit_uword_t * restrict stack_depth_ptr, jmp_buf *jmp_buf_ptr) {{',
            '#define stack_depth (*stack_depth_ptr)',
                gen_body_code(profiling),
            '''\
            #undef stack_depth
            }''',
        )
    ))
    code.extend(run_fn(fn))

print(code)
