from z3 import *
k = 40
bvzero = BitVecVal(0, 32)
bvtwo = BitVecVal(2, 32)

def state_string(state, m):
  (pc,i,j,n,s) = state
  return str(m[pc]) + ": i=" + str(m[i]) + " j=" + str(m[j]) + \
    " n=" + str(m[n]) + " s=" + str(m[s])

def new_state(i):
  stri = str(i)
  return (Int('pc'+stri), BitVec('i'+stri, 32), BitVec('j'+stri, 32),\
          BitVec('n'+stri, 32), BitVec('s'+stri, 32))

# create variables for k+1 different states
states = [new_state(i) for i in range (k+1)]

# predicate for initial states
def initial(state):
  (pc,i,j,n,s) = state
  return pc == 2

# predicate for transition step in program
def trans(state, next_state):
  (pc,i,j,n,s) = state
  (pcx,ix,jx,nx,sx) = next_state
  keep_all = And(n == nx, s == sx, i == ix, j == jx)
  return Or( \
    And(pc == 2, pcx == 3, n == nx), \
    And(pc == 3, pcx == 4, ix == 0, jx == 10, sx == 0, n == nx), \
    And(pc == 4, pcx == 5, i <= n, keep_all), \
    And(pc == 5, pcx == 6, i < j, keep_all), \
    And(pc == 5, pcx == 7, i >= j, keep_all), \
    And(pc == 6, pcx == 7, sx == s + bvtwo, n == nx, i == ix, j == jx), \
    And(pc == 7, pcx == 8, jx == j - 1, n == nx, s == sx, i == ix), \
    And(pc == 8, pcx == 4, ix == i + 1, n == nx, s == sx, j == jx), \
    And(pc == 4, pcx == 9, i > n, keep_all), \
    # add dummy transition to avoid that counterexample is not found
    # because no transition exists:
    And(pc == 9, pcx == 9, keep_all)
  )

# predicate for violation of assertion
def violating(state):
  (pc,i,j,n,s) = state
  return And(pc == 9, Not(s == bvtwo * n), Not (s == bvzero))


solver = Solver()
# the first state must be initial
solver.add(initial(states[0]))
# all other states must be connected to their predecessor by a program step
for i in range(0,k):
  solver.add(trans(states[i], states[i+1]))

# at some state there should be a violation
violated = False
for i in range(0,k):
  violated = Or(violated, violating(states[i+1]))
solver.add(violated)

# check and print result
result = solver.check()
if result == z3.sat:
  m = solver.model()
  print state_string(states[0],m)
  for i in range(1,k+1):
    print "  -> " + state_string(states[i],m)
    if m.eval(violating(states[i])) == True:
      break