Supporting seq.nth (#4723)

This PR adds support for seq.nth operator by eliminating it during expandDefinitions, based on sub-sequences.
Tests that use this operator are also included.

8 files changed, 106 insertions, 1 deletions

diff --git a/src/theory/strings/kinds b/src/theory/strings/kinds
index 226dcbd17..131a48ae9 100644
--- a/src/theory/strings/kinds
+++ b/src/theory/strings/kinds
@@ -83,6 +83,7 @@ constant CONST_SEQUENCE \
     "a sequence of characters"
 
 operator SEQ_UNIT 1 "a sequence of length one"
+operator SEQ_NTH 2 "The nth element of a sequence"
 
 # equal equal / less than / output
 operator STRING_TO_REGEXP 1 "convert string to regexp"
@@ -169,5 +170,6 @@ typerule STRING_TOLOWER "SimpleTypeRule<RString, AString>"
 
 typerule CONST_SEQUENCE ::CVC4::theory::strings::ConstSequenceTypeRule
 typerule SEQ_UNIT ::CVC4::theory::strings::SeqUnitTypeRule
+typerule SEQ_NTH ::CVC4::theory::strings::SeqNthTypeRule
 
 endtheory
diff --git a/src/theory/strings/rewrites.cpp b/src/theory/strings/rewrites.cpp
index 6ea467ae9..82aa21185 100644
--- a/src/theory/strings/rewrites.cpp
+++ b/src/theory/strings/rewrites.cpp
@@ -210,6 +210,7 @@ const char* toString(Rewrite r)
     case Rewrite::LEN_SEQ_UNIT: return "LEN_SEQ_UNIT";
     case Rewrite::CHARAT_ELIM: return "CHARAT_ELIM";
     case Rewrite::SEQ_UNIT_EVAL: return "SEQ_UNIT_EVAL";
+    case Rewrite::SEQ_NTH_EVAL: return "SEQ_NTH_EVAL";      
     default: return "?";
   }
 }
diff --git a/src/theory/strings/rewrites.h b/src/theory/strings/rewrites.h
index bc5de3a8a..c62129385 100644
--- a/src/theory/strings/rewrites.h
+++ b/src/theory/strings/rewrites.h
@@ -212,7 +212,8 @@ enum class Rewrite : uint32_t
   LEN_CONV_INV,
   LEN_SEQ_UNIT,
   CHARAT_ELIM,
-  SEQ_UNIT_EVAL
+  SEQ_UNIT_EVAL,
+  SEQ_NTH_EVAL
 };
 
 /**
diff --git a/src/theory/strings/sequences_rewriter.cpp b/src/theory/strings/sequences_rewriter.cpp
index 292960e6a..f05c32356 100644
--- a/src/theory/strings/sequences_rewriter.cpp
+++ b/src/theory/strings/sequences_rewriter.cpp
@@ -1488,6 +1488,10 @@ RewriteResponse SequencesRewriter::postRewrite(TNode node)
   {
     retNode = rewriteSeqUnit(node);
   }
+  else if (nk == SEQ_NTH)
+  {
+    retNode = rewriteSeqNth(node);
+  }
 
   Trace("sequences-postrewrite")
       << "Strings::SequencesRewriter::postRewrite returning " << retNode
@@ -1507,6 +1511,33 @@ RewriteResponse SequencesRewriter::preRewrite(TNode node)
   return RewriteResponse(REWRITE_DONE, node);
 }
 
+Node SequencesRewriter::rewriteSeqNth(Node node)
+{
+  Assert(node.getKind() == SEQ_NTH);
+  Node ret;
+  Node s = node[0];
+  Node i = node[1];
+  if (s.isConst() && i.isConst())
+  {
+    size_t len = Word::getLength(s);
+    size_t pos = i.getConst<Rational>().getNumerator().toUnsignedInt();
+    if (pos < len)
+    {
+      std::vector<Node> elements = Word::getChars(s);
+      ret = elements[pos];
+      return returnRewrite(node, ret, Rewrite::SEQ_NTH_EVAL);
+    }
+    else
+    {
+      return node;
+    }
+  }
+  else
+  {
+    return node;
+  }
+}
+
 Node SequencesRewriter::rewriteCharAt(Node node)
 {
   Assert(node.getKind() == STRING_CHARAT);
diff --git a/src/theory/strings/sequences_rewriter.h b/src/theory/strings/sequences_rewriter.h
index 47a20a7ca..ed09397d3 100644
--- a/src/theory/strings/sequences_rewriter.h
+++ b/src/theory/strings/sequences_rewriter.h
@@ -264,6 +264,13 @@ class SequencesRewriter : public TheoryRewriter
    */
   Node rewriteSeqUnit(Node node);
 
