2 files changed, 72 insertions, 14 deletions

diff --git a/src/proof/dot/dot_printer.cpp b/src/proof/dot/dot_printer.cpp
index ca85aadd3..08be1728b 100644
--- a/src/proof/dot/dot_printer.cpp
+++ b/src/proof/dot/dot_printer.cpp
@@ -46,23 +46,66 @@ std::string DotPrinter::sanitizeString(const std::string& s)
   return newS;
 }
 
+void DotPrinter::countSubproofs(
+    const ProofNode* pn, std::map<const ProofNode*, size_t>& subpfCounter)
+{
+  std::vector<const ProofNode*> visit;
+  std::unordered_map<const ProofNode*, bool> visited;
+  std::unordered_map<const ProofNode*, bool>::iterator it;
+  const ProofNode* cur;
+  visit.push_back(pn);
+  do
+  {
+    cur = visit.back();
+    visit.pop_back();
+    it = visited.find(cur);
+    if (it == visited.end())
+    {
+      visited[cur] = false;
+      visit.push_back(cur);
+      const std::vector<std::shared_ptr<ProofNode>>& children =
+          cur->getChildren();
+      for (const std::shared_ptr<ProofNode>& c : children)
+      {
+        visit.push_back(c.get());
+      }
+    }
+    else if (!it->second)
+    {
+      visited[cur] = true;
+      size_t counter = 1;
+      const std::vector<std::shared_ptr<ProofNode>>& children =
+          cur->getChildren();
+      for (const std::shared_ptr<ProofNode>& c : children)
+      {
+        counter += subpfCounter[c.get()];
+      }
+      subpfCounter[cur] = counter;
+    }
+  } while (!visit.empty());
+}
+
 void DotPrinter::print(std::ostream& out, const ProofNode* pn)
 {
   uint64_t ruleID = 0;
+  std::map<const ProofNode*, size_t> subpfCounter;
+  countSubproofs(pn, subpfCounter);
 
   // The dot attribute rankdir="BT" sets the direction of the graph layout,
   // placing the root node at the top. The "node [shape..." attribute sets the
   // shape of all nodes to record.
   out << "digraph proof {\n\trankdir=\"BT\";\n\tnode [shape=record];\n";
-  DotPrinter::printInternal(out, pn, ruleID, 0, false);
+  DotPrinter::printInternal(out, pn, ruleID, 0, false, subpfCounter);
   out << "}\n";
 }
 
-void DotPrinter::printInternal(std::ostream& out,
-                               const ProofNode* pn,
-                               uint64_t& ruleID,
-                               uint64_t scopeCounter,
-                               bool inPropositionalView)
+void DotPrinter::printInternal(
+    std::ostream& out,
+    const ProofNode* pn,
+    uint64_t& ruleID,
+    uint64_t scopeCounter,
+    bool inPropositionalView,
+    const std::map<const ProofNode*, size_t>& subpfCounter)
 {
   uint64_t currentRuleID = ruleID;
   const std::vector<std::shared_ptr<ProofNode>>& children = pn->getChildren();
@@ -120,13 +163,17 @@ void DotPrinter::printInternal(std::ostream& out,
   }
   classes << " \"";
   out << classes.str() << colors.str();
+  // add number of subchildren
+  std::map<const ProofNode*, size_t>::const_iterator it = subpfCounter.find(pn);
+  out << ", comment = \"\{\"subProofQty\":" << it->second << "}\"";
   out << " ];\n";
 
   for (const std::shared_ptr<ProofNode>& c : children)
   {
     ++ruleID;
     out << "\t" << ruleID << " -> " << currentRuleID << ";\n";
-    printInternal(out, c.get(), ruleID, scopeCounter, inPropositionalView);
+    printInternal(
+        out, c.get(), ruleID, scopeCounter, inPropositionalView, subpfCounter);
   }
 }
 
@@ -184,4 +231,4 @@ void DotPrinter::ruleArguments(std::ostringstream& currentArguments,
 }
 
 }  // namespace proof
-}  // namespace cvc5
\ No newline at end of file
+}  // namespace cvc5
diff --git a/src/proof/dot/dot_printer.h b/src/proof/dot/dot_printer.h
index 0027d145a..a877efe9b 100644
--- a/src/proof/dot/dot_printer.h
+++ b/src/proof/dot/dot_printer.h
@@ -53,11 +53,13 @@ class DotPrinter
    * @param inPropositionalView flag used to mark the proof node being traversed
    * was generated by the SAT solver and thus belong to the propositional view
    */
-  static void printInternal(std::ostream& out,
-                            const ProofNode* pn,
-                            uint64_t& ruleID,
-                            uint64_t scopeCounter,
-                            bool inPropositionalView);
+  static void printInternal(
+      std::ostream& out,
+      const ProofNode* pn,
+      uint64_t& ruleID,
+      uint64_t scopeCounter,
+      bool inPropositionalView,
+      const std::map<const ProofNode*, size_t>& subpfCounter);
 
   /**
    * Return the arguments of a ProofNode
@@ -73,9 +75,18 @@ class DotPrinter
    * @return The string with the special characters escaped.
    */
   static std::string sanitizeString(const std::string& s);
+
+  /** Traverse proof node and map each proof node to the number of subproofs it
+   * contains, including itself
+   *
+   * @param pn the proof node to be traversed
+   * @param subpfCounter the map to be populated
+   */
+  static void countSubproofs(const ProofNode* pn,
+                             std::map<const ProofNode*, size_t>& subpfCounter);
 };
 
 }  // namespace proof
 }  // namespace cvc5
 
-#endif
\ No newline at end of file
+#endif