1 files changed, 10 insertions, 12 deletions

diff --git a/src/theory/rewriter.cpp b/src/theory/rewriter.cpp
index 6cb3bf3ff..361d90dcd 100644
--- a/src/theory/rewriter.cpp
+++ b/src/theory/rewriter.cpp
@@ -32,13 +32,6 @@ using namespace std;
 namespace cvc5 {
 namespace theory {
 
-/** Attribute true for nodes that have been rewritten with proofs enabled */
-struct RewriteWithProofsAttributeId
-{
-};
-typedef expr::Attribute<RewriteWithProofsAttributeId, bool>
-    RewriteWithProofsAttribute;
-
 // Note that this function is a simplified version of Theory::theoryOf for
 // (type-based) theoryOfMode. We expand and simplify it here for the sake of
 // efficiency.
@@ -173,7 +166,6 @@ Node Rewriter::rewriteTo(theory::TheoryId theoryId,
                          Node node,
                          TConvProofGenerator* tcpg)
 {
-  RewriteWithProofsAttribute rpfa;
 #ifdef CVC5_ASSERTIONS
   bool isEquality = node.getKind() == kind::EQUAL && (!node[0].getType().isBoolean());
 
@@ -187,7 +179,7 @@ Node Rewriter::rewriteTo(theory::TheoryId theoryId,
 
   // Check if it's been cached already
   Node cached = getPostRewriteCache(theoryId, node);
-  if (!cached.isNull() && (tcpg == nullptr || node.getAttribute(rpfa)))
+  if (!cached.isNull() && (tcpg == nullptr || hasRewrittenWithProofs(node)))
   {
     return cached;
   }
@@ -217,7 +209,8 @@ Node Rewriter::rewriteTo(theory::TheoryId theoryId,
       cached = getPreRewriteCache(rewriteStackTop.getTheoryId(),
                                   rewriteStackTop.d_node);
       if (cached.isNull()
-          || (tcpg != nullptr && !rewriteStackTop.d_node.getAttribute(rpfa)))
+          || (tcpg != nullptr
+              && !hasRewrittenWithProofs(rewriteStackTop.d_node)))
       {
         // Rewrite until fix-point is reached
         for(;;) {
@@ -256,7 +249,7 @@ Node Rewriter::rewriteTo(theory::TheoryId theoryId,
                                  rewriteStackTop.d_node);
     // If not, go through the children
     if (cached.isNull()
-        || (tcpg != nullptr && !rewriteStackTop.d_node.getAttribute(rpfa)))
+        || (tcpg != nullptr && !hasRewrittenWithProofs(rewriteStackTop.d_node)))
     {
       // The child we need to rewrite
       unsigned child = rewriteStackTop.d_nextChild++;
@@ -343,7 +336,7 @@ Node Rewriter::rewriteTo(theory::TheoryId theoryId,
       if (tcpg != nullptr)
       {
         // if proofs are enabled, mark that we've rewritten with proofs
-        rewriteStackTop.d_original.setAttribute(rpfa, true);
+        d_tpgNodes.insert(rewriteStackTop.d_original);
         if (!cached.isNull())
         {
           // We may have gotten a different node, due to non-determinism in
@@ -474,5 +467,10 @@ void Rewriter::clearCaches()
   clearCachesInternal();
 }
 
+bool Rewriter::hasRewrittenWithProofs(TNode n) const
+{
+  return d_tpgNodes.find(n) != d_tpgNodes.end();
+}
+
 }  // namespace theory
 }  // namespace cvc5