4 files changed, 27 insertions, 15 deletions

diff --git a/src/smt/preprocessor.cpp b/src/smt/preprocessor.cpp
index 02323561d..912c0ea28 100644
--- a/src/smt/preprocessor.cpp
+++ b/src/smt/preprocessor.cpp
@@ -30,12 +30,14 @@ namespace smt {
 Preprocessor::Preprocessor(SmtEngine& smt,
                            context::UserContext* u,
                            AbstractValues& abs)
-    : d_smt(smt),
+    : d_context(u),
+      d_smt(smt),
       d_absValues(abs),
       d_propagator(true, true),
       d_assertionsProcessed(u, false),
       d_processor(smt, *smt.getResourceManager()),
-      d_rtf(u)
+      d_rtf(u),
+      d_pnm(nullptr)
 {
 }
 
@@ -51,7 +53,7 @@ Preprocessor::~Preprocessor()
 void Preprocessor::finishInit()
 {
   d_ppContext.reset(new preprocessing::PreprocessingPassContext(
-      &d_smt, &d_rtf, &d_propagator));
+      &d_smt, &d_rtf, &d_propagator, d_pnm));
 
   // initialize the preprocessing passes
   d_processor.finishInit(d_ppContext.get());
@@ -62,7 +64,8 @@ bool Preprocessor::process(Assertions& as)
   preprocessing::AssertionPipeline& ap = as.getAssertionPipeline();
 
   // should not be called if empty
-  Assert(ap.size() != 0) << "Can only preprocess a non-empty list of assertions";
+  Assert(ap.size() != 0)
+      << "Can only preprocess a non-empty list of assertions";
 
   if (d_assertionsProcessed && options::incrementalSolving())
   {
@@ -140,7 +143,8 @@ Node Preprocessor::simplify(const Node& node, bool removeItes)
   }
   std::unordered_map<Node, Node, NodeHashFunction> cache;
   Node n = d_processor.expandDefinitions(nas, cache);
-  Node ns = applySubstitutions(n);
+  TrustNode ts = d_ppContext->getTopLevelSubstitutions().apply(n);
+  Node ns = ts.isNull() ? n : ts.getNode();
   if (removeItes)
   {
     // also remove ites if asked
@@ -149,9 +153,11 @@ Node Preprocessor::simplify(const Node& node, bool removeItes)
   return ns;
 }
 
-Node Preprocessor::applySubstitutions(TNode node)
+void Preprocessor::setProofGenerator(PreprocessProofGenerator* pppg)
 {
-  return Rewriter::rewrite(d_ppContext->getTopLevelSubstitutions().apply(node));
+  Assert(pppg != nullptr);
+  d_pnm = pppg->getManager();
+  d_rtf.setProofNodeManager(d_pnm);
 }
 
 }  // namespace smt
diff --git a/src/smt/preprocessor.h b/src/smt/preprocessor.h
index 81757de37..b07e7ec97 100644
--- a/src/smt/preprocessor.h
+++ b/src/smt/preprocessor.h
@@ -100,12 +100,14 @@ class Preprocessor
    */
   RemoveTermFormulas& getTermFormulaRemover();
 
- private:
   /**
-   * Apply substitutions that have been inferred by preprocessing, return the
-   * substituted form of node.
+   * Set proof node manager. Enables proofs in this preprocessor.
    */
-  Node applySubstitutions(TNode node);
+  void setProofGenerator(PreprocessProofGenerator* pppg);
+
+ private:
+  /** A copy of the current context */
+  context::Context* d_context;
   /** Reference to the parent SmtEngine */
   SmtEngine& d_smt;
   /** Reference to the abstract values utility */
@@ -130,6 +132,8 @@ class Preprocessor
    * in term contexts.
    */
   RemoveTermFormulas d_rtf;
+  /** Proof node manager */
+  ProofNodeManager* d_pnm;
 };
 
 }  // namespace smt
diff --git a/src/smt/process_assertions.cpp b/src/smt/process_assertions.cpp
index 719165048..f8af72c3a 100644
--- a/src/smt/process_assertions.cpp
+++ b/src/smt/process_assertions.cpp
@@ -105,9 +105,6 @@ bool ProcessAssertions::apply(Assertions& as)
     return true;
   }
 
-  SubstitutionMap& top_level_substs =
-      d_preprocessingPassContext->getTopLevelSubstitutions();
-
   if (options::bvGaussElim())
   {
     d_passes["bv-gauss"]->apply(&assertions);
@@ -330,6 +327,8 @@ bool ProcessAssertions::apply(Assertions& as)
       // First, find all skolems that appear in the substitution map - their
       // associated iteExpr will need to be moved to the main assertion set
       set<TNode> skolemSet;
+      SubstitutionMap& top_level_substs =
+          d_preprocessingPassContext->getTopLevelSubstitutions().get();
       SubstitutionMap::iterator pos = top_level_substs.begin();
       for (; pos != top_level_substs.end(); ++pos)
       {
diff --git a/src/smt/smt_engine.cpp b/src/smt/smt_engine.cpp
index 2a771ce76..2f9918486 100644
--- a/src/smt/smt_engine.cpp
+++ b/src/smt/smt_engine.cpp
@@ -259,14 +259,17 @@ void SmtEngine::finishInit()
   if (options::proofNew())
   {
     d_pfManager.reset(new PfManager(getUserContext(), this));
+    PreprocessProofGenerator* pppg = d_pfManager->getPreprocessProofGenerator();
     // use this proof node manager
     pnm = d_pfManager->getProofNodeManager();
     // enable proof support in the rewriter
     d_rewriter->setProofNodeManager(pnm);
     // enable it in the assertions pipeline
-    d_asserts->setProofGenerator(d_pfManager->getPreprocessProofGenerator());
+    d_asserts->setProofGenerator(pppg);
     // enable it in the SmtSolver
     d_smtSolver->setProofNodeManager(pnm);
+    // enabled proofs in the preprocessor
+    d_pp->setProofGenerator(pppg);
   }
 
   Trace("smt-debug") << "SmtEngine::finishInit" << std::endl;