Fully Evaluate Nested Arrays for Model Generation in CVC5

src/org/sosy_lab/java_smt/solvers/cvc5/CVC5Model.java

@@ -9,6 +9,7 @@
 package org.sosy_lab.java_smt.solvers.cvc5;
 
 import com.google.common.base.Preconditions;
+import com.google.common.collect.FluentIterable;
 import com.google.common.collect.ImmutableList;
 import com.google.common.collect.ImmutableSet;
 import io.github.cvc5.CVC5ApiException;
@@ -81,7 +82,7 @@ private void recursiveAssignmentFinder(ImmutableSet.Builder<ValueAssignment> bui
         // Vars and UFs, as well as bound vars in UFs!
         // In CVC5 consts are variables! Free variables (in CVC5s notation, we call them bound
         // variables, created with mkVar() can never have a value!)
-        builder.add(getAssignment(expr));
+        builder.addAll(getAssignment(expr));
       } else if (kind == Kind.FORALL || kind == Kind.EXISTS) {
         // Body of the quantifier, with bound vars!
         Term body = expr.getChild(1);
@@ -169,7 +170,95 @@ private ValueAssignment getAssignmentForUf(Term pKeyTerm) {
         keyFormula, valueFormula, equation, nameStr, value, argumentInterpretationBuilder.build());
   }
 
