# Given a path to llvm-objdump and a directory tree, spider the directory tree
# dumping every object file encountered with correct options needed to demangle
# symbols in the object file, and collect statistics about failed / crashed
# demanglings.  Useful for stress testing the demangler against a large corpus
# of inputs.

from __future__ import print_function

import argparse
import functools
import os
import re
import sys
import subprocess
import traceback
from multiprocessing import Pool
import multiprocessing

args = None

def parse_line(line):
    question = line.find('?')
    if question == -1:
        return None, None

    open_paren = line.find('(', question)
    if open_paren == -1:
        return None, None
    close_paren = line.rfind(')', open_paren)
    if open_paren == -1:
        return None, None
    mangled = line[question : open_paren]
    demangled = line[open_paren+1 : close_paren]
    return mangled.strip(), demangled.strip()

class Result(object):
    def __init__(self):
        self.crashed = []
        self.file = None
        self.nsymbols = 0
        self.errors = set()
        self.nfiles = 0

class MapContext(object):
    def __init__(self):
        self.rincomplete = None
        self.rcumulative = Result()
        self.pending_objs = []
        self.npending = 0

def process_file(path, objdump):
    r = Result()
    r.file = path

    popen_args = [objdump, '-t', '-demangle', path]
    p = subprocess.Popen(popen_args, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
    stdout, stderr = p.communicate()
    if p.returncode != 0:
        r.crashed = [r.file]
        return r

    output = stdout.decode('utf-8')

    for line in output.splitlines():
        mangled, demangled = parse_line(line)
        if mangled is None:
            continue
        r.nsymbols += 1
        if "invalid mangled name" in demangled:
            r.errors.add(mangled)
    return r

def add_results(r1, r2):
    r1.crashed.extend(r2.crashed)
    r1.errors.update(r2.errors)
    r1.nsymbols += r2.nsymbols
    r1.nfiles += r2.nfiles

def print_result_row(directory, result):
    print("[{0} files, {1} crashes, {2} errors, {3} symbols]: '{4}'".format(
        result.nfiles, len(result.crashed), len(result.errors), result.nsymbols, directory))

def process_one_chunk(pool, chunk_size, objdump, context):
    objs = []

    incomplete = False
    dir_results = {}
    ordered_dirs = []
    while context.npending > 0 and len(objs) < chunk_size:
        this_dir = context.pending_objs[0][0]
        ordered_dirs.append(this_dir)
        re = Result()
        if context.rincomplete is not None:
            re = context.rincomplete
            context.rincomplete = None

        dir_results[this_dir] = re
        re.file = this_dir

        nneeded = chunk_size - len(objs)
        objs_this_dir = context.pending_objs[0][1]
        navail = len(objs_this_dir)
        ntaken = min(nneeded, navail)
        objs.extend(objs_this_dir[0:ntaken])
        remaining_objs_this_dir = objs_this_dir[ntaken:]
        context.pending_objs[0] = (context.pending_objs[0][0], remaining_objs_this_dir)
        context.npending -= ntaken
        if ntaken == navail:
            context.pending_objs.pop(0)
        else:
            incomplete = True

        re.nfiles += ntaken

    assert(len(objs) == chunk_size or context.npending == 0)

    copier = functools.partial(process_file, objdump=objdump)
    mapped_results = list(pool.map(copier, objs))

    for mr in mapped_results:
        result_dir = os.path.dirname(mr.file)
        result_entry = dir_results[result_dir]
        add_results(result_entry, mr)

    # It's only possible that a single item is incomplete, and it has to be the
    # last item.
    if incomplete:
        context.rincomplete = dir_results[ordered_dirs[-1]]
        ordered_dirs.pop()

    # Now ordered_dirs contains a list of all directories which *did* complete.
    for c in ordered_dirs:
        re = dir_results[c]
        add_results(context.rcumulative, re)
        print_result_row(c, re)

def process_pending_files(pool, chunk_size, objdump, context):
    while context.npending >= chunk_size:
        process_one_chunk(pool, chunk_size, objdump, context)

def go():
    global args

    obj_dir = args.dir
    extensions = args.extensions.split(',')
    extensions = [x if x[0] == '.' else '.' + x for x in extensions]


    pool_size = 48
    pool = Pool(processes=pool_size)

    try:
        nfiles = 0
        context = MapContext()

        for root, dirs, files in os.walk(obj_dir):
            root = os.path.normpath(root)
            pending = []
            for f in files:
                file, ext = os.path.splitext(f)
                if not ext in extensions:
                    continue

                nfiles += 1
                full_path = os.path.join(root, f)
                full_path = os.path.normpath(full_path)
                pending.append(full_path)

            # If this directory had no object files, just print a default
            # status line and continue with the next dir
            if len(pending) == 0:
                print_result_row(root, Result())
                continue

            context.npending += len(pending)
            context.pending_objs.append((root, pending))
            # Drain the tasks, `pool_size` at a time, until we have less than
            # `pool_size` tasks remaining.
            process_pending_files(pool, pool_size, args.objdump, context)

        assert(context.npending < pool_size);
        process_one_chunk(pool, pool_size, args.objdump, context)

        total = context.rcumulative
        nfailed = len(total.errors)
        nsuccess = total.nsymbols - nfailed
        ncrashed = len(total.crashed)

        if (nfailed > 0):
            print("Failures:")
            for m in sorted(total.errors):
                print("  " + m)
        if (ncrashed > 0):
            print("Crashes:")
            for f in sorted(total.crashed):
                print("  " + f)
        print("Summary:")
        spct = float(nsuccess)/float(total.nsymbols)
        fpct = float(nfailed)/float(total.nsymbols)
        cpct = float(ncrashed)/float(nfiles)
        print("Processed {0} object files.".format(nfiles))
        print("{0}/{1} symbols successfully demangled ({2:.4%})".format(nsuccess, total.nsymbols, spct))
        print("{0} symbols could not be demangled ({1:.4%})".format(nfailed, fpct))
        print("{0} files crashed while demangling ({1:.4%})".format(ncrashed, cpct))
            
    except:
        traceback.print_exc()

    pool.close()
    pool.join()

if __name__ == "__main__":
    def_obj = 'obj' if sys.platform == 'win32' else 'o'

    parser = argparse.ArgumentParser(description='Demangle all symbols in a tree of object files, looking for failures.')
    parser.add_argument('dir', type=str, help='the root directory at which to start crawling')
    parser.add_argument('--objdump', type=str, help='path to llvm-objdump.  If not specified ' +
                        'the tool is located as if by `which llvm-objdump`.')
    parser.add_argument('--extensions', type=str, default=def_obj,
                        help='comma separated list of extensions to demangle (e.g. `o,obj`).  ' +
                        'By default this will be `obj` on Windows and `o` otherwise.')

    args = parser.parse_args()


    multiprocessing.freeze_support()
    go()