#!/usr/bin/env python

import argparse
import re
import random
import string

def parse_memory_trace(input_file, stress_level, output_file):
    active_memory_requests = {}

    malloc_pattern = re.compile(r'dmmlib - m (0x\w+) (\d+)')
    realloc_pattern = re.compile(r'dmmlib - r (0x\w+) (0x\w+) (\d+)')
    memalign_pattern = re.compile(r'dmmlib - ma (0x\w+) (\d+) (\d+)')
    free_pattern   = re.compile(r'dmmlib - f (0x\w+)')

    for line in input_file:
        malloc_match = malloc_pattern.search(line)
        if malloc_match:
            if malloc_match.group(1) not in active_memory_requests:
                active_memory_requests[malloc_match.group(1)] = ''.join(random.choice(string.ascii_uppercase + string.digits) for x in range(6))
                variable_declaration = ''.join(['    void *var_', active_memory_requests[malloc_match.group(1)], ';\n'])
                output_file.write(variable_declaration)
            new_line = ''.join(['    var_', active_memory_requests[malloc_match.group(1)], ' = malloc(', malloc_match.group(2), ');\n'])
            output_file.write(new_line)
            if stress_level > 0:
                new_line = ''.join(['    var_', active_memory_requests[malloc_match.group(1)], ' = memset(var_', active_memory_requests[malloc_match.group(1)], ', 0, ', malloc_match.group(2), ');\n'])
                output_file.write(new_line)
            if stress_level > 1:
                new_line = ''.join(['    get_raw_blocks(&systemallocator);\n'])
                output_file.write(new_line)

        realloc_match = realloc_pattern.search(line)
        if realloc_match:
            if realloc_match.group(2) not in active_memory_requests:
                active_memory_requests[realloc_match.group(2)] = ''.join(random.choice(string.ascii_uppercase + string.digits) for x in range(6))
                variable_declaration = ''.join(['    void *var_', active_memory_requests[realloc_match.group(2)], ';\n'])
                output_file.write(variable_declaration)
            new_line = ''.join(['    var_',
                active_memory_requests[realloc_match.group(2)],
                ' = realloc(var_',
                active_memory_requests[realloc_match.group(1)], ', ',
                realloc_match.group(3), ');\n'])
            output_file.write(new_line)
            if stress_level > 0:
                new_line = ''.join(['    var_',
                    active_memory_requests[realloc_match.group(2)],
                    ' = memset(var_',
                    active_memory_requests[realloc_match.group(2)], ', 0, ',
                    realloc_match.group(3), ');\n'])
                output_file.write(new_line)
            if stress_level > 1:
                new_line = ''.join(['    get_raw_blocks(&systemallocator);\n'])
                output_file.write(new_line)

        memalign_match = memalign_pattern.search(line)
        if memalign_match:
            if memalign_match.group(1) not in active_memory_requests:
                active_memory_requests[memalign_match.group(1)] = ''.join(random.choice(string.ascii_uppercase + string.digits) for x in range(6))
                variable_declaration = ''.join(['    void *var_', active_memory_requests[memalign_match.group(1)], ';\n'])
                output_file.write(variable_declaration)
            new_line = ''.join(['    posix_memalign(&var_', active_memory_requests[memalign_match.group(1)], ', ', memalign_match.group(2), ', ', memalign_match.group(3), ');\n'])
            output_file.write(new_line)
            if stress_level > 0:
                new_line = ''.join(['    var_', active_memory_requests[memalign_match.group(1)], ' = memset(var_', active_memory_requests[memalign_match.group(1)], ', 0, ', memalign_match.group(3), ');\n'])
                output_file.write(new_line)
            if stress_level > 1:
                new_line = ''.join(['    get_raw_blocks(&systemallocator);\n'])
                output_file.write(new_line)

        free_match = free_pattern.search(line)
        if free_match:
            if free_match.group(1) in active_memory_requests:
                new_line = ''.join(['    free(var_', active_memory_requests[free_match.group(1)], ');\n'])
                del active_memory_requests[free_match.group(1)]
                output_file.write(new_line)
                if stress_level > 1:
                    new_line = ''.join(['    get_raw_blocks(&systemallocator);\n'])
                    output_file.write(new_line)
            else:
                print "Warning: free call without an actual allocation beforehands"

    input_file.close()

def main():
    parser = argparse.ArgumentParser(
        description='''Generate a C file of dynamic memory operations from a
            dmmlib trace''')
    parser.add_argument('tracefile', type=argparse.FileType('rt'),
        help='dmmlib trace from the application')
    parser.add_argument('-s', '--stress', metavar='stresslevel',
        type=int,
        help='stress level, 0 for none, 1 for data and 2 for data and metadata')
    parser.add_argument('-o', '--output', metavar='outputfile',
        type=argparse.FileType('wt'),
        help='generated C file, will use output.c if not available')

    try:
        args = parser.parse_args()
    except IOError, msg:
        parser.error(str(msg))

    if args.output is None:
        args.output = open('output.c', 'wt')

    if args.stress is None:
        args.stress = 0

    if args.stress < 0:
        args.stress = 0

    if args.stress > 2:
        args.stress = 2

    template_file = open('benchmark_template.c', 'rt')
    for line in template_file:
        if line == '$instructions\n':
            parse_memory_trace(args.tracefile, args.stress, args.output)
        else:
            args.output.write(line)

    args.output.close()

if __name__ == "__main__":
    main()