z3proofternaryadd.py

# mini proof of concept of how to use Z3 to prove the correctness of bit
# manipulation formulas, using the Ternary.__add__ function from
# https://gist.github.com/dougallj/9211fd24c3759f7f340dede28929c659 as an
# example

# partly inspired by Philip Zucker's post about Z3 and ranges:
# https://www.philipzucker.com/more-stupid-z3py-tricks-simple-proofs/

import z3

N_BITS = 64
MASK = (1 << N_BITS) - 1

# this class is a fragment from the gist linked above. an instance of Ternary
# represents a set of integers where some bits are known to be 1, some bits
# known 0, and some bits are unknown.

class Ternary:
    def __init__(self, ones, unknowns):
        self.ones = ones & MASK
        self.unknowns = unknowns & MASK
        # assert removed, doesn't work with the prove
        #assert (self.ones & self.unknowns) == 0, (bin(self.ones), bin(self.unknowns))

    def __add__(self, other):
        x = self.ones + other.ones
        u = self.unknowns | other.unknowns | (x ^ (x + self.unknowns + other.unknowns))
        return Ternary(x & ~u, u)

    def contains(self, value):
        # checks whether the Ternary number (interpreted as a set) contains the
        # contrete value
        value = value & MASK
        return (self.ones & ~self.unknowns) == (value & ~self.unknowns)



# soundness: what we want to prove: for all integers x, y and all valid
# Ternarys a, b with a.contains(x) and b.contains(y) the following should hold:
# (a + b).contains(x + y)

# Z3 is an smt solver, it tries to find whether a formula is satisfiable. this
# means it tries to find concrete values for all the variables in the formula
# so that the formula is true. therefore it can decide queries of the form:
# exists var1, var2, ..., varn: formula(var1, ..., varn)
# and give concrete values for the variables.

# to prove a "forall" statement, we use z3 to check that no counterexample
# exists

def BitVec(name):
    # a machine integer in z3 is a "bit vector"
    return z3.BitVec(name, N_BITS)

def z3_ternary(name):
    # construct a tuple of a Ternary instance where the ones and the unknowns
    # are z3 bitvectors, and a bitvector variable that is the "concrete"
    # number that is supposed to be a member of the Ternary instance
    variable = BitVec(name + "_concrete")
    ones = BitVec(name + "_ones")
    unknowns = BitVec(name + "_unknowns")
    return Ternary(ones, unknowns), variable

def prove(cond):
    # to prove something, try to find a counterexample for the negation
    solver = z3.Solver()
    print("proving", cond)
    z3res = solver.check(z3.Not(cond))
    if z3res == z3.unsat:
        print("worked!") # not possible, therefore the formula holds
    elif z3res == z3.sat:
        # not possible to prove!
        model = solver.model()
        print("couldn't prove", cond, "counterexamples:", model)
        raise ValueError

def prove_implies(*args):
    # helper to check that some antecedents imply a consequence
    last = args[-1]
    prev = args[:-1]
    return prove(z3.Implies(z3.And(*prev), last))

def ternary_well_formed(t):
    # checks whether t is a "valid" Ternary, this is the commented out assert
    # in the original code (but we can't use asserts with z3 variables)
    return (t.ones & t.unknowns) == 0

def test_example_contains():
    # Ternary(0, -2) are the even numbers
    assert Ternary(0, -2).contains(2)
    assert Ternary(0, -2).contains(4)
    assert Ternary(0, -2).contains(-2)
    assert not Ternary(0, -2).contains(1)
    assert not Ternary(0, -2).contains(3)
    assert not Ternary(0, -2).contains(-5)

def test_prove_addition():
    # proof for addition
    # given two Ternary and two integers
    t1, v1 = z3_ternary('self')
    t2, v2 = z3_ternary('other')
    # add the Ternarys together
    t3 = t1 + t2
    # this will construct a new Ternary instance, where the ones and unknowns
    # are z3 formulas! z3 variables (and formulas) have all the Python
    # operators overloaded, so something like this in __add__:
    # x = self.ones + other.ones
    # will construct a new formula term that is the sum of the two ones
    # variables
    prove_implies(
        # we need as a pre-condition that t1 and t2 are well-formed
        ternary_well_formed(t1),
        ternary_well_formed(t2),
        t1.contains(v1),
        t2.contains(v2),
        # we prove that the result is well-formed and that t3 contains the sum
        # of the concrete integer values
        z3.And(ternary_well_formed(t3), t3.contains(v1 + v2))
    )
    # unfortunately this approach only works if there is no control flow at all
    # in __add__, __init__, contains!

