1 files changed, 53 insertions, 1 deletions

diff --git a/src/theory/theory_inference_manager.cpp b/src/theory/theory_inference_manager.cpp
index 4ed01b618..fda6d0881 100644
--- a/src/theory/theory_inference_manager.cpp
+++ b/src/theory/theory_inference_manager.cpp
@@ -15,7 +15,13 @@
 
 #include "theory/theory_inference_manager.h"
 
+#include "options/proof_options.h"
+#include "proof/annotation_proof_generator.h"
+#include "proof/eager_proof_generator.h"
 #include "smt/smt_statistics_registry.h"
+#include "smt/solver_engine_scope.h"
+#include "theory/builtin/proof_checker.h"
+#include "theory/inference_id_proof_annotator.h"
 #include "theory/output_channel.h"
 #include "theory/rewriter.h"
 #include "theory/theory.h"
@@ -56,6 +62,20 @@ TheoryInferenceManager::TheoryInferenceManager(Env& env,
   // don't add true lemma
   Node truen = NodeManager::currentNM()->mkConst(true);
   d_lemmasSent.insert(truen);
+
+  if (isProofEnabled())
+  {
+    context::UserContext* u = userContext();
+    ProofNodeManager* pnm = env.getProofNodeManager();
+    d_defaultPg.reset(
+        new EagerProofGenerator(pnm, u, statsName + "EagerProofGenerator"));
+    if (options().proof.proofAnnotate)
+    {
+      d_iipa.reset(new InferenceIdProofAnnotator(pnm, u));
+      d_apg.reset(new AnnotationProofGenerator(
+          pnm, u, statsName + "AnnotationProofGenerator"));
+    }
+  }
 }
 
 TheoryInferenceManager::~TheoryInferenceManager()
@@ -127,6 +147,11 @@ void TheoryInferenceManager::trustedConflict(TrustNode tconf, InferenceId id)
   resourceManager()->spendResource(id);
   Trace("im") << "(conflict " << id << " " << tconf.getProven() << ")"
               << std::endl;
+  // annotate if the annotation proof generator is active
+  if (d_apg != nullptr)
+  {
+    tconf = annotateId(tconf, id);
+  }
   d_out.trustedConflict(tconf);
   ++d_numConflicts;
 }
@@ -261,7 +286,16 @@ bool TheoryInferenceManager::trustedLemma(const TrustNode& tlem,
   // shouldn't send trivially true or false lemmas
   Assert(!rewrite(tlem.getProven()).isConst());
   d_numCurrentLemmas++;
-  d_out.trustedLemma(tlem, p);
+  // annotate if the annotation proof generator is active
+  if (d_apg != nullptr)
+  {
+    TrustNode tlema = annotateId(tlem, id);
+    d_out.trustedLemma(tlema, p);
+  }
+  else
+  {
+    d_out.trustedLemma(tlem, p);
+  }
   return true;
 }
 
@@ -544,6 +578,24 @@ bool TheoryInferenceManager::cacheLemma(TNode lem, LemmaProperty p)
   return true;
 }
 
+TrustNode TheoryInferenceManager::annotateId(const TrustNode& trn,
+                                             InferenceId id)
+{
+  Assert(d_iipa != nullptr && d_apg != nullptr);
+  Node lemma = trn.getProven();
+  TrustNode trna = trn;
+  // ensure we have a proof generator, make trusted theory lemma if not
+  if (trn.getGenerator() == nullptr)
+  {
+    Node tidn =
+        builtin::BuiltinProofRuleChecker::mkTheoryIdNode(d_theory.getId());
+    trna = d_defaultPg->mkTrustNode(
+        lemma, PfRule::THEORY_LEMMA, {}, {lemma, tidn});
+  }
+  d_iipa->setAnnotation(lemma, id);
+  return d_apg->transform(trna, d_iipa.get());
+}
+
 DecisionManager* TheoryInferenceManager::getDecisionManager()
 {
   return d_decManager;