Visualize Z3's proofs

z3proof.py

#!/usr/bin/env python

# Author: Darius Foo (darius.foo.tw@gmail.com)
# License: GPLv3

from z3 import *

set_param(proof=True)
s = Solver()

x, y = Ints('x y')

s.add(2 * x > 3 * y)
s.add(y > x)
s.add(y > 0)
s.add(x > 0)

a, b = Bools('a b')

# s.add(x > 0)
# s.add(x == 0)

f = a == a
# f = a == Not(Not(a))
s.add(Not(f))

# s.add(Not(a == Not(Not(a))))
# s.add(a == Not(a))

# s.add(And(a, Not(a)))

# s.add(And(a, Not(a)))

# s.add(Not(ForAll([a], Implies(a, a))))
# s.add(Not(Implies(a, a)))

# s.add(Implies(a, b))
# s.add(Not(b))
# s.add(a)

def to_tex(pf, n=0):
    name = pf.decl().name()
    if name in {'mp', 'trans', 'monotonicity', 'iff-false', 'iff-true', 'th-lemma', 'lemma', 'unit-resolution', 'quant-inst', 'and-elim'}:
        hyp = []
        for i in range(pf.num_args() - 1):
            hyp.append(pf.arg(i))
        con = pf.arg(pf.num_args() - 1)
        sep = '\\\\'
        # sep = sep + sep
        hyps = f'\n{"  "*n}{sep}\n{"  "*n}'.join(to_tex(h, n+1) for h in hyp)
        cons = to_tex(con, n+1)
        res = f'\\inferrule* [right={name}]\n{"  "*n}{{{hyps}}}\n{"  "*n}{{{cons}}}'
        # return f'\\fbox{{${res}$}}'
        return res
    # elif name == 'rewrite':
    #     return to_tex(pf.arg(0), n+1)
    elif name in {'asserted', 'rewrite'}:
        # note the extra space in the premise here...
        return f'\\inferrule* [right={name}] {{ }} {{{to_tex(pf.arg(0), n+1)}}}'
        # return to_tex(pf.arg(0))
    elif name in {'or', 'Or'}: # ?
        return f'{to_tex(pf.arg(0))} \\vee {to_tex(pf.arg(1))}'
    elif name == 'and':
        return f'{to_tex(pf.arg(0))} \\wedge {to_tex(pf.arg(1))}'
    elif name == '=>':
        return f'{to_tex(pf.arg(0))} \\to {to_tex(pf.arg(1))}'
    elif name == '<=':
        return f'{to_tex(pf.arg(0))} \\leq {to_tex(pf.arg(1))}'
    elif name == '>=':
        return f'{to_tex(pf.arg(0))} \\geq {to_tex(pf.arg(1))}'
    elif name == '>':
        return f'{to_tex(pf.arg(0))} > {to_tex(pf.arg(1))}'
    elif name == '=':
        # return f'{to_tex(pf.arg(0))} \\leftrightarrow {to_tex(pf.arg(1))}'
        return f'{to_tex(pf.arg(0))} = {to_tex(pf.arg(1))}'
    elif name == 'not':
        thing = to_tex(pf.arg(0))
        if len(thing) > 1 and pf.arg(0).decl().name() not in {'true', 'false'}:
            thing = f'({thing})'
        return f'\\neg {thing}'
    elif name == 'false':
        return f'\\bot'
    elif name == 'true':
        return f'\\top'
    else:
        # print('name', name)
        return str(pf)
        # return f'\\fbox{{${pf}$}}'

if s.check() == sat:
    print(s.model())
else:
    p = s.proof()
    # print(p)
    print(r'''\documentclass{article}
\usepackage{mathpartir}
\usepackage{amsmath}
\usepackage[a0paper, landscape]{geometry}
\begin{document}
$$''')
    print(to_tex(p))
    print(r'''$$
\end{document}
''')