Add new interfaces to term formula removal and theory preprocess (#5717)

This is in preparation for lazy lemmas for term formula removal.

5 files changed, 155 insertions, 15 deletions

diff --git a/src/prop/theory_proxy.cpp b/src/prop/theory_proxy.cpp
index 7559e4015..35602b8b3 100644
--- a/src/prop/theory_proxy.cpp
+++ b/src/prop/theory_proxy.cpp
@@ -174,7 +174,8 @@ theory::TrustNode TheoryProxy::preprocessLemma(
     std::vector<Node>& newSkolems,
     bool doTheoryPreprocess)
 {
-  return d_tpp.preprocessLemma(trn, newLemmas, newSkolems, doTheoryPreprocess);
+  return d_tpp.preprocessLemma(
+      trn, newLemmas, newSkolems, doTheoryPreprocess, true);
 }
 
 theory::TrustNode TheoryProxy::preprocess(
@@ -184,7 +185,7 @@ theory::TrustNode TheoryProxy::preprocess(
     bool doTheoryPreprocess)
 {
   theory::TrustNode pnode =
-      d_tpp.preprocess(node, newLemmas, newSkolems, doTheoryPreprocess);
+      d_tpp.preprocess(node, newLemmas, newSkolems, doTheoryPreprocess, true);
   return pnode;
 }
 
diff --git a/src/smt/term_formula_removal.cpp b/src/smt/term_formula_removal.cpp
index ae3039ffa..42166d074 100644
--- a/src/smt/term_formula_removal.cpp
+++ b/src/smt/term_formula_removal.cpp
@@ -28,7 +28,12 @@ namespace CVC4 {
 
 RemoveTermFormulas::RemoveTermFormulas(context::UserContext* u,
                                        ProofNodeManager* pnm)
-    : d_tfCache(u), d_skolem_cache(u), d_pnm(pnm), d_tpg(nullptr), d_lp(nullptr)
+    : d_tfCache(u),
+      d_skolem_cache(u),
+      d_lemmaCache(u),
+      d_pnm(pnm),
+      d_tpg(nullptr),
+      d_lp(nullptr)
 {
   // enable proofs if necessary
   if (d_pnm != nullptr)
@@ -48,7 +53,7 @@ RemoveTermFormulas::RemoveTermFormulas(context::UserContext* u,
 RemoveTermFormulas::~RemoveTermFormulas() {}
 
 theory::TrustNode RemoveTermFormulas::run(
-    Node assertion,
+    TNode assertion,
     std::vector<theory::TrustNode>& newAsserts,
     std::vector<Node>& newSkolems,
     bool fixedPoint)
@@ -81,6 +86,13 @@ theory::TrustNode RemoveTermFormulas::run(
   return theory::TrustNode::mkTrustRewrite(assertion, itesRemoved, d_tpg.get());
 }
 
+theory::TrustNode RemoveTermFormulas::run(TNode assertion)
+{
+  std::vector<theory::TrustNode> newAsserts;
+  std::vector<Node> newSkolems;
+  return run(assertion, newAsserts, newSkolems, false);
+}
+
 theory::TrustNode RemoveTermFormulas::runLemma(
     theory::TrustNode lem,
     std::vector<theory::TrustNode>& newAsserts,
@@ -121,7 +133,7 @@ theory::TrustNode RemoveTermFormulas::runLemma(
   return theory::TrustNode::mkTrustLemma(newAssertion, d_lp.get());
 }
 
-Node RemoveTermFormulas::runInternal(Node assertion,
+Node RemoveTermFormulas::runInternal(TNode assertion,
                                      std::vector<theory::TrustNode>& output,
                                      std::vector<Node>& newSkolems)
 {
@@ -481,6 +493,9 @@ Node RemoveTermFormulas::runCurrent(std::pair<Node, uint32_t>& curr,
 
       newLem = theory::TrustNode::mkTrustLemma(newAssertion, d_lp.get());
 
+      // store in the lemma cache
+      d_lemmaCache.insert(skolem, newLem);
+
       Trace("rtf-proof-debug") << "Checking closed..." << std::endl;
       newLem.debugCheckClosed("rtf-proof-debug",
                               "RemoveTermFormulas::run:new_assert");
@@ -494,10 +509,10 @@ Node RemoveTermFormulas::runCurrent(std::pair<Node, uint32_t>& curr,
   return Node::null();
 }
 
-Node RemoveTermFormulas::getSkolemForNode(Node node) const
+Node RemoveTermFormulas::getSkolemForNode(Node k) const
 {
   context::CDInsertHashMap<Node, Node, NodeHashFunction>::const_iterator itk =
-      d_skolem_cache.find(node);
+      d_skolem_cache.find(k);
   if (itk != d_skolem_cache.end())
   {
     return itk->second;
@@ -505,6 +520,50 @@ Node RemoveTermFormulas::getSkolemForNode(Node node) const
   return Node::null();
 }
 
+bool RemoveTermFormulas::getSkolems(
+    TNode n, std::unordered_set<Node, NodeHashFunction>& skolems) const
+{
+  // if n was unchanged by term formula removal, just return immediately
+  std::pair<Node, uint32_t> initial(n, d_rtfc.initialValue());
+  TermFormulaCache::const_iterator itc = d_tfCache.find(initial);
+  if (itc != d_tfCache.end())
+  {
+    if (itc->second == n)
+    {
+      return false;
+    }
+  }
+  // otherwise, traverse it
+  bool ret = false;
+  std::unordered_set<TNode, TNodeHashFunction> visited;
+  std::unordered_set<TNode, TNodeHashFunction>::iterator it;
+  std::vector<TNode> visit;
+  TNode cur;
+  visit.push_back(n);
+  do
+  {
+    cur = visit.back();
+    visit.pop_back();
+    it = visited.find(cur);
+    if (it == visited.end())
+    {
+      visited.insert(cur);
+      if (cur.isVar())
+      {
+        if (d_lemmaCache.find(cur) != d_lemmaCache.end())
+        {
+          // technically could already be in skolems if skolems was non-empty,
+          // regardless set return value to true.
+          skolems.insert(cur);
+          ret = true;
+        }
+      }
+      visit.insert(visit.end(), cur.begin(), cur.end());
+    }
+  } while (!visit.empty());
+  return ret;
+}
+
 Node RemoveTermFormulas::getAxiomFor(Node n)
 {
   NodeManager* nm = NodeManager::currentNM();
@@ -516,6 +575,17 @@ Node RemoveTermFormulas::getAxiomFor(Node n)
   return Node::null();
 }
 
+theory::TrustNode RemoveTermFormulas::getLemmaForSkolem(TNode n) const
+{
+  context::CDInsertHashMap<Node, theory::TrustNode, NodeHashFunction>::
+      const_iterator it = d_lemmaCache.find(n);
+  if (it == d_lemmaCache.end())
+  {
+    return theory::TrustNode::null();
+  }
+  return (*it).second;
+}
+
 ProofGenerator* RemoveTermFormulas::getTConvProofGenerator()
 {
   return d_tpg.get();
diff --git a/src/smt/term_formula_removal.h b/src/smt/term_formula_removal.h
index e5453254c..ac2644182 100644
--- a/src/smt/term_formula_removal.h
+++ b/src/smt/term_formula_removal.h
@@ -91,11 +91,17 @@ class RemoveTermFormulas {
    * 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,
+  theory::TrustNode run(TNode assertion,
                         std::vector<theory::TrustNode>& newAsserts,
                         std::vector<Node>& newSkolems,
                         bool fixedPoint = false);
   /**
+   * Same as above, but does not track lemmas, and does not run to fixed point.
+   * The relevant lemmas can be extracted by the caller later using getSkolems
+   * and getLemmaForSkolem.
+   */
+  theory::TrustNode run(TNode assertion);
+  /**
    * Same as above, but transforms a lemma, returning a LEMMA trust node that
    * proves the same formula as lem with term formulas removed.
    */
@@ -120,6 +126,24 @@ class RemoveTermFormulas {
    */
   static Node getAxiomFor(Node n);
 
+  /**
+   * Get the set of skolems introduced by this class that occur in node n,
+   * add them to skolems.
+   *
+   * This method uses an optimization that returns false immediately if n
+   * was unchanged by term formula removal, based on the initial context.
+   *
+   * Return true if any nodes were added to skolems.
+   */
+  bool getSkolems(TNode n,
+                  std::unordered_set<Node, NodeHashFunction>& skolems) const;
+
+  /**
+   * Get the lemma for the skolem, or the null node if k is not a skolem this
+   * class introduced.
+   */
+  theory::TrustNode getLemmaForSkolem(TNode k) const;
+
  private:
   typedef context::CDInsertHashMap<
       std::pair<Node, uint32_t>,
@@ -153,6 +177,11 @@ class RemoveTermFormulas {
    *   d_tfCache[<ite( G, a, b ),0>] = d_tfCache[<ite( G, a, b ),1>] = k.
    */
   context::CDInsertHashMap<Node, Node, NodeHashFunction> d_skolem_cache;
+  /**
+   * Mapping from skolems to their corresponding lemma.
+   */
+  context::CDInsertHashMap<Node, theory::TrustNode, NodeHashFunction>
+      d_lemmaCache;
 
   /** gets the skolem for node
    *
@@ -186,7 +215,7 @@ class RemoveTermFormulas {
    * This uses a term-context-sensitive stack to process assertion. It returns
    * the version of assertion with all term formulas removed.
    */
-  Node runInternal(Node assertion,
+  Node runInternal(TNode assertion,
                    std::vector<theory::TrustNode>& newAsserts,
                    std::vector<Node>& newSkolems);
   /**
diff --git a/src/theory/theory_preprocessor.cpp b/src/theory/theory_preprocessor.cpp
index c56554b53..9a425fc1c 100644
--- a/src/theory/theory_preprocessor.cpp
+++ b/src/theory/theory_preprocessor.cpp
@@ -84,7 +84,8 @@ TheoryPreprocessor::~TheoryPreprocessor() {}
 TrustNode TheoryPreprocessor::preprocess(TNode node,
                                          std::vector<TrustNode>& newLemmas,
                                          std::vector<Node>& newSkolems,
-                                         bool doTheoryPreprocess)
+                                         bool doTheoryPreprocess,
+                                         bool fixedPoint)
 {
   // In this method, all rewriting steps of node are stored in d_tpg.
 
@@ -101,7 +102,7 @@ TrustNode TheoryPreprocessor::preprocess(TNode node,
   }
 
   // Remove the ITEs, fixed point
-  TrustNode ttfr = d_tfr.run(ppNode, newLemmas, newSkolems, true);
+  TrustNode ttfr = d_tfr.run(ppNode, newLemmas, newSkolems, fixedPoint);
   Node rtfNode = ttfr.isNull() ? ppNode : ttfr.getNode();
 
   if (Debug.isOn("lemma-ites"))
@@ -222,15 +223,24 @@ TrustNode TheoryPreprocessor::preprocess(TNode node,
   return tret;
 }
 
+TrustNode TheoryPreprocessor::preprocess(TNode node, bool doTheoryPreprocess)
+{
+  // ignore lemmas, no fixed point
+  std::vector<TrustNode> newLemmas;
+  std::vector<Node> newSkolems;
+  return preprocess(node, newLemmas, newSkolems, doTheoryPreprocess, false);
+}
+
 TrustNode TheoryPreprocessor::preprocessLemma(TrustNode node,
                                               std::vector<TrustNode>& newLemmas,
                                               std::vector<Node>& newSkolems,
-                                              bool doTheoryPreprocess)
+                                              bool doTheoryPreprocess,
+                                              bool fixedPoint)
 {
   // what was originally proven
   Node lemma = node.getProven();
   TrustNode tplemma =
-      preprocess(lemma, newLemmas, newSkolems, doTheoryPreprocess);
+      preprocess(lemma, newLemmas, newSkolems, doTheoryPreprocess, fixedPoint);
   if (tplemma.isNull())
   {
     // no change needed
@@ -267,6 +277,21 @@ TrustNode TheoryPreprocessor::preprocessLemma(TrustNode node,
   return TrustNode::mkTrustLemma(lemmap, d_lp.get());
 }
 
+TrustNode TheoryPreprocessor::preprocessLemma(TrustNode node,
+                                              bool doTheoryPreprocess)
+{
+  // ignore lemmas, no fixed point
+  std::vector<TrustNode> newLemmas;
+  std::vector<Node> newSkolems;
+  return preprocessLemma(
+      node, newLemmas, newSkolems, doTheoryPreprocess, false);
+}
+
+RemoveTermFormulas& TheoryPreprocessor::getRemoveTermFormulas()
+{
+  return d_tfr;
+}
+
 struct preprocess_stack_element
 {
   TNode node;
diff --git a/src/theory/theory_preprocessor.h b/src/theory/theory_preprocessor.h
index 8bda3f5aa..d7c22270d 100644
--- a/src/theory/theory_preprocessor.h
+++ b/src/theory/theory_preprocessor.h
@@ -69,7 +69,13 @@ class TheoryPreprocessor
   TrustNode preprocess(TNode node,
                        std::vector<TrustNode>& newLemmas,
                        std::vector<Node>& newSkolems,
-                       bool doTheoryPreprocess);
+                       bool doTheoryPreprocess,
+                       bool fixedPoint);
+  /**
+   * Same as above, without lemma tracking or fixed point. Lemmas for skolems
+   * can be extracted from the RemoveTermFormulas utility.
+   */
+  TrustNode preprocess(TNode node, bool doTheoryPreprocess);
   /**
    * Same as above, but transforms the proof of node into a proof of the
    * preprocessed node and returns the LEMMA trust node.
@@ -84,7 +90,16 @@ class TheoryPreprocessor
   TrustNode preprocessLemma(TrustNode node,
                             std::vector<TrustNode>& newLemmas,
                             std::vector<Node>& newSkolems,
-                            bool doTheoryPreprocess);
+                            bool doTheoryPreprocess,
+                            bool fixedPoint);
+  /**
+   * Same as above, without lemma tracking or fixed point. Lemmas for skolems
+   * can be extracted from the RemoveTermFormulas utility.
+   */
+  TrustNode preprocessLemma(TrustNode node, bool doTheoryPreprocess);
+
+  /** Get the term formula removal utility */
+  RemoveTermFormulas& getRemoveTermFormulas();
 
  private:
   /**