concolico.py

from z3 import *

def find_path(path, f, *args, **kwargs):
    log = []
    def unwrap(other):
        try:
            return other.v
        except:
            return other
    class Sym:
        def __init__(self, v):
            self.v = v
            
        def __repr__(self):
            return self.v.__repr__()

        # Arithmetic operators
        def __add__(self, other):
            return Sym(self.v.__add__(unwrap(other)))

        def __sub__(self, other):
            return Sym(self.v.__sub__(unwrap(other)))

        def __mul__(self, other):
            return Sym(self.v.__mul__(unwrap(other)))

        def __div__(self, other):
            return Sym(self.v.__div__(unwrap(other)))

        def __neg__(self):
            return Sym(self.v.__neg__())
    
        # Comparison operators
        def __eq__(self, other):
            return Sym(self.v.__eq__(unwrap(other)))

        def __ne__(self, other):
            return Sym(self.v.__ne__(unwrap(other)))

        def __le__(self, other):
            return Sym(self.v.__le__(unwrap(other)))
        
        def __ge__(self, other):
            return Sym(self.v.__ge__(unwrap(other)))

        def __lt__(self, other):
            return Sym(self.v.__lt__(unwrap(other)))

        def __gt__(self, other):
            return Sym(self.v.__gt__(unwrap(other)))

        # Boolean operators
        def __and__(self, other):
            return Sym(And(self.v, unwrap(other)))

        def __or__(self, other):
            return Sym(Or(self.v, unwrap(other)))

        def __not__(self):
            return Sym(Not(self.v))

        # This is the crucial bit: when we evaluate a boolean
        # expression, we check what branch we need to take (true or
        # false), and add a constraint to the solver accordingly
        def __bool__(self):
            if path.pop():
                log.append(self.v)
                return True
            else:
                log.append(Not(self.v))
                return False

    f(*map(Sym,args), **kwargs)
    print(log)
    s = Solver()
    map(lambda x: s.add(x), log)
    if s.check().r == 1:
        print(s.model())
    else:
        print('unsat')

def f(x, y):
    x = x + y*2
    print(x)
    if x == 5:
        z = 0
    else:
        z = 1

find_path([True], f, Real('x'), Real('y'))