trace-as

#!/usr/bin/python3

import sys
import os
import subprocess

JUMP_TARGETS = set()


def set_jump(output):
    output.append('\tsub $128,%rsp\n')
    output.append('\tpush %rax\n')
    output.append('\tcall\t__trace_jump_set\n')
    output.append('\tpop  %rax\n')
    output.append('\tadd $128,%rsp\n')


def trace_jump(output):
    output.append('\tsub $128,%rsp\n')
    output.append('\tpush %rax\n')
    output.append('\tpush %rdi\n')
    output.append('\tpush %rsi\n')
    output.append('\tpush %rdx\n')
    output.append('\tpush %rcx\n')
    output.append('\tpush %r11\n')
    output.append('\tcall\t__trace_jump\n')
    output.append('\tpop  %r11\n')
    output.append('\tpop  %rcx\n')
    output.append('\tpop  %rdx\n')
    output.append('\tpop  %rsi\n')
    output.append('\tpop  %rdi\n')
    output.append('\tpop  %rax\n')
    output.append('\tadd $128,%rsp\n')


def collect_jump_targets(asm_file):
    entry_label = False
    compare_set = False
    file = []

    lines = asm_file.readlines()

    for i in range(len(lines)):
        line = lines[i]
        instruction = line[:-1]
        if entry_label:
            # Case 0: TraceJump the beginning of Main()
            JUMP_TARGETS.add(instruction[:-1])
            entry_label = False

        # The Entry
        if instruction == "main:":
            # instrument the beginning of Main()
            entry_label = True

        # If the next instruction is a conditional jump
        # TraceJumpSet before the current one:
        if (i + 1) < len(lines):
            next_instruction = lines[i + 1]
            if next_instruction.startswith("\t"):
                if (next_instruction[1] == "j") and (
                        next_instruction[2] != "m"):
                    set_jump(file)
                    compare_set = True

        file.append(line)

        # Check each instructions:
        if instruction.startswith("\t"):

            if (instruction[1] == "j") and (instruction[2] != "m"):
                assert compare_set
                compare_set = False
                # Case 1: TraceJump after a label of conditional jump
                JUMP_TARGETS.add(instruction.split("\t")[2])
                # Case 2: TraceJump after conditional jump
                trace_jump(file)
    return file


def instrument_jump_targets(intermediate):
    file = []
    for line in intermediate:
        file.append(line)
        instruction = line[:-2]
        if instruction in JUMP_TARGETS:
            trace_jump(file)

    return file


def instrument(asm, ins):
    asm_file = open(asm, 'rt')
    ins_file = open(ins, 'wt')
    inter = collect_jump_targets(asm_file)
    final = instrument_jump_targets(intermediate=inter)
    ins_file.writelines(final)


if __name__ == "__main__":
    args = sys.argv
    env = os.environ
    use64 = True
    doit = False

    # determine location for temporary files
    if 'TMPDIR' in env:
        tmp = env['TMPDIR']
    elif 'TEMP' in env:
        tmp = env['TEMP']
    elif 'TMP' in env:
        tmp = env['TMP']
    else:
        tmp = '/tmp'

    # determine which assembler to use next
    if 'TRACE_AS' in env:
        as_bin = env['TRACE_AS']
    else:
        as_bin = 'as'

    cmd = list()
    cmd.append(as_bin)

    # loop through all arguments, except last one (file name)
    for arg in args[1:-1]:
        if arg == '--64':
            use64 = True
        elif arg == '--32':
            use64 = False
        cmd.append(arg)

    obj = args[-1]
    instr = None
    if obj == '-version':
        cmd.append(obj)
    elif len(args) > 1:
        instr = "{}/traced-{}".format(tmp, os.path.basename(obj))
        cmd.append(instr)
        doit = True

    # instrument the assembly file
    if doit:
        assert instr
        instrument(obj, instr)

    # run the actual assembler program
    assert use64
    code = subprocess.call(cmd)

    # cleanup
    os.remove(instr)
    sys.exit(code)