[proof-new] Updates to proof node updater (#5156)

2 files changed, 9 insertions, 3 deletions

diff --git a/src/expr/proof_node_updater.cpp b/src/expr/proof_node_updater.cpp
index b97d27c05..295d839e1 100644
--- a/src/expr/proof_node_updater.cpp
+++ b/src/expr/proof_node_updater.cpp
@@ -162,7 +162,7 @@ bool ProofNodeUpdater::runUpdate(std::shared_ptr<ProofNode> cur,
                                  bool& continueUpdate)
 {
   // should it be updated?
-  if (!d_cb.shouldUpdate(cur.get()))
+  if (!d_cb.shouldUpdate(cur, continueUpdate))
   {
     return false;
   }
diff --git a/src/expr/proof_node_updater.h b/src/expr/proof_node_updater.h
index 63e764d20..069b625df 100644
--- a/src/expr/proof_node_updater.h
+++ b/src/expr/proof_node_updater.h
@@ -35,8 +35,14 @@ class ProofNodeUpdaterCallback
  public:
   ProofNodeUpdaterCallback();
   virtual ~ProofNodeUpdaterCallback();
-  /** Should proof pn be updated? */
-  virtual bool shouldUpdate(ProofNode* pn) = 0;
+  /** Should proof pn be updated?
+   *
+   * @param pn the proof node that maybe should be updated
+   * @param continueUpdate whether we should continue recursively updating pn
+   * @return whether we should run the update method on pn
+   */
+  virtual bool shouldUpdate(std::shared_ptr<ProofNode> pn,
+                            bool& continueUpdate) = 0;
   /**
    * Update the proof rule application, store steps in cdp. Return true if
    * the proof changed. It can be assumed that cdp contains proofs of each