(proof-new) Proof production for term formula removal (#4687)

This adds proof support in the term formula removal pass.

It also refactors this class heavily so that its interface is more intuitive and does not depend on AssertionPipeline.

9 files changed, 366 insertions, 123 deletions

diff --git a/src/expr/proof_rule.cpp b/src/expr/proof_rule.cpp
index c51bfb3c7..7844dfccb 100644
--- a/src/expr/proof_rule.cpp
+++ b/src/expr/proof_rule.cpp
@@ -33,6 +33,7 @@ const char* toString(PfRule id)
     case PfRule::MACRO_SR_PRED_TRANSFORM: return "MACRO_SR_PRED_TRANSFORM";
     case PfRule::THEORY_REWRITE: return "THEORY_REWRITE";
     case PfRule::PREPROCESS: return "PREPROCESS";
+    case PfRule::REMOVE_TERM_FORMULA_AXIOM: return "REMOVE_TERM_FORMULA_AXIOM";
     //================================================= Boolean rules
     case PfRule::SPLIT: return "SPLIT";
     case PfRule::EQ_RESOLVE: return "EQ_RESOLVE";
diff --git a/src/expr/proof_rule.h b/src/expr/proof_rule.h
index c76e907c7..62bc77cfb 100644
--- a/src/expr/proof_rule.h
+++ b/src/expr/proof_rule.h
@@ -183,6 +183,13 @@ enum class PfRule : uint32_t
   // based on some preprocessing pass, or otherwise F was added as a new
   // assertion by some preprocessing pass.
   PREPROCESS,
+  // ======== Remove Term Formulas Axiom
+  // Children: none
+  // Arguments: (t)
+  // ---------------------------------------------------------------
+  // Conclusion: RemoveTermFormulas::getAxiomFor(t).
+  REMOVE_TERM_FORMULA_AXIOM,
+
   //================================================= Boolean rules
   // ======== Split
   // Children: none
diff --git a/src/preprocessing/passes/ite_removal.cpp b/src/preprocessing/passes/ite_removal.cpp
index f021fff28..2cc244a34 100644
--- a/src/preprocessing/passes/ite_removal.cpp
+++ b/src/preprocessing/passes/ite_removal.cpp
@@ -34,9 +34,23 @@ PreprocessingPassResult IteRemoval::applyInternal(AssertionPipeline* assertions)
 {
   d_preprocContext->spendResource(ResourceManager::Resource::PreprocessStep);
 
+  IteSkolemMap& imap = assertions->getIteSkolemMap();
   // Remove all of the ITE occurrences and normalize
-  d_preprocContext->getIteRemover()->run(
-      assertions->ref(), assertions->getIteSkolemMap(), true);
+  for (unsigned i = 0, size = assertions->size(); i < size; ++i)
+  {
+    std::vector<theory::TrustNode> newAsserts;
+    std::vector<Node> newSkolems;
+    TrustNode trn = d_preprocContext->getIteRemover()->run(
+        (*assertions)[i], newAsserts, newSkolems, true);
+    // process
+    assertions->replace(i, trn.getNode());
+    Assert(newSkolems.size() == newAsserts.size());
+    for (unsigned j = 0, nnasserts = newAsserts.size(); j < nnasserts; j++)
+    {
+      imap[newSkolems[j]] = assertions->size();
+      assertions->ref().push_back(newAsserts[j].getNode());
+    }
+  }
   for (unsigned i = 0, size = assertions->size(); i < size; ++i)
   {
     assertions->replace(i, Rewriter::rewrite((*assertions)[i]));
diff --git a/src/smt/term_formula_removal.cpp b/src/smt/term_formula_removal.cpp
index 89091d309..3f44e592e 100644
--- a/src/smt/term_formula_removal.cpp
+++ b/src/smt/term_formula_removal.cpp
@@ -18,6 +18,7 @@
 #include <vector>
 
 #include "expr/node_algorithm.h"
+#include "expr/skolem_manager.h"
 #include "options/proof_options.h"
 #include "proof/proof_manager.h"
 
@@ -26,44 +27,51 @@ using namespace std;
 namespace CVC4 {
 
 RemoveTermFormulas::RemoveTermFormulas(context::UserContext* u)
-    : d_tfCache(u), d_skolem_cache(u)
+    : d_tfCache(u),
+      d_skolem_cache(u),
+      d_pnm(nullptr),
+      d_tpg(nullptr),
+      d_lp(nullptr)
 {
 }
 
 RemoveTermFormulas::~RemoveTermFormulas() {}
 
-void RemoveTermFormulas::run(std::vector<Node>& output, IteSkolemMap& iteSkolemMap, bool reportDeps)
+theory::TrustNode RemoveTermFormulas::run(
+    Node assertion,
+    std::vector<theory::TrustNode>& newAsserts,
+    std::vector<Node>& newSkolems,
+    bool reportDeps)
 {
-  size_t n = output.size();
-  for (unsigned i = 0, i_end = output.size(); i < i_end; ++ i) {
-    // Do this in two steps to avoid Node problems(?)
-    // Appears related to bug 512, splitting this into two lines
-    // fixes the bug on clang on Mac OS
-    Node itesRemoved = run(output[i], output, iteSkolemMap, false, false);
-    // In some calling contexts, not necessary to report dependence information.
-    if (reportDeps &&
-        (options::unsatCores() || options::fewerPreprocessingHoles())) {
-      // new assertions have a dependence on the node
-      PROOF( ProofManager::currentPM()->addDependence(itesRemoved, output[i]); )
-      while(n < output.size()) {
-        PROOF( ProofManager::currentPM()->addDependence(output[n], output[i]); )
-        ++n;
-      }
+  Node itesRemoved = run(assertion, newAsserts, newSkolems, false, false);
+  // In some calling contexts, not necessary to report dependence information.
+  if (reportDeps
+      && (options::unsatCores() || options::fewerPreprocessingHoles()))
+  {
+    // new assertions have a dependence on the node
+    PROOF(ProofManager::currentPM()->addDependence(itesRemoved, assertion);)
+    unsigned n = 0;
+    while (n < newAsserts.size())
+    {
+      PROOF(ProofManager::currentPM()->addDependence(newAsserts[n].getProven(),
+                                                     assertion);)
+      ++n;
     }
-    output[i] = itesRemoved;
   }
+  // The rewriting of assertion can be justified by the term conversion proof
+  // generator d_tpg.
+  return theory::TrustNode::mkTrustRewrite(assertion, itesRemoved, d_tpg.get());
 }
 
-Node RemoveTermFormulas::run(TNode node, std::vector<Node>& output,
-                    IteSkolemMap& iteSkolemMap, bool inQuant, bool inTerm) {
+Node RemoveTermFormulas::run(TNode node,
+                             std::vector<theory::TrustNode>& output,
+                             std::vector<Node>& newSkolems,
+                             bool inQuant,
+                             bool inTerm)
+{
   // Current node
   Debug("ite") << "removeITEs(" << node << ")" << " " << inQuant << " " << inTerm << endl;
 
-  //if(node.isVar() || node.isConst()){
-  //   (options::biasedITERemoval() && !containsTermITE(node))){
-  //if(node.isVar()){
-  //  return Node(node);
-  //}
   if( node.getKind()==kind::INST_PATTERN_LIST ){
     return Node(node);
   }
@@ -84,6 +92,8 @@ Node RemoveTermFormulas::run(TNode node, std::vector<Node>& output,
   TypeNode nodeType = node.getType();
   Node skolem;
   Node newAssertion;
+  // the exists form of the assertion
+  ProofGenerator* newAssertionPg = nullptr;
   // Handle non-Boolean ITEs here. Boolean ones (within terms) are handled
   // in the "non-variable Boolean term within term" case below.
   if (node.getKind() == kind::ITE && !nodeType.isBoolean())
@@ -96,15 +106,39 @@ Node RemoveTermFormulas::run(TNode node, std::vector<Node>& output,
       if (skolem.isNull())
       {
         // Make the skolem to represent the ITE
-        skolem = nodeManager->mkSkolem(
+        SkolemManager* sm = nodeManager->getSkolemManager();
+        skolem = sm->mkPurifySkolem(
+            node,
             "termITE",
-            nodeType,
             "a variable introduced due to term-level ITE removal");
         d_skolem_cache.insert(node, skolem);
 
         // The new assertion
         newAssertion = nodeManager->mkNode(
             kind::ITE, node[0], skolem.eqNode(node[1]), skolem.eqNode(node[2]));
+
+        // we justify it internally
+        if (isProofEnabled())
+        {
+          // ---------------------- REMOVE_TERM_FORMULA_AXIOM
+          // (ite node[0]
+          //      (= node node[1])            ------------- MACRO_SR_PRED_INTRO
+          //      (= node node[2]))           node = skolem
+          // ------------------------------------------ MACRO_SR_PRED_TRANSFORM
+          // (ite node[0] (= skolem node[1]) (= skolem node[2]))
+          //
+          // Note that the MACRO_SR_PRED_INTRO step holds due to conversion
+          // of skolem into its witness form, which is node.
+          Node axiom = getAxiomFor(node);
+          d_lp->addStep(axiom, PfRule::REMOVE_TERM_FORMULA_AXIOM, {}, {node});
+          Node eq = node.eqNode(skolem);
+          d_lp->addStep(eq, PfRule::MACRO_SR_PRED_INTRO, {}, {eq});
+          d_lp->addStep(newAssertion,
+                        PfRule::MACRO_SR_PRED_TRANSFORM,
+                        {axiom, eq},
+                        {newAssertion});
+          newAssertionPg = d_lp.get();
+        }
       }
     }
   }
@@ -117,9 +151,10 @@ Node RemoveTermFormulas::run(TNode node, std::vector<Node>& output,
       if (skolem.isNull())
       {
         // Make the skolem to represent the lambda
-        skolem = nodeManager->mkSkolem(
+        SkolemManager* sm = nodeManager->getSkolemManager();
+        skolem = sm->mkPurifySkolem(
+            node,
             "lambdaF",
-            nodeType,
             "a function introduced due to term-level lambda removal");
         d_skolem_cache.insert(node, skolem);
 
@@ -135,6 +170,15 @@ Node RemoveTermFormulas::run(TNode node, std::vector<Node>& output,
         children.push_back(skolem_app.eqNode(node[1]));
         // axiom defining skolem
         newAssertion = nodeManager->mkNode(kind::FORALL, children);
+
+        // Lambda lifting is trivial to justify, hence we don't set a proof
+        // generator here. In particular, replacing the skolem introduced
+        // here with its original lambda ensures the new assertion rewrites
+        // to true.
+        // For example, if (lambda y. t[y]) has skolem k, then this lemma is:
+        //   forall x. k(x)=t[x]
+        // whose witness form rewrites
+        //   forall x. (lambda y. t[y])(x)=t[x] --> forall x. t[x]=t[x] --> true
       }
     }
   }
@@ -148,10 +192,12 @@ Node RemoveTermFormulas::run(TNode node, std::vector<Node>& output,
       skolem = getSkolemForNode(node);
       if (skolem.isNull())
       {
-        // Make the skolem to witness the choice
-        skolem = nodeManager->mkSkolem(
+        // Make the skolem to witness the choice, which notice is handled
+        // as a special case within SkolemManager::mkPurifySkolem.
+        SkolemManager* sm = nodeManager->getSkolemManager();
+        skolem = sm->mkPurifySkolem(
+            node,
             "witnessK",
-            nodeType,
             "a skolem introduced due to term-level witness removal");
         d_skolem_cache.insert(node, skolem);
 
@@ -160,6 +206,27 @@ Node RemoveTermFormulas::run(TNode node, std::vector<Node>& output,
         // The new assertion is the assumption that the body
         // of the witness operator holds for the Skolem
         newAssertion = node[1].substitute(node[0][0], skolem);
+
+        // Get the proof generator, if one exists, which was responsible for
+        // constructing this witness term. This may not exist, in which case
+        // the witness term was trivial to justify. This is the case e.g. for
+        // purification witness terms.
+        if (isProofEnabled())
+        {
+          Node existsAssertion =
+              nodeManager->mkNode(kind::EXISTS, node[0], node[1]);
+          // -------------------- from skolem manager
+          // (exists x. node[1])
+          // -------------------- SKOLEMIZE
+          // node[1] * { x -> skolem }
+          ProofGenerator* expg = sm->getProofGenerator(existsAssertion);
+          if (expg != nullptr)
+          {
+            d_lp->addLazyStep(existsAssertion, expg);
+          }
+          d_lp->addStep(newAssertion, PfRule::SKOLEMIZE, {existsAssertion}, {});
+          newAssertionPg = d_lp.get();
+        }
       }
     }
   }
@@ -175,12 +242,23 @@ Node RemoveTermFormulas::run(TNode node, std::vector<Node>& output,
       // Skolems introduced for Boolean formulas appearing in terms have a
       // special kind (BOOLEAN_TERM_VARIABLE) that ensures they are handled
       // properly in theory combination. We must use this kind here instead of a
-      // generic skolem.
-      skolem = nodeManager->mkBooleanTermVariable();
+      // generic skolem. Notice that the name/comment are currently ignored
+      // within SkolemManager::mkPurifySkolem, since BOOLEAN_TERM_VARIABLE
+      // variables cannot be given names.
+      SkolemManager* sm = nodeManager->getSkolemManager();
+      skolem = sm->mkPurifySkolem(
+          node,
+          "btvK",
+          "a Boolean term variable introduced during term formula removal",
+          NodeManager::SKOLEM_BOOL_TERM_VAR);
       d_skolem_cache.insert(node, skolem);
 
       // The new assertion
       newAssertion = skolem.eqNode(node);
+
+      // Boolean term removal is trivial to justify, hence we don't set a proof
+      // generator here. It is trivial to justify since it is an instance of
+      // purification, which is justified by conversion to witness forms.
     }
   }
 
@@ -192,15 +270,43 @@ Node RemoveTermFormulas::run(TNode node, std::vector<Node>& output,
     // if the skolem was introduced in this call
     if (!newAssertion.isNull())
     {
+      // if proofs are enabled
+      if (isProofEnabled())
+      {
+        // justify the introduction of the skolem
+        // ------------------- MACRO_SR_PRED_INTRO
+        // t = witness x. x=t
+        // The above step is trivial, since the skolems introduced above are
+        // all purification skolems. We record this equality in the term
+        // conversion proof generator.
+        d_tpg->addRewriteStep(node,
+                              skolem,
+                              PfRule::MACRO_SR_PRED_INTRO,
+                              {},
+                              {node.eqNode(skolem)});
+        // justify the axiom that defines the skolem, if not already done so
+        if (newAssertionPg == nullptr)
+        {
+          // Should have trivial justification if not yet provided. This is the
+          // case of lambda lifting and Boolean term removal.
+          // ---------------- MACRO_SR_PRED_INTRO
+          // newAssertion
+          d_lp->addStep(
+              newAssertion, PfRule::MACRO_SR_PRED_INTRO, {}, {newAssertion});
+        }
+      }
       Debug("ite") << "*** term formula removal introduced " << skolem
                    << " for " << node << std::endl;
 
       // Remove ITEs from the new assertion, rewrite it and push it to the
       // output
-      newAssertion = run(newAssertion, output, iteSkolemMap, false, false);
+      newAssertion = run(newAssertion, output, newSkolems, false, false);
+
+      theory::TrustNode trna =
+          theory::TrustNode::mkTrustLemma(newAssertion, d_lp.get());
 
-      iteSkolemMap[skolem] = output.size();
-      output.push_back(newAssertion);
+      output.push_back(trna);
+      newSkolems.push_back(skolem);
     }
 
     // The representation is now the skolem
@@ -225,7 +331,7 @@ Node RemoveTermFormulas::run(TNode node, std::vector<Node>& output,
   }
   // Remove the ITEs from the children
   for(TNode::const_iterator it = node.begin(), end = node.end(); it != end; ++it) {
-    Node newChild = run(*it, output, iteSkolemMap, inQuant, inTerm);
+    Node newChild = run(*it, output, newSkolems, inQuant, inTerm);
     somethingChanged |= (newChild != *it);
     newChildren.push_back(newChild);
   }
@@ -305,4 +411,27 @@ bool RemoveTermFormulas::hasNestedTermChildren( TNode node ) {
          // dont' worry about FORALL or EXISTS (handled separately)
 }
 
+Node RemoveTermFormulas::getAxiomFor(Node n)
+{
+  NodeManager* nm = NodeManager::currentNM();
+  Kind k = n.getKind();
+  if (k == kind::ITE)
+  {
+    return nm->mkNode(kind::ITE, n[0], n.eqNode(n[1]), n.eqNode(n[2]));
+  }
+  return Node::null();
+}
+
+void RemoveTermFormulas::setProofChecker(ProofChecker* pc)
+{
+  if (d_tpg == nullptr)
+  {
+    d_pnm.reset(new ProofNodeManager(pc));
+    d_tpg.reset(new TConvProofGenerator(d_pnm.get()));
+    d_lp.reset(new LazyCDProof(d_pnm.get()));
+  }
+}
+
+bool RemoveTermFormulas::isProofEnabled() const { return d_pnm != nullptr; }
+
 }/* CVC4 namespace */
diff --git a/src/smt/term_formula_removal.h b/src/smt/term_formula_removal.h
index 9ec12cb12..d4c22b78b 100644
--- a/src/smt/term_formula_removal.h
+++ b/src/smt/term_formula_removal.h
@@ -23,8 +23,12 @@
 
 #include "context/cdinsert_hashmap.h"
 #include "context/context.h"
+#include "expr/lazy_proof.h"
 #include "expr/node.h"
+#include "expr/term_conversion_proof_generator.h"
 #include "smt/dump.h"
+#include "theory/eager_proof_generator.h"
+#include "theory/trust_node.h"
 #include "util/bool.h"
 #include "util/hash.h"
 
@@ -33,6 +37,89 @@ namespace CVC4 {
 typedef std::unordered_map<Node, unsigned, NodeHashFunction> IteSkolemMap;
 
 class RemoveTermFormulas {
+ public:
+  RemoveTermFormulas(context::UserContext* u);
+  ~RemoveTermFormulas();
+
+  /**
+   * By introducing skolem variables, this function removes all occurrences of:
+   * (1) term ITEs,
+   * (2) terms of type Boolean that are not Boolean term variables,
+   * (3) lambdas, and
+   * (4) Hilbert choice expressions.
+   * from assertions.
+   * All additional assertions are pushed into assertions. iteSkolemMap
+   * contains a map from introduced skolem variables to the index in
+   * assertions containing the new definition created in conjunction
+   * with that skolem variable.
+   *
+   * As an example of (1):
+   *   f( (ite C 0 1)) = 2
+   * becomes
+   *   f( k ) = 2 ^ ite( C, k=0, k=1 )
+   *
+   * As an example of (2):
+   *   g( (and C1 C2) ) = 3
+   * becomes
+   *   g( k ) = 3 ^ ( k <=> (and C1 C2) )
+   *
+   * As an example of (3):
+   *   (lambda x. t[x]) = f
+   * becomes
+   *   (forall x. k(x) = t[x]) ^ k = f
+   * where k is a fresh skolem function.
+   * This is sometimes called "lambda lifting"
+   *
+   * As an example of (4):
+   *   (witness x. P( x ) ) = t
+   * becomes
+   *   P( k ) ^ k = t
+   * where k is a fresh skolem constant.
+   *
+   * With reportDeps true, report reasoning dependences to the proof
+   * manager (for unsat cores).
+   *
+   * @param assertion The assertion to remove term formulas from
+   * @param newAsserts The new assertions corresponding to axioms for newly
+   * introduced skolems.
+   * @param newSkolems The skolems corresponding to each of the newAsserts.
+   * @param reportDeps Used for unsat cores in the old proof infrastructure.
+   * @return a trust node of kind TrustNodeKind::REWRITE whose
+   * right hand side is assertion after removing term formulas, and the proof
+   * generator (if provided) that can prove the equivalence.
+   */
+  theory::TrustNode run(Node assertion,
+                        std::vector<theory::TrustNode>& newAsserts,
+                        std::vector<Node>& newSkolems,
+                        bool reportDeps = false);
+
+  /**
+   * Substitute under node using pre-existing cache.  Do not remove
+   * any ITEs not seen during previous runs.
+   */
+  Node replace(TNode node, bool inQuant = false, bool inTerm = false) const;
+
+  /** Returns true if e contains a term ite. */
+  bool containsTermITE(TNode e) const;
+
+  /** Garbage collects non-context dependent data-structures. */
+  void garbageCollect();
+
+  /**
+   * Set proof checker, which signals this class to enable proofs using the
+   * given checker.
+   */
+  void setProofChecker(ProofChecker* pc);
+
+  /**
+   * Get axiom for term n. This returns the axiom that this class uses to
+   * eliminate the term n, which is determined by its top-most symbol. For
+   * example, if n is (ite n1 n2 n3), this returns the formula:
+   *   (ite n1 (= (ite n1 n2 n3) n2) (= (ite n1 n2 n3) n3))
+   */
+  static Node getAxiomFor(Node n);
+
+ private:
   typedef context::
       CDInsertHashMap<std::pair<Node, int>,
                       Node,
@@ -77,50 +164,18 @@ class RemoveTermFormulas {
   inline Node getSkolemForNode(Node node) const;
 
   static bool hasNestedTermChildren( TNode node );
-public:
-
-  RemoveTermFormulas(context::UserContext* u);
-  ~RemoveTermFormulas();
 
+  /** A proof node manager */
+  std::unique_ptr<ProofNodeManager> d_pnm;
   /**
-   * By introducing skolem variables, this function removes all occurrences of:
-   * (1) term ITEs,
-   * (2) terms of type Boolean that are not Boolean term variables,
-   * (3) lambdas, and
-   * (4) Hilbert choice expressions.
-   * from assertions.
-   * All additional assertions are pushed into assertions. iteSkolemMap
-   * contains a map from introduced skolem variables to the index in
-   * assertions containing the new definition created in conjunction
-   * with that skolem variable.
-   *
-   * As an example of (1):
-   *   f( (ite C 0 1)) = 2
-   * becomes
-   *   f( k ) = 2 ^ ite( C, k=0, k=1 )
-   *
-   * As an example of (2):
-   *   g( (and C1 C2) ) = 3
-   * becomes
-   *   g( k ) = 3 ^ ( k <=> (and C1 C2) )
-   *
-   * As an example of (3):
-   *   (lambda x. t[x]) = f
-   * becomes
-   *   (forall x. k(x) = t[x]) ^ k = f
-   * where k is a fresh skolem function.
-   * This is sometimes called "lambda lifting"
-   *
-   * As an example of (4):
-   *   (witness x. P( x ) ) = t
-   * becomes
-   *   P( k ) ^ k = t
-   * where k is a fresh skolem constant.
-   *
-   * With reportDeps true, report reasoning dependences to the proof
-   * manager (for unsat cores).
+   * A proof generator for the term conversion.
    */
-  void run(std::vector<Node>& assertions, IteSkolemMap& iteSkolemMap, bool reportDeps = false);
+  std::unique_ptr<TConvProofGenerator> d_tpg;
+  /**
+   * A proof generator for skolems we introduce that are based on axioms that
+   * this class is responsible for.
+   */
+  std::unique_ptr<LazyCDProof> d_lp;
 
   /**
    * Removes terms of the form (1), (2), (3) described above from node.
@@ -133,20 +188,14 @@ public:
    * inTerm is whether we are are processing node in a "term" position, that is, it is a subterm
    *        of a parent term that is not a Boolean connective.
    */
-  Node run(TNode node, std::vector<Node>& additionalAssertions,
-           IteSkolemMap& iteSkolemMap, bool inQuant, bool inTerm);
-
-  /**
-   * Substitute under node using pre-existing cache.  Do not remove
-   * any ITEs not seen during previous runs.
-   */
-  Node replace(TNode node, bool inQuant = false, bool inTerm = false) const;
-
-  /** Returns true if e contains a term ite. */
-  bool containsTermITE(TNode e) const;
-
-  /** Garbage collects non-context dependent data-structures. */
-  void garbageCollect();
+  Node run(TNode node,
+           std::vector<theory::TrustNode>& newAsserts,
+           std::vector<Node>& newSkolems,
+           bool inQuant,
+           bool inTerm);
+
+  /** Whether proofs are enabled */
+  bool isProofEnabled() const;
 };/* class RemoveTTE */
 
 }/* CVC4 namespace */
diff --git a/src/theory/builtin/proof_checker.cpp b/src/theory/builtin/proof_checker.cpp
index 59f405337..817d21fdf 100644
--- a/src/theory/builtin/proof_checker.cpp
+++ b/src/theory/builtin/proof_checker.cpp
@@ -15,6 +15,7 @@
 #include "theory/builtin/proof_checker.h"
 
 #include "expr/skolem_manager.h"
+#include "smt/term_formula_removal.h"
 #include "theory/rewriter.h"
 #include "theory/theory.h"
 
@@ -61,6 +62,7 @@ void BuiltinProofRuleChecker::registerTo(ProofChecker* pc)
   pc->registerChecker(PfRule::MACRO_SR_PRED_TRANSFORM, this);
   pc->registerChecker(PfRule::THEORY_REWRITE, this);
   pc->registerChecker(PfRule::PREPROCESS, this);
+  pc->registerChecker(PfRule::REMOVE_TERM_FORMULA_AXIOM, this);
 }
 
 Node BuiltinProofRuleChecker::applyTheoryRewrite(Node n, bool preRewrite)
@@ -321,6 +323,12 @@ Node BuiltinProofRuleChecker::checkInternal(PfRule id,
     }
     return args[0];
   }
+  else if (id == PfRule::REMOVE_TERM_FORMULA_AXIOM)
+  {
+    Assert(children.empty());
+    Assert(args.size() == 1);
+    return RemoveTermFormulas::getAxiomFor(args[0]);
+  }
   else if (id == PfRule::PREPROCESS)
   {
     Assert(children.empty());
diff --git a/src/theory/theory_engine.cpp b/src/theory/theory_engine.cpp
index 70e744acf..9955be850 100644
--- a/src/theory/theory_engine.cpp
+++ b/src/theory/theory_engine.cpp
@@ -1610,10 +1610,25 @@ theory::LemmaStatus TheoryEngine::lemma(TNode node,
                      << CheckSatCommand(n.toExpr());
   }
 
-  // the assertion pipeline storing the lemmas
-  AssertionPipeline lemmas;
   // call preprocessor
-  d_tpp.preprocess(node, lemmas, preprocess);
+  std::vector<TrustNode> newLemmas;
+  std::vector<Node> newSkolems;
+  TrustNode tlemma = d_tpp.preprocess(node, newLemmas, newSkolems, preprocess);
+
+  // must use an assertion pipeline due to decision engine below
+  AssertionPipeline lemmas;
+  // make the assertion pipeline
+  lemmas.push_back(tlemma.getNode());
+  lemmas.updateRealAssertionsEnd();
+  Assert(newSkolems.size() == newLemmas.size());
+  for (size_t i = 0, nsize = newLemmas.size(); i < nsize; i++)
+  {
+    // store skolem mapping here
+    IteSkolemMap& imap = lemmas.getIteSkolemMap();
+    imap[newSkolems[i]] = lemmas.size();
+    lemmas.push_back(newLemmas[i].getNode());
+  }
+
   // assert lemmas to prop engine
   for (size_t i = 0, lsize = lemmas.size(); i < lsize; ++i)
   {
diff --git a/src/theory/theory_preprocessor.cpp b/src/theory/theory_preprocessor.cpp
index b12916b7d..328c65fcb 100644
--- a/src/theory/theory_preprocessor.cpp
+++ b/src/theory/theory_preprocessor.cpp
@@ -36,41 +36,52 @@ TheoryPreprocessor::~TheoryPreprocessor() {}
 
 void TheoryPreprocessor::clearCache() { d_ppCache.clear(); }
 
-void TheoryPreprocessor::preprocess(TNode node,
-                                    preprocessing::AssertionPipeline& lemmas,
-                                    bool doTheoryPreprocess)
+TrustNode TheoryPreprocessor::preprocess(TNode node,
+                                         std::vector<TrustNode>& newLemmas,
+                                         std::vector<Node>& newSkolems,
+                                         bool doTheoryPreprocess)
 {
   // Run theory preprocessing, maybe
   Node ppNode = doTheoryPreprocess ? theoryPreprocess(node) : Node(node);
 
   // Remove the ITEs
   Trace("te-tform-rm") << "Remove term formulas from " << ppNode << std::endl;
-  lemmas.push_back(ppNode);
-  lemmas.updateRealAssertionsEnd();
-  d_tfr.run(lemmas.ref(), lemmas.getIteSkolemMap());
-  Trace("te-tform-rm") << "..done " << lemmas[0] << std::endl;
+  TrustNode ttfr = d_tfr.run(ppNode, newLemmas, newSkolems, false);
+  Trace("te-tform-rm") << "..done " << ttfr.getNode() << std::endl;
+  Node retNode = ttfr.getNode();
 
   if (Debug.isOn("lemma-ites"))
   {
-    Debug("lemma-ites") << "removed ITEs from lemma: " << ppNode << endl;
-    Debug("lemma-ites") << " + now have the following " << lemmas.size()
+    Debug("lemma-ites") << "removed ITEs from lemma: " << ttfr.getNode()
+                        << endl;
+    Debug("lemma-ites") << " + now have the following " << newLemmas.size()
                         << " lemma(s):" << endl;
-    for (std::vector<Node>::const_iterator i = lemmas.begin();
-         i != lemmas.end();
-         ++i)
+    for (size_t i = 0, lsize = newLemmas.size(); i <= lsize; ++i)
     {
-      Debug("lemma-ites") << " + " << *i << endl;
+      Debug("lemma-ites") << " + " << newLemmas[i].getNode() << endl;
     }
     Debug("lemma-ites") << endl;
   }
 
+  retNode = Rewriter::rewrite(retNode);
+
   // now, rewrite the lemmas
-  Node retLemma;
-  for (size_t i = 0, lsize = lemmas.size(); i < lsize; ++i)
+  for (size_t i = 0, lsize = newLemmas.size(); i < lsize; ++i)
   {
-    Node rewritten = Rewriter::rewrite(lemmas[i]);
-    lemmas.replace(i, rewritten);
+    // get the trust node to process
+    TrustNode trn = newLemmas[i];
+    Assert(trn.getKind() == TrustNodeKind::LEMMA);
+    Node assertion = trn.getNode();
+    // rewrite, which is independent of d_tpg, since additional lemmas
+    // are justified separately.
+    Node rewritten = Rewriter::rewrite(assertion);
+    if (assertion != rewritten)
+    {
+      // not tracking proofs, just make new
+      newLemmas[i] = TrustNode::mkTrustLemma(rewritten, nullptr);
+    }
   }
+  return TrustNode::mkTrustRewrite(node, retNode, nullptr);
 }
 
 struct preprocess_stack_element
diff --git a/src/theory/theory_preprocessor.h b/src/theory/theory_preprocessor.h
index 2488cf162..2f3813e68 100644
--- a/src/theory/theory_preprocessor.h
+++ b/src/theory/theory_preprocessor.h
@@ -46,15 +46,24 @@ class TheoryPreprocessor
   /** Clear the cache of this class */
   void clearCache();
   /**
-   * Preprocesses node and stores it along with lemmas generated by
-   * preprocessing into the assertion pipeline lemmas. The (optional) argument
-   * lcp is the proof that stores a proof of all top-level formulas in lemmas,
-   * assuming that lcp initially contains a proof of node. The flag
-   * doTheoryPreprocess is whether we should run theory-specific preprocessing.
+   * Preprocesses the given assertion node. It returns a TrustNode of kind
+   * TrustNodeKind::REWRITE indicating the preprocessed form of node. It stores
+   * additional lemmas in newLemmas, which are trust nodes of kind
+   * TrustNodeKind::LEMMA. These correspond to e.g. lemmas corresponding to ITE
+   * removal. For each lemma in newLemmas, we add the corresponding skolem that
+   * the lemma defines. The flag doTheoryPreprocess is whether we should run
+   * theory-specific preprocessing.
+   *
+   * @param node The assertion to preprocess,
+   * @param newLemmas The lemmas to add to the set of assertions,
+   * @param newSkolems The skolems that newLemmas correspond to,
+   * @param doTheoryPreprocess whether to run theory-specific preprocessing.
+   * @return The trust node corresponding to rewriting node via preprocessing.
    */
-  void preprocess(TNode node,
-                  preprocessing::AssertionPipeline& lemmas,
-                  bool doTheoryPreprocess);
+  TrustNode preprocess(TNode node,
+                       std::vector<TrustNode>& newLemmas,
+                       std::vector<Node>& newSkolems,
+                       bool doTheoryPreprocess);
   /**
    * Runs theory specific preprocessing on the non-Boolean parts of
    * the formula.  This is only called on input assertions, after ITEs