path: root/tools/klee-stats/klee-stats

#!/usr/bin/env python
# -*- encoding: utf-8 -*-
"""Output statistics logged by Klee."""

# use '/' to mean true division and '//' to mean floor division
from __future__ import division
from __future__ import print_function

import os
import re
import sys
import argparse

from operator import itemgetter
try:
    from tabulate import TableFormat, Line, DataRow, tabulate
except:
    print('Error: Package "tabulate" required for table formatting. '
          'Please install it using "pip" or your package manager.',
          file=sys.stderr)
    exit(1)

Legend = [
    ('Instrs', 'number of executed instructions'),
    ('Time', 'total wall time (s)'),
    ('TUser', 'total user time'),
    ('ICov', 'instruction coverage in the LLVM bitcode (%)'),
    ('BCov', 'branch coverage in the LLVM bitcode (%)'),
    ('ICount', 'total static instructions in the LLVM bitcode'),
    ('TSolver', 'time spent in the constraint solver'),
    ('States', 'number of currently active states'),
    ('Mem', 'megabytes of memory currently used'),
    ('Queries', 'number of queries issued to STP'),
    ('AvgQC', 'average number of query constructs per query'),
    ('Tcex', 'time spent in the counterexample caching code'),
    ('Tfork', 'time spent forking'),
    ('TResolve', 'time spent in object resolution'),
]

KleeTable = TableFormat(lineabove=Line("-", "-", "-", "-"),
                        linebelowheader=Line("-", "-", "-", "-"),
                        linebetweenrows=None,
                        linebelow=Line("-", "-", "-", "-"),
                        headerrow=DataRow("|", "|", "|"),
                        datarow=DataRow("|", "|", "|"),
                        padding=0,
                        with_header_hide=None)

def getLogFile(path):
    """Return the path to run.stats."""
    return os.path.join(path, 'run.stats')


class LazyEvalList:
    """Store all the lines in run.stats and eval() when needed."""
    def __init__(self, lines):
        # The first line in the records contains headers.
        self.lines = lines[1:]

    def __getitem__(self, index):
        if isinstance(self.lines[index], str):
            self.lines[index] = eval(self.lines[index])
        return self.lines[index]

    def __len__(self):
        return len(self.lines)


def getMatchedRecordIndex(records, column, target):
    """Find target from the specified column in records."""
    target = int(target)
    lo = 0
    hi = len(records) - 1
    while lo < hi:
        mid = (lo + hi) // 2
        if column(records[mid]) <= target:
            lo = mid + 1
        else:
            hi = mid
    return lo


def aggregateRecords(records):
    # index for memUsage and stateCount in run.stats
    memIndex = 6
    stateIndex = 5

    # maximum and average memory usage
    memValues = list(map(itemgetter(memIndex), records))
    maxMem = max(memValues) / 1024 / 1024
    avgMem = sum(memValues) / len(memValues) / 1024 / 1024

    # maximum and average number of states
    stateValues = list(map(itemgetter(stateIndex), records))
    maxStates = max(stateValues)
    avgStates = sum(stateValues) / len(stateValues)

    return (maxMem, avgMem, maxStates, avgStates)


def stripCommonPathPrefix(paths):
    paths = map(os.path.normpath, paths)
    paths = [p.split('/') for p in paths]
    zipped = zip(*paths)
    i = 0
    for i, elts in enumerate(zipped):
        if len(set(elts)) > 1:
            break
    return ['/'.join(p[i:]) for p in paths]


def getKeyIndex(key, labels):
    """Get the index of the specified key in labels."""
    def normalizeKey(key):
        return re.split('\W', key)[0]

    for i, title in enumerate(labels):
        if normalizeKey(title) == normalizeKey(key):
            return i
    else:
        raise ValueError('invalid key: {0}'.format(key))


def getKleeOutDirs(dirs):
    kleeOutDirs = []
    for dir in dirs:
        if os.path.exists(os.path.join(dir, 'info')):
            kleeOutDirs.append(dir)
        else:
            for root, subdirs, _ in os.walk(dir):
                for d in subdirs:
                    path = os.path.join(root, d)
                    if os.path.exists(os.path.join(path, 'info')):
                        kleeOutDirs.append(path)
    return kleeOutDirs


