3 files changed, 30 insertions, 17 deletions
diff --git a/src/proof/proof_manager.cpp b/src/proof/proof_manager.cpp
index cad2baaa5..99e3010b4 100644
--- a/src/proof/proof_manager.cpp
+++ b/src/proof/proof_manager.cpp
@@ -565,6 +565,30 @@ void LFSCProof::toStream(std::ostream& out, const ProofLetMap& map) const
   Unreachable();
 }
 
+void collectAtoms(TNode node, std::set<Node>& seen, CnfProof* cnfProof)
+{
+  Debug("pf::pm::atoms") << "collectAtoms: Colleting atoms from " << node
+                         << "\n";
+  if (seen.find(node) != seen.end())
+  {
+    Debug("pf::pm::atoms") << "collectAtoms:\t already seen\n";
+    return;
+  }
+  // if I have a SAT literal for a node, save it, unless this node is a
+  // negation, in which case its underlying will be collected downstream
+  if (cnfProof->hasLiteral(node) && node.getKind() != kind::NOT)
+  {
+    Debug("pf::pm::atoms") << "collectAtoms: has SAT literal, save\n";
+    seen.insert(node);
+  }
+  for (unsigned i = 0; i < node.getNumChildren(); ++i)
+  {
+    Debug("pf::pm::atoms") << push;
+    collectAtoms(node[i], seen, cnfProof);
+    Debug("pf::pm::atoms") << pop;
+  }
+}
+
 void LFSCProof::toStream(std::ostream& out) const
 {
   TimerStat::CodeTimer proofProductionTimer(
@@ -683,10 +707,15 @@ void LFSCProof::toStream(std::ostream& out) const
     d_cnfProof->collectAtomsForClauses(used_lemmas, atoms);
 
     // collects the atoms in the assertions
+    Debug("pf::pm") << std::endl
+                    << "LFSCProof::toStream: Colleting atoms from assertions "
+                    << used_assertions << "\n"
+                    << push;
     for (TNode used_assertion : used_assertions)
     {
-      utils::collectAtoms(used_assertion, atoms);
+      collectAtoms(used_assertion, atoms, d_cnfProof);
     }
+    Debug("pf::pm") << pop;
 
     std::set<Node>::iterator atomIt;
     Debug("pf::pm") << std::endl
diff --git a/src/proof/proof_utils.cpp b/src/proof/proof_utils.cpp
index 9f73fed80..cad56db6a 100644
--- a/src/proof/proof_utils.cpp
+++ b/src/proof/proof_utils.cpp
@@ -21,19 +21,6 @@
 namespace CVC4 {
 namespace utils {
 
-void collectAtoms(TNode node, std::set<Node>& seen) {
-  if (seen.find(node) != seen.end())
-    return;
-  if (theory::Theory::theoryOf(node) != theory::THEORY_BOOL || node.isVar()) {
-      seen.insert(node);
-      return;
-  }
-
-  for (unsigned i = 0; i < node.getNumChildren(); ++i) {
-    collectAtoms(node[i], seen);
-  }
-}
-
 std::string toLFSCKind(Kind kind) {
   switch(kind) {
     // core kinds
diff --git a/src/proof/proof_utils.h b/src/proof/proof_utils.h
index 147835ecd..e54edd8b7 100644
--- a/src/proof/proof_utils.h
+++ b/src/proof/proof_utils.h
@@ -225,8 +225,5 @@ inline const bool getBit(Expr expr, unsigned i) {
   return (bit == 1u);
 }
 
-void collectAtoms(TNode node, std::set<Node>& seen);
-
-
 }
 }