Search code examples
pythonz3z3py

How to extract the value of ArithRef from within the solver constraint in Z3py?


I am trying to add a constraint from within the Z3py code that requires extracting of a value for an ArithRef variable but I'm unable to do so. I have attached the Z3 snippet below:

from z3 import *
import random

R = 3
H = 4

Dest = [[2,3,4], [0,3,2,5], [1,4,5], [1,4,2], [3,1], [2,3,1]]

s = Solver()

T =[[ Int('r%d_t%d' % (k,t))
         for k in range (R)] for t in range (H)]


for t in range(H):
    for r in range(R):
        s.add(If(t==0, If(r==0, T[t][r]==0, T[t][r]==-1),
                 If(T[t-1][r]==0, T[t][r]==random.choice(Dest[0]), T[t][r]==-1)))

Here I get the solution as follows:

[[ 0  2 -1 -1]
 [-1 -1 -1 -1]
 [-1 -1 -1 -1]]

However, I try to generalize the constraint to the following,

for t in range(H):
        for r in range(R):
            s.add(If(t==0, If(r==0, T[t][r]==0, T[t][r]==-1),
                     If(T[t-1][r]==0, T[t][r]==random.choice(Dest[T[t-1][r]]), T[t][r]==-1)))

I get the error:

list indices must be integers or slices, not ArithRef

How can this issue be resolved ?


Solution

  • In general, you cannot mix and match indexing into a Python list or calling random on a list when symbolic variables are involved. That is, it is illegal to call A[k] when k is a symbolic integer. That's what the error message is telling you: Until you actually run the solver (via a call to s.check()) and grab the value from a model, you cannot index a list via a symbolic variable. Similarly, calling random.choice isn't going to work either when you need to symbolically pick which list to pick from.

    What you need to do, instead, is to create symbolic index values that correspond to your random values, and do a symbolic walk down the list. Something like this:

    from z3 import *
    
    R = 3
    H = 4
    
    Dest = [[2,3,4], [0,3,2,5], [1,4,5], [1,4,2], [3,1], [2,3,1]]
    
    s = Solver()
    
    T = [[Int('r%d_t%d' % (k,t)) for k in range (R)] for t in range (H)]
    
    # Symbolically walk over a list, grabbing the ith element of it
    # The idea here is that i is not just an integer, but can be a symbolic
    # integer with other constraitns on it. We also take a continuation, k,
    # which we call on the result, for convenience. It's assumed that the
    # caller makes sure lst[i] is always within bounds; i.e., i >= 0, and
    # i < len(lst).
    def SymbolicWalk(i, lst, k):
       if len(lst) == 1:
           return k(lst[0])
       else:
           return If(i == 0, k(lst[0]), SymbolicWalk(i-1, lst[1:], k));
    
    # Pick a random element of the list
    # This is the symbolic version of random.choice
    def SymbolicChoice(lst):
        i = FreshInt('internal_choice')
        s.add(i >= 0)
        s.add(i <  len(lst))
        return SymbolicWalk(i, lst, lambda x: x)
    
    # Grab the element at given symbolic index, and then call the symbolic choice on it
    def index(i, lst):
        return SymbolicWalk(i, lst, SymbolicChoice)
    
    for t in range(H):
        for r in range(R):
            s.add(If(t==0, If(r==0, T[t][r]==0, T[t][r]==-1),
                     If(T[t-1][r]==0, T[t][r]==index(T[t-1][r], Dest), T[t][r]==-1)))
    
    print(s.check())
    print(s.model())
    

    I've added some comments above which I hope should help. But the work-horse here is the function SymbolicWalk which achieves the indexing of a list with a symbolic value. Note the creation of "random"-indexes via the call to FreshInt.

    When I run the above it produces:

    [ r2_t3 = -1,
     r1_t3 = -1,
     r0_t3 = -1,
     r2_t2 = -1,
     r1_t2 = -1,
     r0_t2 = -1,
     r2_t1 = -1,
     r1_t1 = -1,
     r0_t1 = 2,
     r2_t0 = -1,
     r1_t0 = -1,
     r0_t0 = 0]
    

    which should be what you are looking for. Note that I elided the outputs to variables internal_choice!k which are internally generated in calls to FreshInt; you shouldn't refer to those values; they are used internally to get the random value from your Dest list.