Rewrite witness terms at prerewrite (#5419)

Ensures (witness ((x Int)) (= x (+ 1 a)) is rewritten to (+ 1 a), instead of e.g. (witness ((x Int)) (= a (- x 1)).

This is required for supporting purification skolems for arithmetic terms in the new proof architecture.

2 files changed, 62 insertions, 44 deletions

diff --git a/src/theory/builtin/theory_builtin_rewriter.cpp b/src/theory/builtin/theory_builtin_rewriter.cpp
index 17270932d..99e2bfd74 100644
--- a/src/theory/builtin/theory_builtin_rewriter.cpp
+++ b/src/theory/builtin/theory_builtin_rewriter.cpp
@@ -92,43 +92,26 @@ RewriteResponse TheoryBuiltinRewriter::postRewrite(TNode node) {
     }
     return RewriteResponse(REWRITE_DONE, node);
   }
-  else if (node.getKind() == kind::WITNESS)
+  // otherwise, do the default call
+  return doRewrite(node);
+}
+
+RewriteResponse TheoryBuiltinRewriter::doRewrite(TNode node)
+{
+  switch (node.getKind())
   {
-    if (node[1].getKind() == kind::EQUAL)
-    {
-      for (unsigned i = 0; i < 2; i++)
-      {
-        // (witness ((x T)) (= x t)) ---> t
-        if (node[1][i] == node[0][0])
-        {
-          Trace("builtin-rewrite") << "Witness rewrite: " << node << " --> "
-                                   << node[1][1 - i] << std::endl;
-          // also must be a legal elimination: the other side of the equality
-          // cannot contain the variable, and it must be a subtype of the
-          // variable
-          if (!expr::hasSubterm(node[1][1 - i], node[0][0]) &&
-              node[1][i].getType().isSubtypeOf(node[0][0].getType()))
-          {
-            return RewriteResponse(REWRITE_DONE, node[1][1 - i]);
-          }
-        }
-      }
-    }
-    else if (node[1] == node[0][0])
+    case kind::WITNESS:
     {
-      // (witness ((x Bool)) x) ---> true
-      return RewriteResponse(REWRITE_DONE,
-                             NodeManager::currentNM()->mkConst(true));
+      // it is important to run this rewriting at prerewrite and postrewrite,
+      // since e.g. arithmetic rewrites equalities in ways that may make an
+      // equality not in solved form syntactically, e.g. (= x (+ 1 a)) rewrites
+      // to (= a (- x 1)), where x no longer is in solved form.
+      Node rnode = rewriteWitness(node);
+      return RewriteResponse(REWRITE_DONE, rnode);
     }
-    else if (node[1].getKind() == kind::NOT && node[1][0] == node[0][0])
-    {
-      // (witness ((x Bool)) (not x)) ---> false
-      return RewriteResponse(REWRITE_DONE,
-                             NodeManager::currentNM()->mkConst(false));
-    }
-    return RewriteResponse(REWRITE_DONE, node);
-  }else{ 
-    return doRewrite(node);
+    case kind::DISTINCT:
+      return RewriteResponse(REWRITE_DONE, blastDistinct(node));
+    default: return RewriteResponse(REWRITE_DONE, node);
   }
 }
 
@@ -468,6 +451,42 @@ Node TheoryBuiltinRewriter::getArrayRepresentationForLambdaRec(TNode n,
   }
 }
 
+Node TheoryBuiltinRewriter::rewriteWitness(TNode node)
+{
+  Assert(node.getKind() == kind::WITNESS);
+  if (node[1].getKind() == kind::EQUAL)
+  {
+    for (size_t i = 0; i < 2; i++)
+    {
+      // (witness ((x T)) (= x t)) ---> t
+      if (node[1][i] == node[0][0])
+      {
+        Trace("builtin-rewrite") << "Witness rewrite: " << node << " --> "
+                                 << node[1][1 - i] << std::endl;
+        // also must be a legal elimination: the other side of the equality
+        // cannot contain the variable, and it must be a subtype of the
+        // variable
+        if (!expr::hasSubterm(node[1][1 - i], node[0][0])
+            && node[1][i].getType().isSubtypeOf(node[0][0].getType()))
+        {
+          return node[1][1 - i];
+        }
+      }
+    }
+  }
+  else if (node[1] == node[0][0])
+  {
+    // (witness ((x Bool)) x) ---> true
+    return NodeManager::currentNM()->mkConst(true);
+  }
+  else if (node[1].getKind() == kind::NOT && node[1][0] == node[0][0])
+  {
+    // (witness ((x Bool)) (not x)) ---> false
+    return NodeManager::currentNM()->mkConst(false);
+  }
+  return node;
+}
+
 Node TheoryBuiltinRewriter::getArrayRepresentationForLambda(TNode n)
 {
   Assert(n.getKind() == kind::LAMBDA);
diff --git a/src/theory/builtin/theory_builtin_rewriter.h b/src/theory/builtin/theory_builtin_rewriter.h
index 3ccb6b4af..5fe4b159d 100644
--- a/src/theory/builtin/theory_builtin_rewriter.h
+++ b/src/theory/builtin/theory_builtin_rewriter.h
@@ -33,16 +33,6 @@ class TheoryBuiltinRewriter : public TheoryRewriter
 
  public:
 
-  static inline RewriteResponse doRewrite(TNode node)
-  {
-    switch (node.getKind())
-    {
-      case kind::DISTINCT:
-        return RewriteResponse(REWRITE_DONE, blastDistinct(node));
-      default: return RewriteResponse(REWRITE_DONE, node);
-    }
-  }
-
   RewriteResponse postRewrite(TNode node) override;
 
   RewriteResponse preRewrite(TNode node) override { return doRewrite(node); }
@@ -56,6 +46,15 @@ class TheoryBuiltinRewriter : public TheoryRewriter
   static Node getArrayRepresentationForLambdaRec(TNode n, TypeNode retType);
 
  public:
+  /**
+   * The default rewriter for rewrites that occur at both pre and post rewrite.
+   */
+  static RewriteResponse doRewrite(TNode node);
+  /**
+   * Main entry point for rewriting terms of the form (witness ((x T)) (P x)).
+   * Returns the rewritten form of node.
+   */
+  static Node rewriteWitness(TNode node);
   /** Get function type for array type
    *
    * This returns the function type of terms returned by the function