-  private ValueAssignment getAssignment(Term pKeyTerm) {
+  /** Takes a (nested) select statement and returns its indices. */
+  private Iterable<Term> getArgs(Term array) throws CVC5ApiException {
+    if (array.getKind().equals(Kind.SELECT)) {
+      return FluentIterable.concat(getArgs(array.getChild(0)), ImmutableList.of(array.getChild(1)));
+    } else {
+      return ImmutableList.of();
+    }
+  }
+
+  /** Takes a select statement with multiple indices and returns the variable name at the bottom. */
+  private String getVar(Term array) throws CVC5ApiException {
+    if (array.getKind().equals(Kind.SELECT)) {
+      return getVar(array.getChild(0));
+    } else {
+      return array.getSymbol();
+    }
+  }
+
+  /** Build assignment for an array value. */
+  private Iterable<ValueAssignment> buildArrayAssignment(Term expr, Term value) {
+    // CVC5 returns values such as "(Store (Store ... i1,1 e1,1) i1,0 e1,0)" where the i1,x match
+    // the first index of the array and the elements e1,Y can again be arrays (if there is more
+    // than one index). We need "pattern match" this values to extract assignments from it.
+    // Initially we have:
+    //  arr = (Store (Store ... i1,1 e1,1) i1,0 e1,0)
+    // where 'arr' is the name of the variable. By applying (Select i1,0 ...) to both side we get:
+    // (Select i1,0 arr) = (Select i1,0 (Store (Store ... i1,1 e1,1) i1,0 e1,0))
+    // The right side simplifies to e1,0 as the index matches. We need to continue with this for
+    // all other matches to the first index, that is
+    //  (Select i1,1 arr) = (Select i1,0 (Store (Store ... i1,1 e1,1) i1,0 e1,0))
+    //                    = (Select i1,0 (Store ... i1,1 e1,1))
+    //                    = e1,1
+    // until all matches are explored and the bottom of the Store chain is reached. If the array
+    // has more than one dimension we also have to descend into the elements to apply the next
+    // index there. For instance, let's consider a 2-dimensional array with type (Array Int ->
+    // (Array Int -> Int)). After matching the first index just as in the above example we might
+    // have:
+    //  (Select i1,0 arr) = (Select i1,0 (Store (Store ... i1,1 e1,1) i1,0 e1,0)) = e1,0
+    // In this case e1,0 is again a Store chain, let's say e1,0 := (Store (Store ... i2,1 e2,1)
+    // i2,0 e2,0), and we now continue with the second index:
+    //  (Select i2,1 (Select i1,0 arr)) = (Select i2,1 (Store (Store ... i2,1 e2,1) i2,0 e2,0)) =
+    //                                  = e2,1
+    // This again has to be done for all possible matches.
+    // Once we've reached the last index, the successor element will be a non-array value. We
+    // then create the final assignments and return:
+    //  (Select iK,mK ... (Select i2,1 (Select i1,0 arr)) = eik,mK
+    try {
+      if (value.getKind().equals(Kind.STORE)) {
+        // This is a Store node for the current index. We need to follow the chain downwards to
+        // match this index, while also exploring the successor for the other indices
+        Term index = value.getChild(1);
+        Term element = value.getChild(2);
+
+        Term select = creator.getEnv().mkTerm(Kind.SELECT, expr, index);
+
+        Iterable<ValueAssignment> current;
+        if (expr.getSort().getArrayElementSort().isArray()) {
+          current = buildArrayAssignment(select, element);
+        } else {
+          Term equation = creator.getEnv().mkTerm(Kind.EQUAL, select, element);
+
+          current =
+              FluentIterable.of(
+                  new ValueAssignment(
+                      creator.encapsulate(creator.getFormulaType(element), select),
+                      creator.encapsulate(creator.getFormulaType(element), element),
+                      creator.encapsulateBoolean(equation),
+                      getVar(expr),
+                      creator.convertValue(element, element),
+                      FluentIterable.from(getArgs(select)).transform(this::evaluateImpl).toList()));
+        }
+        return FluentIterable.concat(current, buildArrayAssignment(expr, value.getChild(0)));
+
+      } else if (value.getKind().equals(Kind.CONST_ARRAY)) {
+        // We've reached the end of the Store chain
+        return ImmutableList.of();
+
+      } else {
+        // Should be unreachable
+        // We assume that array values are made up of "const" and "store" nodes with non-array
+        // constants as leaves
+        throw new IllegalArgumentException();
+      }
+    } catch (CVC5ApiException e) {
+      throw new RuntimeException(e);
+    }
+  }
+
+  private Iterable<ValueAssignment> getAssignment(Term pKeyTerm) {
     ImmutableList.Builder<Object> argumentInterpretationBuilder = ImmutableList.builder();
     for (int i = 0; i < pKeyTerm.getNumChildren(); i++) {
       try {
@@ -193,13 +282,23 @@ private ValueAssignment getAssignment(Term pKeyTerm) {
     }
 
     Term valueTerm = solver.getValue(pKeyTerm);
-    Formula keyFormula = creator.encapsulateWithTypeOf(pKeyTerm);
-    Formula valueFormula = creator.encapsulateWithTypeOf(valueTerm);
-    BooleanFormula equation =
-        creator.encapsulateBoolean(termManager.mkTerm(Kind.EQUAL, pKeyTerm, valueTerm));
-    Object value = creator.convertValue(pKeyTerm, valueTerm);
-    return new ValueAssignment(
-        keyFormula, valueFormula, equation, nameStr, value, argumentInterpretationBuilder.build());
+    if (valueTerm.getSort().isArray()) {
+      return buildArrayAssignment(pKeyTerm, valueTerm);
+    } else {
+      Formula keyFormula = creator.encapsulateWithTypeOf(pKeyTerm);
+      Formula valueFormula = creator.encapsulateWithTypeOf(valueTerm);
+      BooleanFormula equation =
+          creator.encapsulateBoolean(termManager.mkTerm(Kind.EQUAL, pKeyTerm, valueTerm));
+      Object value = creator.convertValue(pKeyTerm, valueTerm);
+      return ImmutableList.of(
+          new ValueAssignment(
+              keyFormula,
+              valueFormula,
+              equation,
+              nameStr,
+              value,
+              argumentInterpretationBuilder.build()));
+    }
   }
 
   @Override

src/org/sosy_lab/java_smt/test/ModelTest.java

@@ -14,6 +14,7 @@
 import static com.google.common.truth.TruthJUnit.assume;
 import static org.junit.Assert.assertThrows;
 import static org.sosy_lab.java_smt.api.FormulaType.IntegerType;
+import static org.sosy_lab.java_smt.api.FormulaType.getBitvectorTypeWithSize;
 import static org.sosy_lab.java_smt.test.ProverEnvironmentSubject.assertThat;
 
 import com.google.common.collect.ImmutableList;
@@ -79,6 +80,8 @@ public class ModelTest extends SolverBasedTest0.ParameterizedSolverBasedTest0 {
 
   private static final ArrayFormulaType<IntegerFormula, IntegerFormula> ARRAY_TYPE_INT_INT =
       FormulaType.getArrayType(IntegerType, IntegerType);
+  private static final ArrayFormulaType<BitvectorFormula, BitvectorFormula> ARRAY_TYPE_BV32_BV32 =
+      FormulaType.getArrayType(getBitvectorTypeWithSize(32), getBitvectorTypeWithSize(32));
 
   private static final ImmutableList<Solvers> SOLVERS_WITH_PARTIAL_MODEL =
       ImmutableList.of(Solvers.Z3);
@@ -1145,6 +1148,124 @@ public void testGetArrays6() throws SolverException, InterruptedException {
     testModelIterator(f);
   }
 
+  @Test
+  public void testGetArrays3IntegerNoParsing() throws SolverException, InterruptedException {
+    requireIntegers();
+    requireArrays();
+    requireArrayModel();
+
+    assume()
+        .withMessage("Solvers have problems with multi-dimensional arrays")
+        .that(solverToUse())
+        .isNoneOf(Solvers.PRINCESS, Solvers.CVC4, Solvers.YICES2);
+
+    // (= (select (select (select arr 5) 3) 1) x)
+    // (= x 123)"
+    ArrayFormulaType<IntegerFormula, IntegerFormula> innerType =
+        FormulaType.getArrayType(IntegerType, IntegerType);
+    ArrayFormulaType<IntegerFormula, ArrayFormula<IntegerFormula, IntegerFormula>> middleType =
+        FormulaType.getArrayType(IntegerType, innerType);
+    ArrayFormulaType<
+            IntegerFormula,
+            ArrayFormula<IntegerFormula, ArrayFormula<IntegerFormula, IntegerFormula>>>
+        type = FormulaType.getArrayType(IntegerType, middleType);
+
+    IntegerFormula num1 = imgr.makeNumber(1);
+    IntegerFormula num3 = imgr.makeNumber(3);
+    IntegerFormula num5 = imgr.makeNumber(5);
+    IntegerFormula num123 = imgr.makeNumber(123);
+    IntegerFormula x = imgr.makeVariable("x");
+    ArrayFormula<
+            IntegerFormula,
+            ArrayFormula<IntegerFormula, ArrayFormula<IntegerFormula, IntegerFormula>>>
+        arr = amgr.makeArray("arr", type);
+
+    BooleanFormula f =
+        bmgr.and(
+            imgr.equal(x, amgr.select(amgr.select(amgr.select(arr, num5), num3), num1)),
+            imgr.equal(x, num123));
+
+    testModelIterator(f);
+    testModelGetters(f, imgr.makeVariable("x"), BigInteger.valueOf(123), "x");
+    ArrayFormulaType<
+            IntegerFormula,
+            ArrayFormula<IntegerFormula, ArrayFormula<IntegerFormula, IntegerFormula>>>
+        arrType =
+            FormulaType.getArrayType(
+                IntegerType, FormulaType.getArrayType(IntegerType, ARRAY_TYPE_INT_INT));
+    testModelGetters(
+        f,
+        amgr.select(
+            amgr.select(
+                amgr.select(amgr.makeArray("arr", arrType), imgr.makeNumber(5)),
+                imgr.makeNumber(3)),
+            imgr.makeNumber(1)),
+        BigInteger.valueOf(123),
+        "arr",
+        true,
+        ImmutableList.of(BigInteger.valueOf(5), BigInteger.valueOf(3), BigInteger.ONE));
+  }
+
+  @Test
+  public void testGetArrays3BitvectorNoParsing() throws SolverException, InterruptedException {
+    requireBitvectors();
+    requireArrays();
+    requireArrayModel();
+
+    assume()
+        .withMessage("Solvers have problems with multi-dimensional arrays")
+        .that(solverToUse())
+        .isNoneOf(
+            Solvers.PRINCESS, Solvers.CVC4, Solvers.YICES2, Solvers.BOOLECTOR, Solvers.BITWUZLA);
+
+    // (= (select (select (select arr 5) 3) 1) x)
+    // (= x 123)"
+    BitvectorType bvType = getBitvectorTypeWithSize(32);
+    ArrayFormulaType<BitvectorFormula, BitvectorFormula> innerType =
+        FormulaType.getArrayType(bvType, bvType);
+    ArrayFormulaType<BitvectorFormula, ArrayFormula<BitvectorFormula, BitvectorFormula>>
+        middleType = FormulaType.getArrayType(bvType, innerType);
+    ArrayFormulaType<
+            BitvectorFormula,
+            ArrayFormula<BitvectorFormula, ArrayFormula<BitvectorFormula, BitvectorFormula>>>
+        type = FormulaType.getArrayType(bvType, middleType);
+
+    BitvectorFormula num1 = bvmgr.makeBitvector(32, 1);
+    BitvectorFormula num3 = bvmgr.makeBitvector(32, 3);
+    BitvectorFormula num5 = bvmgr.makeBitvector(32, 5);
+    BitvectorFormula num123 = bvmgr.makeBitvector(32, 123);
+    BitvectorFormula x = bvmgr.makeVariable(32, "x");
+    ArrayFormula<
+            BitvectorFormula,
+            ArrayFormula<BitvectorFormula, ArrayFormula<BitvectorFormula, BitvectorFormula>>>
+        arr = amgr.makeArray("arr", type);
+
+    BooleanFormula f =
+        bmgr.and(
+            bvmgr.equal(x, amgr.select(amgr.select(amgr.select(arr, num5), num3), num1)),
+            bvmgr.equal(x, num123));
+
+    testModelIterator(f);
+    testModelGetters(f, bvmgr.makeVariable(32, "x"), BigInteger.valueOf(123), "x");
+    ArrayFormulaType<
+            BitvectorFormula,
+            ArrayFormula<BitvectorFormula, ArrayFormula<BitvectorFormula, BitvectorFormula>>>
+        arrType =
+            FormulaType.getArrayType(
+                bvType, FormulaType.getArrayType(bvType, ARRAY_TYPE_BV32_BV32));
+    testModelGetters(
+        f,
+        amgr.select(
+            amgr.select(
+                amgr.select(amgr.makeArray("arr", arrType), bvmgr.makeBitvector(32, 5)),
+                bvmgr.makeBitvector(32, 3)),
+            bvmgr.makeBitvector(32, 1)),
+        BigInteger.valueOf(123),
+        "arr",
+        true,
+        ImmutableList.of(BigInteger.valueOf(5), BigInteger.valueOf(3), BigInteger.ONE));
+  }
+
   @Test
   public void testGetArrays3() throws SolverException, InterruptedException {
     requireParser();