Incorporate rewriting on demand in the evaluator (#3549)

2 files changed, 168 insertions, 26 deletions

diff --git a/src/theory/evaluator.cpp b/src/theory/evaluator.cpp
index ffb5b9fd4..ae573956f 100644
--- a/src/theory/evaluator.cpp
+++ b/src/theory/evaluator.cpp
@@ -17,6 +17,7 @@
 #include "theory/evaluator.h"
 
 #include "theory/bv/theory_bv_utils.h"
+#include "theory/rewriter.h"
 #include "theory/theory.h"
 #include "util/integer.h"
 
@@ -119,16 +120,33 @@ Node Evaluator::eval(TNode n,
 {
   Trace("evaluator") << "Evaluating " << n << " under substitution " << args
                      << " " << vals << std::endl;
-  return evalInternal(n, args, vals).toNode();
+  std::unordered_map<TNode, Node, NodeHashFunction> evalAsNode;
+  Node ret = evalInternal(n, args, vals, evalAsNode).toNode();
+  if (!ret.isNull())
+  {
+    // maybe it was stored in the evaluation-as-node map
+    std::unordered_map<TNode, Node, NodeHashFunction>::iterator itn =
+        evalAsNode.find(n);
+    if (itn != evalAsNode.end())
+    {
+      return itn->second;
+    }
+  }
+  return ret;
 }
 
-EvalResult Evaluator::evalInternal(TNode n,
-                                   const std::vector<Node>& args,
-                                   const std::vector<Node>& vals)
+EvalResult Evaluator::evalInternal(
+    TNode n,
+    const std::vector<Node>& args,
+    const std::vector<Node>& vals,
+    std::unordered_map<TNode, Node, NodeHashFunction>& evalAsNode)
 {
   std::unordered_map<TNode, EvalResult, TNodeHashFunction> results;
   std::vector<TNode> queue;
   queue.emplace_back(n);
+  std::unordered_map<TNode, Node, NodeHashFunction>::iterator itn;
+  std::unordered_map<TNode, EvalResult, TNodeHashFunction>::iterator itr;
+  NodeManager* nm = NodeManager::currentNM();
 
   while (queue.size() != 0)
   {
@@ -140,30 +158,98 @@ EvalResult Evaluator::evalInternal(TNode n,
       continue;
     }
 
+    bool doProcess = true;
     bool doEval = true;
     for (const auto& currNodeChild : currNode)
     {
-      if (results.find(currNodeChild) == results.end())
+      itr = results.find(currNodeChild);
+      if (itr == results.end())
       {
         queue.emplace_back(currNodeChild);
+        doProcess = false;
+      }
+      else if (itr->second.d_tag == EvalResult::INVALID)
+      {
+        // we cannot evaluate since there was an invalid child
         doEval = false;
       }
     }
+    Trace("evaluator") << "Evaluator: visit " << currNode
+                       << ", process = " << doProcess
+                       << ", evaluate = " << doEval << std::endl;
 
-    if (doEval)
+    if (doProcess)
     {
       queue.pop_back();
 
       Node currNodeVal = currNode;
+
+      // The code below should either:
+      // (1) store a valid EvalResult into results[currNode], or
+      // (2) store an invalid EvalResult into results[currNode] and
+      // store the result of substitution + rewriting currNode { args -> vals }
+      // into evalAsNode[currNode].
+
+      // If we did not successfully evaluate all children
+      if (!doEval)
+      {
+        // Reconstruct the node with a combination of the children that
+        // successfully evaluated, and the children that did not.
+        Trace("evaluator") << "Evaluator: collect arguments" << std::endl;
+        std::vector<Node> echildren;
+        if (currNode.getMetaKind() == kind::metakind::PARAMETERIZED)
+        {
+          echildren.push_back(currNode.getOperator());
+        }
+        for (const auto& currNodeChild : currNode)
+        {
+          itr = results.find(currNodeChild);
+          if (itr->second.d_tag == EvalResult::INVALID)
+          {
+            // could not evaluate this child, look in the node cache
+            itn = evalAsNode.find(currNodeChild);
+            Assert(itn != evalAsNode.end());
+            echildren.push_back(itn->second);
+          }
+          else
+          {
+            // otherwise, use the evaluation
+            echildren.push_back(itr->second.toNode());
+          }
+        }
+        // The value is the result of our (partially) successful evaluation
+        // of the children.
+        currNodeVal = nm->mkNode(currNode.getKind(), echildren);
+        Trace("evaluator") << "Evaluator: partially evaluated " << currNodeVal
+                           << std::endl;
+        // Use rewriting. Notice we do not need to substitute here since
+        // all substitutions should already have been applied recursively.
+        currNodeVal = Rewriter::rewrite(currNodeVal);
+        Trace("evaluator") << "Evaluator: now after substitution + rewriting: "
+                           << currNodeVal << std::endl;
+        if (currNodeVal.getNumChildren() > 0)
+        {
+          // We may continue with a valid EvalResult at this point only if
+          // we have no children. We must otherwise fail here since some of
+          // our children may not have successful evaluations.
+          results[currNode] = EvalResult();
+          evalAsNode[currNode] = currNodeVal;
+          continue;
+        }
+        // Otherwise, we may be able to turn the overall result into an
+        // valid EvalResult and continue. We fallthrough and continue with the
+        // block of code below.
+      }
+
       if (currNode.isVar())
       {
         const auto& it = std::find(args.begin(), args.end(), currNode);
-
         if (it == args.end())
         {
-          return EvalResult();
+          evalAsNode[currNode] = currNode;
+          results[currNode] = EvalResult();
+          continue;
         }
-
         ptrdiff_t pos = std::distance(args.begin(), it);
         currNodeVal = vals[pos];
       }
@@ -189,12 +275,17 @@ EvalResult Evaluator::evalInternal(TNode n,
         }
 
         // Lambdas are evaluated in a recursive fashion because each evaluation
-        // requires different substitutions
-        results[currNode] = evalInternal(op[1], lambdaArgs, lambdaVals);
+        // requires different substitutions. We use a fresh cache since the
+        // evaluation of op[1] is under a new substitution and thus should not
+        // be cached. We could alternatively copy evalAsNode to evalAsNodeC but
+        // favor avoiding this copy for performance reasons.
+        std::unordered_map<TNode, Node, NodeHashFunction> evalAsNodeC;
+        results[currNode] =
+            evalInternal(op[1], lambdaArgs, lambdaVals, evalAsNodeC);
         if (results[currNode].d_tag == EvalResult::INVALID)
         {
-          // evaluation was invalid, we fail
-          return results[currNode];
+          // evaluation was invalid, we take the node of op[1] as the result
+          evalAsNode[currNode] = evalAsNodeC[op[1]];
         }
         continue;
       }
@@ -304,7 +395,12 @@ EvalResult Evaluator::evalInternal(TNode n,
           results[currNode] = EvalResult(x < y);
           break;
         }
-
+        case kind::ABS:
+        {
+          const Rational& x = results[currNode[0]].d_rat;
+          results[currNode] = EvalResult(x.abs());
+          break;
+        }
         case kind::CONST_STRING:
           results[currNode] = EvalResult(currNodeVal.getConst<String>());
           break;
@@ -568,6 +664,40 @@ EvalResult Evaluator::evalInternal(TNode n,
           results[currNode] = EvalResult(res);
           break;
         }
+        case kind::BITVECTOR_UDIV:
+        case kind::BITVECTOR_UDIV_TOTAL:
+        {
+          if (currNodeVal.getKind() == kind::BITVECTOR_UDIV_TOTAL
+              || results[currNode[1]].d_bv.getValue() != 0)
+          {
+            BitVector res = results[currNode[0]].d_bv;
+            res = res.unsignedDivTotal(results[currNode[1]].d_bv);
+            results[currNode] = EvalResult(res);
+          }
+          else
+          {
+            results[currNode] = EvalResult();
+            evalAsNode[currNode] = currNodeVal;
+          }
+          break;
+        }
+        case kind::BITVECTOR_UREM:
+        case kind::BITVECTOR_UREM_TOTAL:
+        {
+          if (currNodeVal.getKind() == kind::BITVECTOR_UREM_TOTAL
+              || results[currNode[1]].d_bv.getValue() != 0)
+          {
+            BitVector res = results[currNode[0]].d_bv;
+            res = res.unsignedRemTotal(results[currNode[1]].d_bv);
+            results[currNode] = EvalResult(res);
+          }
+          else
+          {
+            results[currNode] = EvalResult();
+            evalAsNode[currNode] = currNodeVal;
+          }
+          break;
+        }
 
         case kind::EQUAL:
         {
@@ -604,7 +734,8 @@ EvalResult Evaluator::evalInternal(TNode n,
             {
               Trace("evaluator") << "Theory " << Theory::theoryOf(currNode[0])
                                  << " not supported" << std::endl;
-              return EvalResult();
+              results[currNode] = EvalResult();
+              evalAsNode[currNode] = currNodeVal;
               break;
             }
           }
@@ -629,7 +760,8 @@ EvalResult Evaluator::evalInternal(TNode n,
         {
           Trace("evaluator") << "Kind " << currNodeVal.getKind()
                              << " not supported" << std::endl;
-          return EvalResult();
+          results[currNode] = EvalResult();
+          evalAsNode[currNode] = currNodeVal;
         }
       }
     }
diff --git a/src/theory/evaluator.h b/src/theory/evaluator.h
index ce19a1f67..a65ffdf19 100644
--- a/src/theory/evaluator.h
+++ b/src/theory/evaluator.h
@@ -85,9 +85,9 @@ class Evaluator
  public:
   /**
    * Evaluates node `n` under the substitution described by the variable names
-   * `args` and the corresponding values `vals`. The function returns a null
-   * node if there is a subterm that is not constant under the substitution or
-   * if an operator is not supported by the evaluator.
+   * `args` and the corresponding values `vals`. This method uses evaluation
+   * for subterms that evaluate to constants supported in the EvalResult
+   * class and substitution with rewriting for the others.
    *
    * The nodes in the vals must be constant values, that is, they must return
    * true for isConst().
@@ -100,14 +100,24 @@ class Evaluator
   /**
    * Evaluates node `n` under the substitution described by the variable names
    * `args` and the corresponding values `vals`. The internal version returns
-   * an EvalResult which has slightly less overhead for recursive calls. The
-   * function returns an invalid result if there is a subterm that is not
-   * constant under the substitution or if an operator is not supported by the
-   * evaluator.
+   * an EvalResult which has slightly less overhead for recursive calls.
+   *
+   * The method returns an invalid EvalResult if the result of the substitution
+   * on n does not result in a constant value that is one of those supported in
+   * the EvalResult class. Notice that e.g. datatype constants are not supported
+   * in this class.
+   *
+   * This method additionally adds subterms of n that could not be evaluated
+   * (in the above sense) to the map evalAsNode. For each such subterm n', we
+   * store the node corresponding to the result of applying the substitution
+   * `args` to `vals` and rewriting. Notice that this map contains an entry
+   * for n in the case that it cannot be evaluated.
    */
-  EvalResult evalInternal(TNode n,
-                          const std::vector<Node>& args,
-                          const std::vector<Node>& vals);
+  EvalResult evalInternal(
+      TNode n,
+      const std::vector<Node>& args,
+      const std::vector<Node>& vals,
+      std::unordered_map<TNode, Node, NodeHashFunction>& evalAsNode);
 };
 
 }  // namespace theory