Add getNumIndices to Op (#6386)

Add getNumIndices to Op

3 files changed, 102 insertions, 0 deletions

diff --git a/src/api/cpp/cvc5.cpp b/src/api/cpp/cvc5.cpp
index aafc5331a..9fb008bfd 100644
--- a/src/api/cpp/cvc5.cpp
+++ b/src/api/cpp/cvc5.cpp
@@ -428,6 +428,7 @@ const static std::unordered_map<cvc5::Kind, Kind, cvc5::kind::KindHashFunction>
         {cvc5::Kind::TO_INTEGER, TO_INTEGER},
         {cvc5::Kind::TO_REAL, TO_REAL},
         {cvc5::Kind::PI, PI},
+        {cvc5::Kind::IAND_OP, IAND},
         /* BV -------------------------------------------------------------- */
         {cvc5::Kind::CONST_BITVECTOR, CONST_BITVECTOR},
         {cvc5::Kind::BITVECTOR_CONCAT, BITVECTOR_CONCAT},
@@ -1842,6 +1843,52 @@ bool Op::isIndexed() const
   CVC5_API_TRY_CATCH_END;
 }
 
+size_t Op::getNumIndices() const
+{
+  CVC5_API_TRY_CATCH_BEGIN;
+  CVC5_API_CHECK_NOT_NULL;
+  if (!isIndexedHelper())
+  {
+    return 0;
+  }
+
+  Kind k = intToExtKind(d_node->getKind());
+  size_t size = 0;
+  switch (k)
+  {
+    case DIVISIBLE: size = 1; break;
+    case RECORD_UPDATE: size = 1; break;
+    case BITVECTOR_REPEAT: size = 1; break;
+    case BITVECTOR_ZERO_EXTEND: size = 1; break;
+    case BITVECTOR_SIGN_EXTEND: size = 1; break;
+    case BITVECTOR_ROTATE_LEFT: size = 1; break;
+    case BITVECTOR_ROTATE_RIGHT: size = 1; break;
+    case INT_TO_BITVECTOR: size = 1; break;
+    case IAND: size = 1; break;
+    case FLOATINGPOINT_TO_UBV: size = 1; break;
+    case FLOATINGPOINT_TO_SBV: size = 1; break;
+    case TUPLE_UPDATE: size = 1; break;
+    case REGEXP_REPEAT: size = 1; break;
+    case BITVECTOR_EXTRACT: size = 2; break;
+    case FLOATINGPOINT_TO_FP_IEEE_BITVECTOR: size = 2; break;
+    case FLOATINGPOINT_TO_FP_FLOATINGPOINT: size = 2; break;
+    case FLOATINGPOINT_TO_FP_REAL: size = 2; break;
+    case FLOATINGPOINT_TO_FP_SIGNED_BITVECTOR: size = 2; break;
+    case FLOATINGPOINT_TO_FP_UNSIGNED_BITVECTOR: size = 2; break;
+    case FLOATINGPOINT_TO_FP_GENERIC: size = 2; break;
+    case REGEXP_LOOP: size = 2; break;
+    case TUPLE_PROJECT:
+      size = d_node->getConst<TupleProjectOp>().getIndices().size();
+      break;
+    default: CVC5_API_CHECK(false) << "Unhandled kind " << kindToString(k);
+  }
+
+  //////// all checks before this line
+  return size;
+  ////////
+  CVC5_API_TRY_CATCH_END;
+}
+
 template <>
 std::string Op::getIndices() const
 {
diff --git a/src/api/cpp/cvc5.h b/src/api/cpp/cvc5.h
index aaf19c30b..c8f1b8d4a 100644
--- a/src/api/cpp/cvc5.h
+++ b/src/api/cpp/cvc5.h
@@ -819,6 +819,11 @@ class CVC4_EXPORT Op
   bool isIndexed() const;
 
   /**
+   * @return the number of indices of this op
+   */
+  size_t getNumIndices() const;
+
+  /**
    * Get the indices used to create this Op.
    * Supports the following template arguments:
    *   - string
diff --git a/test/unit/api/op_black.cpp b/test/unit/api/op_black.cpp
index e7b83f455..7183fa7a7 100644
--- a/test/unit/api/op_black.cpp
+++ b/test/unit/api/op_black.cpp
@@ -46,6 +46,56 @@ TEST_F(TestApiBlackOp, opFromKind)
   ASSERT_THROW(d_solver.mkOp(BITVECTOR_EXTRACT), CVC5ApiException);
 }
 
+TEST_F(TestApiBlackOp, getNumIndices)
+{
+  Op plus = d_solver.mkOp(PLUS);
+  Op divisible = d_solver.mkOp(DIVISIBLE, 4);
+  Op record_update = d_solver.mkOp(RECORD_UPDATE, "test");
+  Op bitvector_repeat = d_solver.mkOp(BITVECTOR_REPEAT, 5);
+  Op bitvector_zero_extend = d_solver.mkOp(BITVECTOR_ZERO_EXTEND, 6);
+  Op bitvector_sign_extend = d_solver.mkOp(BITVECTOR_SIGN_EXTEND, 7);
+  Op bitvector_rotate_left = d_solver.mkOp(BITVECTOR_ROTATE_LEFT, 8);
+  Op bitvector_rotate_right = d_solver.mkOp(BITVECTOR_ROTATE_RIGHT, 9);
+  Op int_to_bitvector = d_solver.mkOp(INT_TO_BITVECTOR, 10);
+  Op iand = d_solver.mkOp(IAND, 3);
+  Op floatingpoint_to_ubv = d_solver.mkOp(FLOATINGPOINT_TO_UBV, 11);
+  Op floatingopint_to_sbv = d_solver.mkOp(FLOATINGPOINT_TO_SBV, 13);
+  Op tuple_update = d_solver.mkOp(TUPLE_UPDATE, 5);
+  Op floatingpoint_to_fp_ieee_bitvector =
+      d_solver.mkOp(FLOATINGPOINT_TO_FP_IEEE_BITVECTOR, 4, 25);
+  Op floatingpoint_to_fp_floatingpoint =
+      d_solver.mkOp(FLOATINGPOINT_TO_FP_FLOATINGPOINT, 4, 25);
+  Op floatingpoint_to_fp_real = d_solver.mkOp(FLOATINGPOINT_TO_FP_REAL, 4, 25);
+  Op floatingpoint_to_fp_signed_bitvector =
+      d_solver.mkOp(FLOATINGPOINT_TO_FP_SIGNED_BITVECTOR, 4, 25);
+  Op floatingpoint_to_fp_unsigned_bitvector =
+      d_solver.mkOp(FLOATINGPOINT_TO_FP_UNSIGNED_BITVECTOR, 4, 25);
+  Op floatingpoint_to_fp_generic =
+      d_solver.mkOp(FLOATINGPOINT_TO_FP_GENERIC, 4, 25);
+  Op regexp_loop = d_solver.mkOp(REGEXP_LOOP, 2, 3);
+
+  ASSERT_EQ(0, plus.getNumIndices());
+  ASSERT_EQ(1, divisible.getNumIndices());
+  ASSERT_EQ(1, record_update.getNumIndices());
+  ASSERT_EQ(1, bitvector_repeat.getNumIndices());
+  ASSERT_EQ(1, bitvector_zero_extend.getNumIndices());
+  ASSERT_EQ(1, bitvector_sign_extend.getNumIndices());
+  ASSERT_EQ(1, bitvector_rotate_left.getNumIndices());
+  ASSERT_EQ(1, bitvector_rotate_right.getNumIndices());
+  ASSERT_EQ(1, int_to_bitvector.getNumIndices());
+  ASSERT_EQ(1, iand.getNumIndices());
+  ASSERT_EQ(1, floatingpoint_to_ubv.getNumIndices());
+  ASSERT_EQ(1, floatingopint_to_sbv.getNumIndices());
+  ASSERT_EQ(1, tuple_update.getNumIndices());
+  ASSERT_EQ(2, floatingpoint_to_fp_ieee_bitvector.getNumIndices());
+  ASSERT_EQ(2, floatingpoint_to_fp_floatingpoint.getNumIndices());
+  ASSERT_EQ(2, floatingpoint_to_fp_real.getNumIndices());
+  ASSERT_EQ(2, floatingpoint_to_fp_signed_bitvector.getNumIndices());
+  ASSERT_EQ(2, floatingpoint_to_fp_unsigned_bitvector.getNumIndices());
+  ASSERT_EQ(2, floatingpoint_to_fp_generic.getNumIndices());
+  ASSERT_EQ(2, regexp_loop.getNumIndices());
+}
+
 TEST_F(TestApiBlackOp, getIndicesString)
 {
   Op x;