lean_profile_to_json.py

#!/usr/bin/env python3
import argparse
import re, json
from pathlib import Path
from collections import deque

parser = argparse.ArgumentParser(description='Convert a log file generated from running lean --profile to JSON format')
parser.add_argument('-i', metavar='path/to/profile.log', type=str, nargs='?',
                    help='path to profile file', default='profile.log')
parser.add_argument('-o', metavar='path/to/profile.json', type=str, nargs='?',
                    help='output json file', default='profile.json')
parser.add_argument('-n', action='store_true', help='if this flag is present, the data will be exported as ndjson rather than json')

args = parser.parse_args()
profile = str(Path(args.i).resolve())
output = str(Path(args.o).resolve())

with open(profile) as f:
    read_data = f.read()

read_array = read_data.split('\n')

def timeInMs(timeStr):
    if (timeStr.endswith('ms')):
        return float(timeStr[:-2])
    else:
        return 1000*float(timeStr[:-1])

def parse(arr):
    arrD = deque(arr)
    out = []
    while len(arrD) > 0:
        line = arrD.popleft()
        if len(line) == 0:
            return out
        first = line[0]
        lineSplit = re.split('\s+', line)
        if first == 'p':
            # parsing took {time}
            out.append({
                'type': 'pars',
                't': timeInMs(lineSplit[-1])
            })
            continue
        elif first == 't':
            # type checking of {decl} took {time}
            out.append({
                'type': 'tc',
                't': timeInMs(lineSplit[-1]),
                'decl': lineSplit[3]
            })
            continue
        elif first == 'c':
            # compilation of {decl} took {time}
            out.append({
                'type': 'comp',
                't': timeInMs(lineSplit[-1]),
                'decl': lineSplit[2]
            })
            continue
        elif first == 'd':
            # decl post-processing of {decl} took {time}
            out.append({
                'type': 'decl',
                't': timeInMs(lineSplit[-1]),
                'decl': lineSplit[3]
            })
            continue
        elif first == 'e':
            # elaboration...
            tacList = []
            tacObj = {'type': 'elab', 'lines': tacList}
            out.append(tacObj)
            while line[11] == ':':
                # elaboration: tactic ...
                if lineSplit[2][0] == 'c':
                    # elaboration: tactic compilation took {time}
                    tacList.append({
                        'type': 'comp',
                        't': timeInMs(lineSplit[-1])
                    })
                    line = arrD.popleft()
                    lineSplit = re.split('\s+', line)
                    continue
                elif lineSplit[2][0] == 'e':
                    # elaboration: tactic execution took {time}
                    execList = []
                    tacList.append({
                        'type': 'exec',
                        't': timeInMs(lineSplit[-1]),
                        'tac': execList
                    })
                    line = arrD.popleft()
                    lineSplit = re.split('\s+', line)
                    if line[0] == 'n':
                        while line[0] == 'n':
                            # num. allocated ...
                            allocatedFirst = lineSplit[2][0]
                            if allocatedFirst == 'o':
                                # num. allocated objects: {nObj}
                                tacObj['nO'] = int(lineSplit[-1])
                            elif allocatedFirst == 'c':
                                # num. allocated closures: {nClos}
                                tacObj['nC'] = int(lineSplit[-1])
                            elif allocatedFirst == 'b':
                                # num. allocated big nums: {nBig}
                                tacObj['nB'] = int(lineSplit[-1])
                            else:
                                raise Exception('unexpected after "num. allocated"', line)
                            line = arrD.popleft()
                            lineSplit = re.split('\s+', line.strip())
                        while (len(lineSplit) == 3) and lineSplit[0][0] != 'p':
                            # exclude "parsing took {time}"
                            # {time}   {pct}   {decl}
                            execList.append({
                                't': timeInMs(lineSplit[0]),
                                'pct': float(lineSplit[1][:-1]),
                                'decl': lineSplit[-1]
                            })
                            line = arrD.popleft()
                            lineSplit = re.split('\s+', line.strip())
                else:
                    raise Exception('unexpected after "elaboration: tactic"', line)
            if line[11] == ' ':
                # elaboration of {decl} took {time}
                tacObj['t'] = timeInMs(lineSplit[-1])
                tacObj['decl'] = lineSplit[2]
            elif line[0] != 'e':
                # an extra line was read in the "elaboration: tactic" loop, so we should put it back
                arrD.appendleft(line)
            else:
                raise Exception('unexpected after "elaboration"', line)
            continue
        else:
            raise Exception('unexpected line', line)
    return out

parsed_array = parse(read_array)

if args.n:
    json_data = "\n".join(map(json.dumps, parsed_array))

else:
    json_data = json.dumps(parsed_array, indent='')


with open(output, 'w') as f:
    print(json_data, file=f)