4 files changed, 34 insertions, 1 deletions
diff --git a/src/expr/term_context.cpp b/src/expr/term_context.cpp
index d339a3d28..7f8aa9eac 100644
--- a/src/expr/term_context.cpp
+++ b/src/expr/term_context.cpp
@@ -15,6 +15,8 @@
 
 #include "expr/term_context.h"
 
+#include "theory/theory.h"
+
 namespace cvc5 {
 
 uint32_t TermContext::computeValueOp(TNode t, uint32_t tval) const
@@ -133,4 +135,13 @@ void PolarityTermContext::getFlags(uint32_t val, bool& hasPol, bool& pol)
   pol = val == 2;
 }
 
+uint32_t TheoryLeafTermContext::initialValue() const { return 0; }
+
+uint32_t TheoryLeafTermContext::computeValue(TNode t,
+                                             uint32_t tval,
+                                             size_t index) const
+{
+  return theory::Theory::isLeafOf(t, d_theoryId) ? 1 : tval;
+}
+
 }  // namespace cvc5
diff --git a/src/expr/term_context.h b/src/expr/term_context.h
index 37d2ef6bd..db08372ef 100644
--- a/src/expr/term_context.h
+++ b/src/expr/term_context.h
@@ -19,6 +19,7 @@
 #define CVC5__EXPR__TERM_CONTEXT_H
 
 #include "expr/node.h"
+#include "theory/theory_id.h"
 
 namespace cvc5 {
 
@@ -164,6 +165,23 @@ class PolarityTermContext : public TermContext
   static void getFlags(uint32_t val, bool& hasPol, bool& pol);
 };
 
+/**
+ * Similar to InQuantTermContext, but computes whether we are below a theory
+ * leaf of given theory id.
+ */
+class TheoryLeafTermContext : public TermContext
+{
+ public:
+  TheoryLeafTermContext(theory::TheoryId id) : d_theoryId(id) {}
+  /** The initial value: not beneath a theory leaf. */
+  uint32_t initialValue() const override;
+  /** Compute the value of the index^th child of t whose hash is tval */
+  uint32_t computeValue(TNode t, uint32_t tval, size_t index) const override;
+
+ private:
+  theory::TheoryId d_theoryId;
+};
+
 }  // namespace cvc5
 
 #endif /* CVC5__EXPR__TERM_CONVERSION_PROOF_GENERATOR_H */
diff --git a/src/theory/bv/bitblast/proof_bitblaster.cpp b/src/theory/bv/bitblast/proof_bitblaster.cpp
index ea7ba1a13..a1b856d88 100644
--- a/src/theory/bv/bitblast/proof_bitblaster.cpp
+++ b/src/theory/bv/bitblast/proof_bitblaster.cpp
@@ -70,6 +70,7 @@ std::unordered_map<Kind, PfRule, kind::KindHashFunction>
 BBProof::BBProof(TheoryState* state, ProofNodeManager* pnm, bool fineGrained)
     : d_bb(new BBSimple(state)),
       d_pnm(pnm),
+      d_tcontext(new TheoryLeafTermContext(theory::THEORY_BV)),
       d_tcpg(pnm ? new TConvProofGenerator(
                  pnm,
                  nullptr,
@@ -80,7 +81,7 @@ BBProof::BBProof(TheoryState* state, ProofNodeManager* pnm, bool fineGrained)
                  /* STATIC to get the same ProofNode for a shared subterm. */
                  TConvCachePolicy::STATIC,
                  "BBProof::TConvProofGenerator",
-                 nullptr,
+                 d_tcontext.get(),
                  false)
                  : nullptr),
       d_recordFineGrainedProofs(fineGrained)
diff --git a/src/theory/bv/bitblast/proof_bitblaster.h b/src/theory/bv/bitblast/proof_bitblaster.h
index 86cbeae81..428581fe0 100644
--- a/src/theory/bv/bitblast/proof_bitblaster.h
+++ b/src/theory/bv/bitblast/proof_bitblaster.h
@@ -17,6 +17,7 @@
 #ifndef CVC5__THEORY__BV__BITBLAST__PROOF_BITBLASTER_H
 #define CVC5__THEORY__BV__BITBLAST__PROOF_BITBLASTER_H
 
+#include "expr/term_context.h"
 #include "theory/bv/bitblast/simple_bitblaster.h"
 
 namespace cvc5 {
@@ -59,6 +60,8 @@ class BBProof
   std::unique_ptr<BBSimple> d_bb;
   /** The associated proof node manager. */
   ProofNodeManager* d_pnm;
+  /** Term context for d_tcpg to not rewrite below BV leafs. */
+  std::unique_ptr<TermContext> d_tcontext;
   /** The associated term conversion proof generator. */
   std::unique_ptr<TConvProofGenerator> d_tcpg;
   /** Map bit-vector nodes to bit-blasted nodes. */