Add support for indexed functions

src/org/sosy_lab/java_smt/api/FunctionDeclaration.java

@@ -29,6 +29,11 @@ public interface FunctionDeclaration<E extends Formula> {
    */
   String getName();
 
+  /**
+   * @return List of indices for the function
+   */
+  ImmutableList<Integer> getIndices();
+
   /**
    * @return Sort of the function output.
    */

src/org/sosy_lab/java_smt/api/FunctionDeclarationKind.java

@@ -320,5 +320,130 @@ public enum FunctionDeclarationKind {
    * Solvers support a lot of different built-in theories. We enforce standardization only across a
    * small subset.
    */
-  OTHER
+  OTHER;
+
+  public static int getArity(FunctionDeclarationKind pKind) {
+    // TODO Add missing kinds
+    switch (pKind) {
+      case AND:
+      case OR:
+      case IFF:
+      case XOR:
+      case IMPLIES:
+      case DISTINCT:
+      case SUB:
+      case ADD:
+      case DIV:
+      case MUL:
+      case LT:
+      case LTE:
+      case GT:
+      case GTE:
+      case EQ:
+      case BV_CONCAT:
+      case BV_OR:
+      case BV_AND:
+      case BV_XOR:
+      case BV_SUB:
+      case BV_ADD:
+      case BV_MUL:
+        return -1;
+
+      case FP_ROUND_EVEN:
+      case FP_ROUND_AWAY:
+      case FP_ROUND_POSITIVE:
+      case FP_ROUND_NEGATIVE:
+      case FP_ROUND_ZERO:
+        return 0;
+
+      case NOT:
+      case UMINUS:
+      case EQ_ZERO:
+      case GTE_ZERO:
+      case FLOOR:
+      case TO_REAL:
+      case CONST:
+      case BV_EXTRACT:
+      case BV_SIGN_EXTENSION:
+      case BV_ZERO_EXTENSION:
+      case BV_NOT:
+      case BV_NEG:
+      case BV_ROTATE_LEFT_BY_INT:
+      case BV_ROTATE_RIGHT_BY_INT:
+      case FP_NEG:
+      case FP_ABS:
+      case FP_IS_NAN:
+      case FP_IS_INF:
+      case FP_IS_ZERO:
+      case FP_IS_NEGATIVE:
+      case FP_IS_SUBNORMAL:
+      case FP_IS_NORMAL:
+      case FP_AS_IEEEBV:
+      case FP_FROM_IEEEBV:
+        return 1;
+
+      case SELECT:
+      case MODULO:
+      case BV_SDIV:
+      case BV_UDIV:
+      case BV_SREM:
+      case BV_UREM:
+      case BV_SMOD:
+      case BV_ULT:
+      case BV_SLT:
+      case BV_ULE:
+      case BV_SLE:
+      case BV_UGT:
+      case BV_SGT:
+      case BV_UGE:
+      case BV_SGE:
+      case BV_SHL:
+      case BV_LSHR:
+      case BV_ASHR:
+      case BV_ROTATE_LEFT:
+      case BV_ROTATE_RIGHT:
+      case BV_UCASTTO_FP:
+      case BV_SCASTTO_FP:
+      case FP_MAX:
+      case FP_MIN:
+      case FP_SQRT:
+      case FP_REM:
+      case FP_LT:
+      case FP_LE:
+      case FP_GE:
+      case FP_GT:
+      case FP_EQ:
+      case FP_ROUND_TO_INTEGRAL:
+      case FP_CASTTO_FP:
+      case FP_CASTTO_SBV:
+      case FP_CASTTO_UBV:
+        return 2;
+
+      case ITE:
+      case STORE:
+      case FP_SUB:
+      case FP_ADD:
+      case FP_DIV:
+      case FP_MUL:
+        return 3;
+
+      default:
+        throw new IllegalArgumentException(String.format("Unsupported kind: \"%s\"", pKind));
+    }
+  }
+
+  public static int getNumIndices(FunctionDeclarationKind pKind) {
+    // TODO Add missing kinds
+    switch (pKind) {
+      case BV_ROTATE_LEFT_BY_INT:
+      case BV_ROTATE_RIGHT_BY_INT:
+      case BV_SIGN_EXTENSION:
+      case BV_ZERO_EXTENSION:
+        return 1;
+      case BV_EXTRACT:
+        return 2;
+      default:
+        return 0;
+    }
+  }
 }

src/org/sosy_lab/java_smt/basicimpl/FunctionDeclarationImpl.java

@@ -26,11 +26,26 @@ public abstract class FunctionDeclarationImpl<F extends Formula, T>
   public static <F extends Formula, T> FunctionDeclaration<F> of(
       String name,
       FunctionDeclarationKind kind,
+      List<Integer> pIndices,
       List<FormulaType<?>> pArgumentTypes,
       FormulaType<F> pReturnType,
       T pDeclaration) {
     return new AutoValue_FunctionDeclarationImpl<>(
-        kind, name, pReturnType, ImmutableList.copyOf(pArgumentTypes), pDeclaration);
+        kind,
+        name,
+        ImmutableList.copyOf(pIndices),
+        pReturnType,
+        ImmutableList.copyOf(pArgumentTypes),
+        pDeclaration);
+  }
+
+  public static <F extends Formula, T> FunctionDeclaration<F> of(
+      String name,
+      FunctionDeclarationKind kind,
+      List<FormulaType<?>> pArgumentTypes,
+      FormulaType<F> pReturnType,
+      T pDeclaration) {
+    return of(name, kind, ImmutableList.of(), pArgumentTypes, pReturnType, pDeclaration);
   }
 
   /**

src/org/sosy_lab/java_smt/solvers/bitwuzla/BitwuzlaFormulaCreator.java

@@ -484,7 +484,12 @@ public <R> R visit(FormulaVisitor<R> visitor, Formula formula, Term f)
           formula,
           arguments.build(),
           FunctionDeclarationImpl.of(
-              name, getDeclarationKind(f), argumentTypes.build(), getFormulaType(f), decl));
+              name,
+              getDeclarationKind(f),
+              ImmutableList.copyOf(f.indices()),
+              argumentTypes.build(),
+              getFormulaType(f),
+              decl));
     }
   }
 

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

@@ -452,7 +452,17 @@ public <R> R visit(FormulaVisitor<R> visitor, Formula formula, final Term f) {
 
         List<FormulaType<?>> argsTypes = new ArrayList<>();
 
-        // Term operator = normalize(f.getSort());
+        // Collect indices
+        Op operator = f.getOp();
+        ImmutableList.Builder<Integer> indexBuilder = ImmutableList.builder();
+        for (int p = 0; p < operator.getNumIndices(); p++) {
+          Term index = operator.get(p);
+          if (index.isIntegerValue()) {
+            indexBuilder.add(index.getIntegerValue().intValueExact());
+          }
+        }
+
+        // Collect arguments
         Kind kind = f.getKind();
         if (sort.isFunction() || kind == Kind.APPLY_UF) {
           // The arguments are all children except the first one
@@ -468,10 +478,6 @@ public <R> R visit(FormulaVisitor<R> visitor, Formula formula, final Term f) {
           }
         }
 
-        // TODO some operations (BV_SIGN_EXTEND, BV_ZERO_EXTEND, maybe more) encode information as
-        // part of the operator itself, thus the arity is one too small and there might be no
-        // possibility to access the information from user side. Should we encode such information
-        // as additional parameters? We do so for some methods of Princess.
         if (sort.isFunction()) {
           return visitor.visitFunction(
               formula,
@@ -494,7 +500,12 @@ public <R> R visit(FormulaVisitor<R> visitor, Formula formula, final Term f) {
               formula,
               argsBuilder.build(),
               FunctionDeclarationImpl.of(
-                  getName(f), getDeclarationKind(f), argsTypes, getFormulaType(f), normalize(f)));
+                  getName(f),
+                  getDeclarationKind(f),
+                  indexBuilder.build(),
+                  argsTypes,
+                  getFormulaType(f),
+                  normalize(f)));
         }
       }
     } catch (CVC5ApiException e) {

src/org/sosy_lab/java_smt/solvers/mathsat5/Mathsat5FormulaCreator.java

@@ -112,6 +112,7 @@
 import static org.sosy_lab.java_smt.solvers.mathsat5.Mathsat5NativeApi.msat_type_repr;
 
 import com.google.common.base.Preconditions;
+import com.google.common.base.Splitter;
 import com.google.common.collect.ImmutableList;
 import com.google.common.collect.ImmutableSet;
 import com.google.common.primitives.Longs;
@@ -348,6 +349,18 @@ public <R> R visit(FormulaVisitor<R> visitor, Formula formula, final Long f) {
         return visitor.visitFreeVariable(formula, name);
       }
 
+      // Collect indices: We need to extract them from the term name in MathSAT
+      FunctionDeclarationKind kind = getDeclarationKind(f);
+      List<String> tokens =
+          Splitter.on('_')
+              .splitToList(name.startsWith("`") ? name.substring(1, name.length() - 1) : name);
+
+      ImmutableList.Builder<Integer> indices = ImmutableList.builder();
+      for (int p = 1; p <= FunctionDeclarationKind.getNumIndices(kind); p++) {
+        indices.add(Integer.valueOf(tokens.get(p)));
+      }
+
+      // Now collect the arguments
       ImmutableList.Builder<Formula> args = ImmutableList.builder();
       ImmutableList.Builder<FormulaType<?>> argTypes = ImmutableList.builder();
       for (int i = 0; i < arity; i++) {
@@ -365,7 +378,8 @@ public <R> R visit(FormulaVisitor<R> visitor, Formula formula, final Long f) {
           args.build(),
           FunctionDeclarationImpl.of(
               name,
-              getDeclarationKind(f),
+              kind,
+              indices.build(),
               argTypes.build(),
               getFormulaType(f),
               msat_term_get_decl(f)));

src/org/sosy_lab/java_smt/solvers/princess/PrincessFormulaCreator.java

@@ -488,16 +488,13 @@ public <R> R visit(FormulaVisitor<R> visitor, final Formula f, final IExpression
       }
     }
 
-    ImmutableList.Builder<Formula> args = ImmutableList.builder();
-    ImmutableList.Builder<FormulaType<?>> argTypes = ImmutableList.builder();
-    int arity = input.length();
-    int arityStart = 0;
+    int numSortParameters = getNumSortParameters(kind);
+    int numIndices = FunctionDeclarationKind.getNumIndices(kind);
 
     PrincessFunctionDeclaration solverDeclaration;
-    if (isBitvectorOperationWithAdditionalArgument(kind)) {
+    if (numSortParameters > 0) {
       // the first argument is the bitsize, and it is not relevant for the user.
       // we do not want type/sort information as arguments.
-      arityStart = 1;
       if (input instanceof IAtom) {
         solverDeclaration = new PrincessBitvectorToBooleanDeclaration(((IAtom) input).pred());
       } else if (input instanceof IFunApp) {
@@ -518,7 +515,17 @@ public <R> R visit(FormulaVisitor<R> visitor, final Formula f, final IExpression
       solverDeclaration = new PrincessByExampleDeclaration(input);
     }
 
-    for (int i = arityStart; i < arity; i++) {
+    ImmutableList.Builder<Integer> indices = ImmutableList.builder();
+
+    for (int i = numSortParameters; i < numSortParameters + numIndices; i++) {
+      IExpression term = input.apply(i);
+      indices.add(((IIntLit) term).value().intValue());
+    }
+
+    ImmutableList.Builder<Formula> args = ImmutableList.builder();
+    ImmutableList.Builder<FormulaType<?>> argTypes = ImmutableList.builder();
+
+    for (int i = numSortParameters + numIndices; i < input.length(); i++) {
       IExpression arg = input.apply(i);
       FormulaType<?> argumentType = getFormulaType(arg);
       args.add(encapsulate(argumentType, arg));
@@ -529,7 +536,12 @@ public <R> R visit(FormulaVisitor<R> visitor, final Formula f, final IExpression
         f,
         args.build(),
         FunctionDeclarationImpl.of(
-            getName(input), kind, argTypes.build(), getFormulaType(f), solverDeclaration));
+            getName(input),
+            kind,
+            indices.build(),
+            argTypes.build(),
+            getFormulaType(f),
+            solverDeclaration));
   }
 
   private <R> R visitQuantifier(FormulaVisitor<R> visitor, BooleanFormula f, IQuantified input) {
@@ -569,7 +581,8 @@ private String getFreshVariableNameForBody(IFormula body) {
         body, k -> "__JAVASMT__BOUND_VARIABLE_" + boundVariableNames.size());
   }
 
-  private boolean isBitvectorOperationWithAdditionalArgument(FunctionDeclarationKind kind) {
+  private int getNumSortParameters(FunctionDeclarationKind kind) {
+    // TODO Check that his is still right
     switch (kind) {
       case BV_NOT:
       case BV_NEG:
@@ -593,9 +606,14 @@ private boolean isBitvectorOperationWithAdditionalArgument(FunctionDeclarationKi
       case BV_UGE:
       case BV_SGE:
       case BV_EQ:
-        return true;
+      case BV_LSHR:
+      case BV_ASHR:
+      case BV_SHL:
+        return 1;
+      case BV_CONCAT:
+        return 2;
       default:
-        return false;
+        return 0;
     }
   }
 

src/org/sosy_lab/java_smt/solvers/z3/Z3FormulaCreator.java

@@ -22,6 +22,7 @@
 import com.microsoft.z3.Z3Exception;
 import com.microsoft.z3.enumerations.Z3_ast_kind;
 import com.microsoft.z3.enumerations.Z3_decl_kind;
+import com.microsoft.z3.enumerations.Z3_param_kind;
 import com.microsoft.z3.enumerations.Z3_sort_kind;
 import com.microsoft.z3.enumerations.Z3_symbol_kind;
 import java.lang.ref.PhantomReference;
@@ -492,7 +493,8 @@ public <R> R visit(FormulaVisitor<R> visitor, final Formula formula, final Long
         return visitor.visitConstant(formula, convertValue(f));
       case Z3_APP_AST:
         int arity = Native.getAppNumArgs(environment, f);
-        int declKind = Native.getDeclKind(environment, Native.getAppDecl(environment, f));
+        long decl = Native.getAppDecl(environment, f);
+        int declKind = Native.getDeclKind(environment, decl);
 
         if (arity == 0) {
           // constants
@@ -538,6 +540,14 @@ public <R> R visit(FormulaVisitor<R> visitor, final Formula formula, final Long
           }
         }
 
+        ImmutableList.Builder<Integer> indexBuilder = ImmutableList.builder();
+        for (int p = 0; p < Native.getDeclNumParameters(environment, decl); p++) {
+          int kind = Native.getDeclParameterKind(environment, decl, p);
+          if (kind == Z3_param_kind.Z3_PK_UINT.toInt()) {
+            indexBuilder.add(Native.getDeclIntParameter(environment, decl, p));
+          }
+        }
+
         // Function application with zero or more parameters
         ImmutableList.Builder<Formula> args = ImmutableList.builder();
         ImmutableList.Builder<FormulaType<?>> argTypes = ImmutableList.builder();
@@ -553,6 +563,7 @@ public <R> R visit(FormulaVisitor<R> visitor, final Formula formula, final Long
             FunctionDeclarationImpl.of(
                 getAppName(f),
                 getDeclarationKind(f),
+                indexBuilder.build(),
                 argTypes.build(),
                 getFormulaType(f),
                 Native.getAppDecl(environment, f)));