def getLabels(pr):
    if pr == 'all':
        labels = ('Path', 'Instrs', 'Time(s)', 'ICov(%)', 'BCov(%)', 'ICount',
                  'TSolver(%)', 'States', 'maxStates', 'avgStates', 'Mem(MB)',
                  'maxMem(MB)', 'avgMem(MB)', 'Queries', 'AvgQC', 'Tcex(%)',
                  'Tfork(%)')
    elif pr == 'reltime':
        labels = ('Path', 'Time(s)', 'TUser(%)', 'TSolver(%)',
                  'Tcex(%)', 'Tfork(%)', 'TResolve(%)')
    elif pr == 'abstime':
        labels = ('Path', 'Time(s)', 'TUser(s)', 'TSolver(s)',
                  'Tcex(s)', 'Tfork(s)', 'TResolve(s)')
    elif pr == 'more':
        labels = ('Path', 'Instrs', 'Time(s)', 'ICov(%)', 'BCov(%)', 'ICount',
                  'TSolver(%)', 'States', 'maxStates', 'Mem(MB)', 'maxMem(MB)')
    else:
        labels = ('Path', 'Instrs', 'Time(s)', 'ICov(%)',
                  'BCov(%)', 'ICount', 'TSolver(%)')
    return labels


def getRow(record, stats, pr):
    """Compose data for the current run into a row."""
    I, BFull, BPart, BTot, T, St, Mem, QTot, QCon,\
        _, Treal, SCov, SUnc, _, Ts, Tcex, Tf, Tr = record
    maxMem, avgMem, maxStates, avgStates = stats

    # special case for straight-line code: report 100% branch coverage
    if BTot == 0:
        BFull = BTot = 1

    Mem = Mem / 1024 / 1024
    AvgQC = int(QCon / max(1, QTot))

    if pr == 'all':
        row = (I, Treal, 100 * SCov / (SCov + SUnc),
               100 * (2 * BFull + BPart) / (2 * BTot), SCov + SUnc,
               100 * Ts / Treal, St, maxStates, avgStates,
               Mem, maxMem, avgMem, QTot, AvgQC,
               100 * Tcex / Treal, 100 * Tf / Treal)
    elif pr == 'reltime':
        row = (Treal, 100 * T / Treal, 100 * Ts / Treal,
               100 * Tcex / Treal, 100 * Tf / Treal,
               100 * Tr / Treal)
    elif pr == 'abstime':
        row = (Treal, T, Ts, Tcex, Tf, Tr)
    elif pr == 'more':
        row = (I, Treal, 100 * SCov / (SCov + SUnc),
               100 * (2 * BFull + BPart) / (2 * BTot),
               SCov + SUnc, 100 * Ts / Treal,
               St, maxStates, Mem, maxMem)
    else:
        row = (I, Treal, 100 * SCov / (SCov + SUnc),
               100 * (2 * BFull + BPart) / (2 * BTot),
               SCov + SUnc, 100 * Ts / Treal)
    return row


def drawLineChart(vectors, titles):
    """Draw a line chart based on data from vectors.

    Args:
        vectors: A list of vectors. Each vector is drawn in a separate
                 subchart.
        titles: A list of strings. Each string is an explanation of the
                corresponding vector.
    """
    try:
        import matplotlib
        matplotlib.use('Agg')  # plot without X-server connection
        import matplotlib.pyplot as plt
    except:
        print('Error: Package "matplotlib" required for figure plotting. '
              'Please install it using "pip" or your package manager.',
              file=sys.stderr)
        exit(1)
    from math import sqrt, ceil

    nFigs = len(vectors)       # number of subfigures
    nRow = int(sqrt(nFigs))    # number of rows in the figure
    nCol = ceil(nFigs / nRow)  # number of columns in the figure

    fig = plt.figure()
    for i in range(nFigs):
        ax = fig.add_subplot(nRow, nCol, i)
        ax.plot(vectors[i])
        ax.set_title(titles[i])
        #ax.set_xlabel()
        #ax.set_ylabel()

    #fig.show()
    #raw_input('Press any key to continue..')

    def getFileName():
        """Get a unused file name in current directory."""
        i = 0
        fileName = 'fig-klee-stats-{0}.png'.format(i)
        while os.path.exists(fileName):
            i = i + 1
            fileName = 'fig-klee-stats-{0}.png'.format(i)
        return fileName

    plt.savefig(getFileName(),
                dpi=150,  # dots-per-inch, default: 100
                bbox_inches='tight',  # tighter bbox
                transparent=False)


