Add rewrite for solving stoi (#2532)

4 files changed, 179 insertions, 9 deletions

diff --git a/src/theory/strings/theory_strings.cpp b/src/theory/strings/theory_strings.cpp
index 1acb5ec95..bc9c19815 100644
--- a/src/theory/strings/theory_strings.cpp
+++ b/src/theory/strings/theory_strings.cpp
@@ -1659,12 +1659,30 @@ void TheoryStrings::checkExtfInference( Node n, Node nr, ExtfInfoTmp& in, int ef
     return;
   }
   NodeManager* nm = NodeManager::currentNM();
-
-  Node exp = n.eqNode(in.d_const);
-  exp = Rewriter::rewrite(exp);
+  Trace("strings-extf-infer") << "checkExtfInference: " << n << " : " << nr
+                              << " == " << in.d_const << std::endl;
 
   // add original to explanation
-  in.d_exp.push_back(exp);
+  if (n.getType().isBoolean())
+  {
+    // if Boolean, it's easy
+    in.d_exp.push_back(in.d_const.getConst<bool>() ? n : n.negate());
+  }
+  else
+  {
+    // otherwise, must explain via base node
+    Node r = getRepresentative(n);
+    // we have that:
+    //   d_eqc_to_const_exp[r] => d_eqc_to_const_base[r] = in.d_const
+    // thus:
+    //   n = d_eqc_to_const_base[r] ^ d_eqc_to_const_exp[r] => n = in.d_const
+    Assert(d_eqc_to_const_base.find(r) != d_eqc_to_const_base.end());
+    addToExplanation(n, d_eqc_to_const_base[r], in.d_exp);
+    Assert(d_eqc_to_const_exp.find(r) != d_eqc_to_const_exp.end());
+    in.d_exp.insert(in.d_exp.end(),
+                    d_eqc_to_const_exp[r].begin(),
+                    d_eqc_to_const_exp[r].end());
+  }
 
   // d_extf_infer_cache stores whether we have made the inferences associated
   // with a node n,
@@ -1793,6 +1811,25 @@ void TheoryStrings::checkExtfInference( Node n, Node nr, ExtfInfoTmp& in, int ef
         getExtTheory()->markReduced(n);
       }
     }
+    return;
+  }
+
+  // If it's not a predicate, see if we can solve the equality n = c, where c
+  // is the constant that extended term n is equal to.
+  Node inferEq = nr.eqNode(in.d_const);
+  Node inferEqr = Rewriter::rewrite(inferEq);
+  Node inferEqrr = inferEqr;
+  if (inferEqr.getKind() == EQUAL)
+  {
+    // try to use the extended rewriter for equalities
+    inferEqrr = TheoryStringsRewriter::rewriteEqualityExt(inferEqr);
+  }
+  if (inferEqrr != inferEqr)
+  {
+    inferEqrr = Rewriter::rewrite(inferEqrr);
+    Trace("strings-extf-infer") << "checkExtfInference: " << inferEq
+                                << " ...reduces to " << inferEqrr << std::endl;
+    sendInternalInference(in.d_exp, inferEqrr, "EXTF_equality_rew");
   }
 }
 
@@ -3809,6 +3846,58 @@ void TheoryStrings::registerTerm( Node n, int effort ) {
   }
 }
 
