import re


class Link(object):
    def __init__(self, src, dst, label):
        self._src = 'bb%i' % (src,)
        self._dst = dst if isinstance(dst, str) else 'bb%i' % (dst,)
        self._label = label

def main():

    #cratename,pat = 'rustc_lint','fn .*expr_refers_to_this_method.*'
    cratename,pat = 'std','fn resize.*HashMap'
    #cratename,pat = 'rustc', 'fn tables.*::"rustc"::ty::context::TyCtxt'

    fp = open('output/lib'+cratename+'.hir_3_mir.rs');
    start_pat = re.compile(pat)
    def_line = None
    for line in fp:
        line = line.strip()
        if start_pat.match(line) != None:
            print "#",line
            def_line = line
            break

    if def_line is None:
        return

    for line in fp:
        if line.strip() == "bb0: {":
            break

    bbs = []
    cur_bb_lines = []
    level = 2
    for line in fp:
        line = line.strip()
        if line == "}":
            level -= 1
            if level == 0:
                break
            else:
                bbs.append( cur_bb_lines )
                cur_bb_lines = []
                continue

        if "bb" in line and ": {" in line:
            level += 1
            continue

        outstr = ""
        comment_level = 0
        i = 0
        while i < len(line):
            if comment_level > 0:
                if line[i:i+2] == '*/':
                    comment_level -= 1
                    i += 2
                    continue
            if line[i:i+2] == '/*':
                comment_level += 1
                i += 2
                continue 
            if comment_level == 0:
                outstr += line[i]
            i += 1
        print "#",len(bbs),outstr

        cur_bb_lines.append(outstr)

    goto_regex = re.compile('goto bb(\d+);$')
    call_regex = re.compile('.*goto bb(\d+) else bb(\d+)$')
    if_regex = re.compile('.*goto bb(\d+); } else { goto bb(\d+); }$')
    switch_regex = re.compile('(\d+) => bb(\d+),')

    links = []
    for idx,bb in enumerate(bbs):
        if bb[-1] == 'return;':
            links.append( Link(idx, 'return', "return") )
            continue
        if bb[-1] == 'diverge;':
            #links.append( Link(idx, 'panic', "diverge") )
            continue
        m = goto_regex.match(bb[-1])
        if m != None:
            links.append( Link(idx, int(m.group(1)), "") )
            continue
        m = call_regex.match(bb[-1])
        if m != None:
            links.append( Link(idx, int(m.group(1)), "ret") )
            #links.append( Link(idx, int(m.group(2)), "panic") )
            continue
        m = if_regex.match(bb[-1])
        if m != None:
            links.append( Link(idx, int(m.group(1)), "true") )
            links.append( Link(idx, int(m.group(2)), "false") )
            continue


        for m in switch_regex.finditer(bb[-1]):
            links.append( Link(idx, int(m.group(2)), "var%s" % (m.group(1),) ) )



    print "digraph {"
    for l in links:
        print '"%s" -> "%s" [label="%s"];' % (l._src, l._dst, l._label)

    print ""
    for idx,bb in enumerate(bbs):
        print '"bb%i" [label="BB%i:' % (idx,idx,),
        for stmt in bb:
            print '\\n',stmt.replace('"', '\\"'),
        print '"];'
    print "}"


main()