def main():
    # function for sanitizing arguments
    def isPositiveInt(value):
        try:
            value = int(value)
        except ValueError:
            raise argparse.ArgumentTypeError(
                'integer expected: {0}'.format(value))
        if value <= 0:
            raise argparse.ArgumentTypeError(
                'positive integer expected: {0}'.format(value))
        return value

    parser = argparse.ArgumentParser(
        description='output statistics logged by klee',
        epilog='LEGEND\n' + tabulate(Legend),
        formatter_class=argparse.RawDescriptionHelpFormatter)

    parser.add_argument('dir', nargs='+', help='klee output directory')

    parser.add_argument('--precision',
                        dest='precision', type=isPositiveInt,
                        default=2, metavar='n',
                        help='Floating point numbers display precision '
                        '(default: 2).')
    parser.add_argument('--table-format',
                        choices=['plain', 'simple', 'grid', 'pipe', 'orgtbl',
                                 'rst', 'mediawiki', 'latex', 'klee'],
                        dest='tableFormat', default='klee',
                        help='Table format for the summary.')
    parser.add_argument('--draw-line-chart',
                        dest='drawLineChart', metavar='header-list',
                        help='Draw line chart for a list of columns. '
                        'Columns must be chosen from the headers of the '
                        'table outputted and separated by comma (e.g., '
                        '--draw-line-chart=Instrs,Time). Data points '
                        'on x-axis correspond to lines in run.stats.')
    parser.add_argument('--sample-interval', dest='sampleInterv',
                        type=isPositiveInt, default='10', metavar='n',
                        help='Sample a data point every n lines for a '
                        'run.stats (default: 10)')

    # argument group for controlling output verboseness
    pControl = parser.add_mutually_exclusive_group(required=False)
    pControl.add_argument('--print-all',
                          action='store_true', dest='pAll',
                          help='Print all available information.')
    pControl.add_argument('--print-rel-times',
                          action='store_true', dest='pRelTimes',
                          help='Print only values of measured times. '
                          'Values are relative to the measured system '
                          'execution time.')
    pControl.add_argument('--print-abs-times',
                          action='store_true', dest='pAbsTimes',
                          help='Print only values of measured times. '
                          'Absolute values (in seconds) are printed.')
    pControl.add_argument('--print-more',
                          action='store_true', dest='pMore',
                          help='Print extra information (needed when '
                          'monitoring an ongoing run).')

    # arguments for sorting
    parser.add_argument('--sort-by', dest='sortBy', metavar='header',
                        help='Key value to sort by. Must be chosen from '
                        'the headers of the table outputted  (e.g., '
                        '--sort-by=Instrs).')
    parser.add_argument('--ascending',
                        dest='ascending', action='store_true',
                        help='Sort in ascending order (default: False).')

    # arguments for comparing
    parser.add_argument('--compare-by', dest='compBy', metavar='header',
                        help='Key value on which to compare runs to the '
                        'reference one (which is the first one). Must be '
                        'chosen from the headers of the table oputputted. '
                        'e.g., --compare-by=Instrs shows how each run '
                        'compares to the reference run after executing the '
                        'same number of instructions as the reference run. '
                        "If a run hasn't executed as many instructions as "
                        'the reference one, we simply print the statistics '
                        'at the end of that run.')
    parser.add_argument('--compare-at', dest='compAt', metavar='value',
                        help='Value to compare the runs at. Can be special '
                        "value 'last' to compare at the last point which "
                        'makes sense. Use in conjunction with --compare-by.')

    args = parser.parse_args()

    # get print controls
    pr = 'NONE'
    if args.pAll:
        pr = 'all'
    elif args.pRelTimes:
        pr = 'reltime'
    elif args.pAbsTimes:
        pr = 'abstime'
    elif args.pMore:
        pr = 'more'

    dirs = getKleeOutDirs(args.dir)
    if len(dirs) == 0:
        print('no klee output dir found', file=sys.stderr)
        exit(1)
    # read contents from every run.stats file into LazyEvalList
    data = [LazyEvalList(list(open(getLogFile(d)))) for d in dirs]
    if len(data) > 1:
        dirs = stripCommonPathPrefix(dirs)
    # attach the stripped path
    data = list(zip(dirs, data))

    labels = getLabels(pr)
    # labels in the same order as in the run.stats file. used by --compare-by.
    # current impl needs monotonic values, so only keep the ones making sense.
    rawLabels = ('Instrs', '', '', '', '', '', '', 'Queries',
                 '', '', 'Time', 'ICov', '', '', '', '', '', '')

    if args.compBy:
        # index in the record of run.stats
        compIndex = getKeyIndex(args.compBy, rawLabels)
        if args.compAt:
            if args.compAt == 'last':
                # [records][last-record][compare-by-index]
                refValue = min(map(lambda r: r[1][-1][compIndex], data))
            else:
                refValue = args.compAt
        else:
            refValue = data[0][1][-1][compIndex]

    # build the main body of the table
    table = []
    totRecords = []  # accumulated records
    totStats = []    # accumulated stats
    for path, records in data:
        row = [path]
        if args.compBy:
            matchIndex = getMatchedRecordIndex(
                records, itemgetter(compIndex), refValue)
            stats = aggregateRecords(LazyEvalList(records[:matchIndex + 1]))
            totStats.append(stats)
            row.extend(getRow(records[matchIndex], stats, pr))
            totRecords.append(records[matchIndex])
        else:
            stats = aggregateRecords(records)
            totStats.append(stats)
            row.extend(getRow(records[-1], stats, pr))
            totRecords.append(records[-1])
        table.append(row)
    # calculate the total
    totRecords = [sum(e) for e in zip(*totRecords)]
    totStats = [sum(e) for e in zip(*totStats)]
    totalRow = ['Total ({0})'.format(len(table))]
    totalRow.extend(getRow(totRecords, totStats, pr))

    if args.sortBy:
        table = sorted(table, key=itemgetter(getKeyIndex(args.sortBy, labels)),
                       reverse=(not args.ascending))

    if len(data) > 1:
        table.append(totalRow)
    table.insert(0, labels)

    if args.tableFormat != 'klee':
        print(tabulate(
            table, headers='firstrow',
            tablefmt=args.tableFormat,
            floatfmt='.{p}f'.format(p=args.precision),
            numalign='right', stralign='center'))
    else:
        stream = tabulate(
            table, headers='firstrow',
            tablefmt=KleeTable,
            floatfmt='.{p}f'.format(p=args.precision),
            numalign='right', stralign='center')
        # add a line separator before the total line
        if len(data) > 1:
            stream = stream.splitlines()
            stream.insert(-2, stream[-1])
            stream = '\n'.join(stream)
        print(stream)

    if args.drawLineChart:
        if len(dirs) != 1:
            print('--draw-line-chart only supports using a single file',
                  file=sys.stderr)
            exit(1)

        # sample according to the given interval
        samples = [r for i, r in enumerate(data[0][1])
                   if i % args.sampleInterv == 0]
        vectors = []
        for i in range(len(samples)):
            # aggregate all the samples upto the i-th one
            stats = aggregateRecords(samples[:i + 1])
            vectors.append(getRow(samples[i], stats, pr))

        titles = args.drawLineChart.split(',')
        # Index returned by getKeyIndex() against labels is starting from
        # 'Path' column.  vectors here doesn't have the 'Path' column.
        indices = [getKeyIndex(e, labels) - 1 for e in titles]
        # get one more column to make sure itemgetter always return tuples
        indices.append(0)
        # strip columns not specified by the user
        vectors = map(itemgetter(*indices), vectors)
        # combine elements in the same column into the same tuple
        vectors = list(zip(*vectors))
        # remove the column we get but not needed
        vectors = vectors[:-1]

        drawLineChart(vectors, titles)


if __name__ == '__main__':
    main()