1 files changed, 36 insertions, 18 deletions

diff --git a/src/theory/evaluator.cpp b/src/theory/evaluator.cpp
index cdf5c4e5a..852dede78 100644
--- a/src/theory/evaluator.cpp
+++ b/src/theory/evaluator.cpp
@@ -117,16 +117,18 @@ Node EvalResult::toNode() const
 
 Node Evaluator::eval(TNode n,
                      const std::vector<Node>& args,
-                     const std::vector<Node>& vals) const
+                     const std::vector<Node>& vals,
+                     bool useRewriter) const
 {
   std::unordered_map<Node, Node, NodeHashFunction> visited;
-  return eval(n, args, vals, visited);
+  return eval(n, args, vals, visited, useRewriter);
 }
 Node Evaluator::eval(
     TNode n,
     const std::vector<Node>& args,
     const std::vector<Node>& vals,
-    const std::unordered_map<Node, Node, NodeHashFunction>& visited) const
+    const std::unordered_map<Node, Node, NodeHashFunction>& visited,
+    bool useRewriter) const
 {
   Trace("evaluator") << "Evaluating " << n << " under substitution " << args
                      << " " << vals << " with visited size = " << visited.size()
@@ -137,28 +139,38 @@ Node Evaluator::eval(
   for (const std::pair<const Node, Node>& p : visited)
   {
     Trace("evaluator") << "Add " << p.first << " == " << p.second << std::endl;
-    results[p.first] = evalInternal(p.second, args, vals, evalAsNode, results);
+    results[p.first] = evalInternal(p.second, args, vals, evalAsNode, results, useRewriter);
     if (results[p.first].d_tag == EvalResult::INVALID)
     {
       // could not evaluate, use the evalAsNode map
-      evalAsNode[p.first] = evalAsNode[p.second];
+      std::unordered_map<TNode, Node, NodeHashFunction>::iterator itn =
+          evalAsNode.find(p.second);
+      Assert(itn != evalAsNode.end());
+      Node val = itn->second;
+      if (useRewriter)
+      {
+        val = Rewriter::rewrite(val);
+      }
+      evalAsNode[p.first] = val;
     }
   }
   Trace("evaluator") << "Run eval internal..." << std::endl;
-  Node ret = evalInternal(n, args, vals, evalAsNode, results).toNode();
+  Node ret = evalInternal(n, args, vals, evalAsNode, results, useRewriter).toNode();
   // if we failed to evaluate
-  if (ret.isNull())
+  if (ret.isNull() && useRewriter)
   {
     // should be stored in the evaluation-as-node map
     std::unordered_map<TNode, Node, NodeHashFunction>::iterator itn =
         evalAsNode.find(n);
     Assert(itn != evalAsNode.end());
-    ret = itn->second;
+    ret = Rewriter::rewrite(itn->second);
   }
-  // should be the same as substitution + rewriting
-  Assert(ret
-         == Rewriter::rewrite(
-             n.substitute(args.begin(), args.end(), vals.begin(), vals.end())));
+  // should be the same as substitution + rewriting, or possibly null if
+  // useRewriter is false
+  Assert((ret.isNull() && !useRewriter)
+         || ret
+                == Rewriter::rewrite(n.substitute(
+                       args.begin(), args.end(), vals.begin(), vals.end())));
   return ret;
 }
 
@@ -167,7 +179,8 @@ EvalResult Evaluator::evalInternal(
     const std::vector<Node>& args,
     const std::vector<Node>& vals,
     std::unordered_map<TNode, Node, NodeHashFunction>& evalAsNode,
-    std::unordered_map<TNode, EvalResult, TNodeHashFunction>& results) const
+    std::unordered_map<TNode, EvalResult, TNodeHashFunction>& results,
+    bool useRewriter) const
 {
   std::vector<TNode> queue;
   queue.emplace_back(n);
@@ -243,6 +256,12 @@ EvalResult Evaluator::evalInternal(
         // successfully evaluated, and the children that did not.
         Trace("evaluator") << "Evaluator: collect arguments" << std::endl;
         currNodeVal = reconstruct(currNodeVal, results, evalAsNode);
+        if (useRewriter)
+        {
+          // Rewrite the result now, if we use the rewriter. We will see below
+          // if we are able to turn it into a valid EvalResult.
+          currNodeVal = Rewriter::rewrite(currNodeVal);
+        }
         needsReconstruct = false;
         Trace("evaluator") << "Evaluator: now after substitution + rewriting: "
                            << currNodeVal << std::endl;
@@ -274,7 +293,6 @@ EvalResult Evaluator::evalInternal(
         currNodeVal = vals[pos];
         // Don't need to reconstruct since range of substitution should already
         // be normalized.
-        Assert(vals[pos] == Rewriter::rewrite(vals[pos]));
         needsReconstruct = false;
       }
       else if (currNode.getKind() == kind::APPLY_UF)
@@ -318,7 +336,7 @@ EvalResult Evaluator::evalInternal(
         std::unordered_map<TNode, Node, NodeHashFunction> evalAsNodeC;
         std::unordered_map<TNode, EvalResult, TNodeHashFunction> resultsC;
         results[currNode] =
-            evalInternal(op[1], lambdaArgs, lambdaVals, evalAsNodeC, resultsC);
+            evalInternal(op[1], lambdaArgs, lambdaVals, evalAsNodeC, resultsC, useRewriter);
         Trace("evaluator") << "Evaluated via arguments to "
                            << results[currNode].d_tag << std::endl;
         if (results[currNode].d_tag == EvalResult::INVALID)
@@ -894,9 +912,9 @@ Node Evaluator::reconstruct(
   // of the children.
   Node nn = nm->mkNode(n.getKind(), echildren);
   Trace("evaluator") << "Evaluator: reconstructed " << nn << std::endl;
-  // Use rewriting. Notice we do not need to substitute here since
-  // all substitutions should already have been applied recursively.
-  return Rewriter::rewrite(nn);
+  // Return node, without rewriting. Notice we do not need to substitute here
+  // since all substitutions should already have been applied recursively.
+  return nn;
 }
 
 }  // namespace theory