+void TheoryStrings::sendInternalInference(std::vector<Node>& exp,
+                                          Node conc,
+                                          const char* c)
+{
+  if (conc.getKind() == AND)
+  {
+    for (const Node& cc : conc)
+    {
+      sendInternalInference(exp, cc, c);
+    }
+    return;
+  }
+  bool pol = conc.getKind() != NOT;
+  Node lit = pol ? conc : conc[0];
+  if (lit.getKind() == EQUAL)
+  {
+    for (unsigned i = 0; i < 2; i++)
+    {
+      if (!lit[i].isConst() && !hasTerm(lit[i]))
+      {
+        // introduces a new non-constant term, do not infer
+        return;
+      }
+    }
+    // does it already hold?
+    if (pol ? areEqual(lit[0], lit[1]) : areDisequal(lit[0], lit[1]))
+    {
+      return;
+    }
+  }
+  else if (lit.isConst())
+  {
+    if (lit.getConst<bool>())
+    {
+      Assert(pol);
+      // trivially holds
+      return;
+    }
+  }
+  else if (!hasTerm(lit))
+  {
+    // introduces a new non-constant term, do not infer
+    return;
+  }
+  else if (areEqual(lit, pol ? d_true : d_false))
+  {
+    // already holds
+    return;
+  }
+  sendInference(exp, conc, c);
+}
+
 void TheoryStrings::sendInference( std::vector< Node >& exp, std::vector< Node >& exp_n, Node eq, const char * c, bool asLemma ) {
   eq = eq.isNull() ? d_false : Rewriter::rewrite( eq );
   if( eq!=d_true ){
diff --git a/src/theory/strings/theory_strings.h b/src/theory/strings/theory_strings.h
index 2fd7b3525..6c653bf05 100644
--- a/src/theory/strings/theory_strings.h
+++ b/src/theory/strings/theory_strings.h
@@ -580,7 +580,12 @@ private:
   void doPendingFacts();
   void doPendingLemmas();
   bool hasProcessed();
+  /**
+   * Adds equality a = b to the vector exp if a and b are distinct terms. It
+   * must be the case that areEqual( a, b ) holds in this context.
+   */
   void addToExplanation(Node a, Node b, std::vector<Node>& exp);
+  /** Adds lit to the vector exp if it is non-null */
   void addToExplanation(Node lit, std::vector<Node>& exp);
 
   /** Register term
@@ -602,6 +607,22 @@ private:
    * effort, the call to this method does nothing.
    */
   void registerTerm(Node n, int effort);
+  //-------------------------------------send inferences
+  /** send internal inferences
+   *
+   * This is called when we have inferred exp => conc, where exp is a set
+   * of equalities and disequalities that hold in the current equality engine.
+   * This method adds equalities and disequalities ~( s = t ) via
+   * sendInference such that both s and t are either constants or terms
+   * that already occur in the equality engine, and ~( s = t ) is a consequence
+   * of conc. This function can be seen as a "conservative" version of
+   * sendInference below in that it does not introduce any new non-constant
+   * terms to the state.
+   *
+   * The argument c is a string identifying the reason for the interference.
+   * This string is used for debugging purposes.
+   */
+  void sendInternalInference(std::vector<Node>& exp, Node conc, const char* c);
   // send lemma
   void sendInference(std::vector<Node>& exp,
                      std::vector<Node>& exp_n,
@@ -615,6 +636,7 @@ private:
   void sendLemma(Node ant, Node conc, const char* c);
   void sendInfer(Node eq_exp, Node eq, const char* c);
   bool sendSplit(Node a, Node b, const char* c, bool preq = true);
+  //-------------------------------------end send inferences
 
   /** mkConcat **/
   inline Node mkConcat(Node n1, Node n2);
diff --git a/src/theory/strings/theory_strings_rewriter.cpp b/src/theory/strings/theory_strings_rewriter.cpp
index 9813848a1..de5effd23 100644
--- a/src/theory/strings/theory_strings_rewriter.cpp
+++ b/src/theory/strings/theory_strings_rewriter.cpp
@@ -321,10 +321,21 @@ Node TheoryStringsRewriter::rewriteEquality(Node node)
 Node TheoryStringsRewriter::rewriteEqualityExt(Node node)
 {
   Assert(node.getKind() == EQUAL);
-  if (!node[0].getType().isString())
+  if (node[0].getType().isInteger())
   {
-    return node;
+    return rewriteArithEqualityExt(node);
   }
+  if (node[0].getType().isString())
+  {
+    return rewriteStrEqualityExt(node);
+  }
+  return node;
+}
+
+Node TheoryStringsRewriter::rewriteStrEqualityExt(Node node)
+{
+  Assert(node.getKind() == EQUAL && node[0].getType().isString());
+
   NodeManager* nm = NodeManager::currentNM();
   std::vector<Node> c[2];
   Node new_ret;
@@ -577,6 +588,38 @@ Node TheoryStringsRewriter::rewriteEqualityExt(Node node)
   return node;
 }
 
+Node TheoryStringsRewriter::rewriteArithEqualityExt(Node node)
+{
+  Assert(node.getKind() == EQUAL && node[0].getType().isInteger());
+
+  NodeManager* nm = NodeManager::currentNM();
+
+  // cases where we can solve the equality
+  for (unsigned i = 0; i < 2; i++)
+  {
+    if (node[i].isConst())
+    {
+      Node on = node[1 - i];
+      Kind onk = on.getKind();
+      if (onk == STRING_STOI)
+      {
+        Rational r = node[i].getConst<Rational>();
+        int sgn = r.sgn();
+        Node onEq;
+        std::stringstream ss;
+        if (sgn >= 0)
+        {
+          ss << r.getNumerator();
+        }
+        Node new_ret = on[0].eqNode(nm->mkConst(String(ss.str())));
+        return returnRewrite(node, new_ret, "stoi-solve");
+      }
+    }
+  }
+
+  return node;
+}
+
 // TODO (#1180) add rewrite
 //  str.++( str.substr( x, n1, n2 ), str.substr( x, n1+n2, n3 ) ) --->
 //  str.substr( x, n1, n2+n3 )
diff --git a/src/theory/strings/theory_strings_rewriter.h b/src/theory/strings/theory_strings_rewriter.h
index c0aa91360..91d87769c 100644
--- a/src/theory/strings/theory_strings_rewriter.h
+++ b/src/theory/strings/theory_strings_rewriter.h
@@ -98,6 +98,19 @@ class TheoryStringsRewriter {
    * a is in rewritten form.
    */
   static bool checkEntailArithInternal(Node a);
+  /** rewrite string equality extended
+   *
+   * This method returns a formula that is equivalent to the equality between
+   * two strings s = t, given by node. It is called by rewriteEqualityExt.
+   */
+  static Node rewriteStrEqualityExt(Node node);
+  /** rewrite arithmetic equality extended
+   *
+   * This method returns a formula that is equivalent to the equality between
+   * two arithmetic string terms s = t, given by node. t is called by
+   * rewriteEqualityExt.
+   */
+  static Node rewriteArithEqualityExt(Node node);
   /**
    * Called when node rewrites to ret.
    *
@@ -129,9 +142,12 @@ class TheoryStringsRewriter {
   /** rewrite equality extended
    *
    * This method returns a formula that is equivalent to the equality between
-   * two strings s = t, given by node. Specifically, this function performs
-   * rewrites whose conclusion is not necessarily one of
-   * { s = t, t = s, true, false }.
+   * two terms s = t, given by node, where s and t are terms in the signature
+   * of the theory of strings. Notice that s and t may be of string type or
+   * of Int type.
+   *
+   * Specifically, this function performs rewrites whose conclusion is not
+   * necessarily one of { s = t, t = s, true, false }.
    */
   static Node rewriteEqualityExt(Node node);
   /** rewrite concat