+  /** rewrite seq.nth
+   * This is the entry point for post-rewriting terms n of the form
+   *   seq.nth(s, i)
+   * Returns the rewritten form of node.
+   */
+  Node rewriteSeqNth(Node node);
+
   /** length preserving rewrite
    *
    * Given input n, this returns a string n' whose length is equivalent to n.
diff --git a/src/theory/strings/skolem_cache.h b/src/theory/strings/skolem_cache.h
index 302c69e83..755f70833 100644
--- a/src/theory/strings/skolem_cache.h
+++ b/src/theory/strings/skolem_cache.h
@@ -138,6 +138,11 @@ class SkolemCache
     //   where b is a regular expression, n is the number of occurrences of b
     //   in a, and k(0)=0.
     SK_OCCUR_LEN,
+    // For function k: ((Seq U) x Int) -> U
+    // exists k.
+    // forall s, n.
+    //  k(s, n) is some undefined value of sort U
+    SK_NTH,
   };
   /**
    * Returns a skolem of type string that is cached for (a,b,id) and has
diff --git a/src/theory/strings/theory_strings.cpp b/src/theory/strings/theory_strings.cpp
index 1fffd9638..84fd08e7c 100644
--- a/src/theory/strings/theory_strings.cpp
+++ b/src/theory/strings/theory_strings.cpp
@@ -594,6 +594,40 @@ TrustNode TheoryStrings::expandDefinition(Node node)
     return TrustNode::mkTrustRewrite(node, ret, nullptr);
   }
 
+  if (node.getKind() == SEQ_NTH)
+  {
+    // str.nth(s,i) --->
+    //   ite(0<=i<=len(s), witness k. 0<=i<=len(s)->unit(k) = seq.at(s,i),
+    //   uf(s,i))
+    NodeManager* nm = NodeManager::currentNM();
+    Node s = node[0];
+    Node i = node[1];
+    Node len = nm->mkNode(STRING_LENGTH, s);
+    Node cond =
+        nm->mkNode(AND, nm->mkNode(LEQ, d_zero, i), nm->mkNode(LT, i, len));
+    TypeNode elemType = s.getType().getSequenceElementType();
+    Node k = nm->mkBoundVar(elemType);
+    Node bvl = nm->mkNode(BOUND_VAR_LIST, k);
+    std::vector<TypeNode> argTypes;
+    argTypes.push_back(s.getType());
+    argTypes.push_back(nm->integerType());
+    TypeNode ufType = nm->mkFunctionType(argTypes, elemType);
+    SkolemCache* sc = d_termReg.getSkolemCache();
+    Node uf = sc->mkTypedSkolemCached(
+        ufType, Node::null(), Node::null(), SkolemCache::SK_NTH, "Uf");
+    Node ret = nm->mkNode(
+        ITE,
+        cond,
+        nm->mkNode(WITNESS,
+                   bvl,
+                   nm->mkNode(IMPLIES,
+                              cond,
+                              nm->mkNode(SEQ_UNIT, k)
+                                  .eqNode(nm->mkNode(STRING_CHARAT, s, i)))),
+        nm->mkNode(APPLY_UF, uf, s, i));
+    return TrustNode::mkTrustRewrite(node, ret, nullptr);
+  }
+
   return TrustNode::null();
 }
 
diff --git a/src/theory/strings/theory_strings_type_rules.h b/src/theory/strings/theory_strings_type_rules.h
index 12ddb8a3d..b9d0899f2 100644
--- a/src/theory/strings/theory_strings_type_rules.h
+++ b/src/theory/strings/theory_strings_type_rules.h
@@ -376,6 +376,30 @@ class SeqUnitTypeRule
   }
 };
 
+class SeqNthTypeRule
+{
+ public:
+  static TypeNode computeType(NodeManager* nodeManager, TNode n, bool check)
+  {
+    TypeNode t = n[0].getType(check);
+    TypeNode t1 = t.getSequenceElementType();
+    if (check)
+    {
+      if (!t.isSequence())
+      {
+        throw TypeCheckingExceptionPrivate(n, "expecting a sequence in nth");
+      }
+      TypeNode t2 = n[1].getType(check);
+      if (!t2.isInteger())
+      {
+        throw TypeCheckingExceptionPrivate(
+            n, "expecting an integer start term in nth");
+      }
+    }
+    return t1;
+  }
+};
+
 /** Properties of the sequence type */
 struct SequenceProperties
 {