Expand `seq.nth` lazily (#5287)

Our first support for seq.nth eliminated it eagerly during expandDefinitions.
This PR changes that, by eliminating it lazily, as done with other extended string operators.

6 files changed, 57 insertions, 39 deletions

diff --git a/src/theory/strings/extf_solver.cpp b/src/theory/strings/extf_solver.cpp
index 88edaeff7..53cd92ac2 100644
--- a/src/theory/strings/extf_solver.cpp
+++ b/src/theory/strings/extf_solver.cpp
@@ -65,6 +65,7 @@ ExtfSolver::ExtfSolver(SolverState& s,
   d_extt.addFunctionKind(kind::STRING_TOUPPER);
   d_extt.addFunctionKind(kind::STRING_REV);
   d_extt.addFunctionKind(kind::SEQ_UNIT);
+  d_extt.addFunctionKind(kind::SEQ_NTH);
 
   d_true = NodeManager::currentNM()->mkConst(true);
   d_false = NodeManager::currentNM()->mkConst(false);
@@ -188,7 +189,7 @@ bool ExtfSolver::doReduction(int effort, Node n)
     NodeManager* nm = NodeManager::currentNM();
     Assert(k == STRING_SUBSTR || k == STRING_UPDATE || k == STRING_STRCTN
            || k == STRING_STRIDOF || k == STRING_ITOS || k == STRING_STOI
-           || k == STRING_STRREPL || k == STRING_STRREPLALL
+           || k == STRING_STRREPL || k == STRING_STRREPLALL || k == SEQ_NTH
            || k == STRING_REPLACE_RE || k == STRING_REPLACE_RE_ALL
            || k == STRING_LEQ || k == STRING_TOLOWER || k == STRING_TOUPPER
            || k == STRING_REV);
diff --git a/src/theory/strings/sequences_rewriter.cpp b/src/theory/strings/sequences_rewriter.cpp
index 43918d6b1..2cefe6b09 100644
--- a/src/theory/strings/sequences_rewriter.cpp
+++ b/src/theory/strings/sequences_rewriter.cpp
@@ -18,6 +18,7 @@
 
 #include "expr/attribute.h"
 #include "expr/node_builder.h"
+#include "expr/sequence.h"
 #include "theory/rewriter.h"
 #include "theory/strings/arith_entail.h"
 #include "theory/strings/regexp_entail.h"
@@ -1525,7 +1526,7 @@ Node SequencesRewriter::rewriteSeqNth(Node node)
     size_t pos = i.getConst<Rational>().getNumerator().toUnsignedInt();
     if (pos < len)
     {
-      std::vector<Node> elements = Word::getChars(s);
+      std::vector<Node> elements = s.getConst<Sequence>().getVec();
       ret = elements[pos];
       return returnRewrite(node, ret, Rewrite::SEQ_NTH_EVAL);
     }
diff --git a/src/theory/strings/term_registry.cpp b/src/theory/strings/term_registry.cpp
index 71671d8f1..b5b2a5a13 100644
--- a/src/theory/strings/term_registry.cpp
+++ b/src/theory/strings/term_registry.cpp
@@ -141,7 +141,7 @@ void TermRegistry::preRegisterTerm(TNode n)
   {
     if (k == STRING_STRIDOF || k == STRING_ITOS || k == STRING_STOI
         || k == STRING_STRREPL || k == STRING_SUBSTR || k == STRING_STRREPLALL
-        || k == STRING_REPLACE_RE || k == STRING_REPLACE_RE_ALL
+        || k == SEQ_NTH || k == STRING_REPLACE_RE || k == STRING_REPLACE_RE_ALL
         || k == STRING_STRCTN || k == STRING_LEQ || k == STRING_TOLOWER
         || k == STRING_TOUPPER || k == STRING_REV || k == STRING_UPDATE)
     {
diff --git a/src/theory/strings/theory_strings.cpp b/src/theory/strings/theory_strings.cpp
index 0375fd311..2b2c25333 100644
--- a/src/theory/strings/theory_strings.cpp
+++ b/src/theory/strings/theory_strings.cpp
@@ -118,6 +118,8 @@ void TheoryStrings::finishInit()
   d_equalityEngine->addFunctionKind(kind::STRING_IN_REGEXP, eagerEval);
   d_equalityEngine->addFunctionKind(kind::STRING_TO_CODE, eagerEval);
   d_equalityEngine->addFunctionKind(kind::SEQ_UNIT, eagerEval);
+  // `seq.nth` is not always defined, and so we do not evaluate it eagerly.
+  d_equalityEngine->addFunctionKind(kind::SEQ_NTH, false);
   // extended functions
   d_equalityEngine->addFunctionKind(kind::STRING_STRCTN, eagerEval);
   d_equalityEngine->addFunctionKind(kind::STRING_LEQ, eagerEval);
@@ -578,41 +580,6 @@ TrustNode TheoryStrings::expandDefinition(Node node)
             ITE, cond, t.eqNode(nm->mkNode(STRING_TO_CODE, k)), k.eqNode(emp)));
     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_preprocess.cpp b/src/theory/strings/theory_strings_preprocess.cpp
index 966964bc8..a8c1a6573 100644
--- a/src/theory/strings/theory_strings_preprocess.cpp
+++ b/src/theory/strings/theory_strings_preprocess.cpp
@@ -418,6 +418,55 @@ Node StringsPreprocess::reduce(Node t,
 
     retNode = stoit;
   }
+  else if (t.getKind() == kind::SEQ_NTH)
+  {
+    // processing term:  str.nth( s, n)
+    // similar to substr.
+    Node s = t[0];
+    Node n = t[1];
+    Node skt = sc->mkSkolemCached(t, SkolemCache::SK_PURIFY, "sst");
+    Node t12 = nm->mkNode(PLUS, n, one);
+    Node lt0 = nm->mkNode(STRING_LENGTH, s);
+    // start point is greater than or equal zero
+    Node c1 = nm->mkNode(GEQ, n, zero);
+    // start point is less than end of string
+    Node c2 = nm->mkNode(GT, lt0, n);
+    // check whether this application of seq.nth is defined.
+    Node cond = nm->mkNode(AND, c1, c2);
+
+    // nodes for the case where `seq.nth` is defined.
+    Node sk1 = sc->mkSkolemCached(s, n, SkolemCache::SK_PREFIX, "sspre");
+    Node sk2 = sc->mkSkolemCached(s, t12, SkolemCache::SK_SUFFIX_REM, "sssufr");
+    Node unit = nm->mkNode(SEQ_UNIT, skt);
+    Node b11 = s.eqNode(nm->mkNode(STRING_CONCAT, sk1, unit, sk2));
+    // length of first skolem is second argument
+    Node b12 = nm->mkNode(STRING_LENGTH, sk1).eqNode(n);
+    Node lsk2 = nm->mkNode(STRING_LENGTH, sk2);
+    Node b13 = nm->mkNode(EQUAL, lsk2, nm->mkNode(MINUS, lt0, t12));
+    Node b1 = nm->mkNode(AND, b11, b12, b13);
+
+    // nodes for the case where `seq.nth` is undefined.
+    std::vector<TypeNode> argTypes;
+    argTypes.push_back(s.getType());
+    argTypes.push_back(nm->integerType());
+    TypeNode elemType = s.getType().getSequenceElementType();
+    TypeNode ufType = nm->mkFunctionType(argTypes, elemType);
+    Node uf = sc->mkTypedSkolemCached(
+        ufType, Node::null(), Node::null(), SkolemCache::SK_NTH, "Uf");
+    Node b2 = nm->mkNode(EQUAL, skt, nm->mkNode(APPLY_UF, uf, s, n));
+
+    // the full ite, split on definedness of `seq.nth`
+    Node lemma = nm->mkNode(ITE, cond, b1, b2);
+
+    // assert:
+    // IF    n >=0 AND n < len( s )
+    // THEN: s = sk1 ++ unit(skt) ++ sk2 AND
+    //       len( sk1 ) = n AND
+    //       ( len( sk2 ) = len( s )- (n+1)
+    // ELSE: skt = Uf(s, n), where Uf is a cached skolem function.
+    asserts.push_back(lemma);
+    retNode = skt;
+  }
   else if (t.getKind() == kind::STRING_STRREPL)
   {
     // processing term: replace( x, y, z )
diff --git a/src/theory/strings/theory_strings_utils.cpp b/src/theory/strings/theory_strings_utils.cpp
index d48f081e1..286c0dc04 100644
--- a/src/theory/strings/theory_strings_utils.cpp
+++ b/src/theory/strings/theory_strings_utils.cpp
@@ -382,7 +382,7 @@ TypeNode getOwnerStringType(Node n)
   TypeNode tn;
   Kind k = n.getKind();
   if (k == STRING_STRIDOF || k == STRING_LENGTH || k == STRING_STRCTN
-      || k == STRING_PREFIX || k == STRING_SUFFIX)
+      || k == SEQ_NTH || k == STRING_PREFIX || k == STRING_SUFFIX)
   {
     // owning string type is the type of first argument
     tn = n[0].getType();