test_example_contains()
test_prove_addition()


# running it prints this huge formula which is successfully proven for all
# values of the variables:

# Implies(And(self_ones &
#             18446744073709551615 &
#             self_unknowns &
#             18446744073709551615 ==
#             0,
#             other_ones &
#             18446744073709551615 &
#             other_unknowns &
#             18446744073709551615 ==
#             0,
#             self_ones &
#             18446744073709551615 &
#             ~(self_unknowns & 18446744073709551615) ==
#             self_concrete &
#             18446744073709551615 &
#             ~(self_unknowns & 18446744073709551615),
#             other_ones &
#             18446744073709551615 &
#             ~(other_unknowns & 18446744073709551615) ==
#             other_concrete &
#             18446744073709551615 &
#             ~(other_unknowns & 18446744073709551615)),
#         And((self_ones & 18446744073709551615) +
#             (other_ones & 18446744073709551615) &
#             ~(self_unknowns & 18446744073709551615 |
#               other_unknowns & 18446744073709551615 |
#               (self_ones & 18446744073709551615) +
#               (other_ones & 18446744073709551615) ^
#               (self_ones & 18446744073709551615) +
#               (other_ones & 18446744073709551615) +
#               (self_unknowns & 18446744073709551615) +
#               (other_unknowns & 18446744073709551615)) &
#             18446744073709551615 &
#             (self_unknowns & 18446744073709551615 |
#              other_unknowns & 18446744073709551615 |
#              (self_ones & 18446744073709551615) +
#              (other_ones & 18446744073709551615) ^
#              (self_ones & 18446744073709551615) +
#              (other_ones & 18446744073709551615) +
#              (self_unknowns & 18446744073709551615) +
#              (other_unknowns & 18446744073709551615)) &
#             18446744073709551615 ==
#             0,
#             (self_ones & 18446744073709551615) +
#             (other_ones & 18446744073709551615) &
#             ~(self_unknowns & 18446744073709551615 |
#               other_unknowns & 18446744073709551615 |
#               (self_ones & 18446744073709551615) +
#               (other_ones & 18446744073709551615) ^
#               (self_ones & 18446744073709551615) +
#               (other_ones & 18446744073709551615) +
#               (self_unknowns & 18446744073709551615) +
#               (other_unknowns & 18446744073709551615)) &
#             18446744073709551615 &
#             ~((self_unknowns & 18446744073709551615 |
#                other_unknowns & 18446744073709551615 |
#                (self_ones & 18446744073709551615) +
#                (other_ones & 18446744073709551615) ^
#                (self_ones & 18446744073709551615) +
#                (other_ones & 18446744073709551615) +
#                (self_unknowns & 18446744073709551615) +
#                (other_unknowns & 18446744073709551615)) &
#               18446744073709551615) ==
#             self_concrete + other_concrete &
#             18446744073709551615 &
#             ~((self_unknowns & 18446744073709551615 |
#                other_unknowns & 18446744073709551615 |
#                (self_ones & 18446744073709551615) +
#                (other_ones & 18446744073709551615) ^
#                (self_ones & 18446744073709551615) +
#                (other_ones & 18446744073709551615) +
#                (self_unknowns & 18446744073709551615) +
#                (other_unknowns & 18446744073709551615)) &
#               18446744073709551615)))

def test_prove_addition_commutative():
    t1, v1 = z3_ternary('self')
    t2, v2 = z3_ternary('other')
    t3a = t1 + t2
    t3b = t2 + t1
    prove(t3a.ones == t3b.ones)
    prove(t3a.unknowns == t3b.unknowns)

def test_prove_addition_associative():
    t1, v1 = z3_ternary('self')
    t2, v2 = z3_ternary('other')
    t3, v3 = z3_ternary('third')
    t4a = (t1 + t2) + t3
    t4b = t1 + (t2 + t3)
    # this fails! to understand the counterexample, you
    # can set N_BITS=4
    prove(t4a.ones == t4b.ones)
    prove(t4a.unknowns == t4b.unknowns)

test_prove_addition_commutative()
test_prove_addition_associative()


# note that all our proofs are about safety (soundness), not precision!