2 files changed, 91 insertions, 7 deletions

diff --git a/src/theory/arith/nl/cad/cdcac.cpp b/src/theory/arith/nl/cad/cdcac.cpp
index d076438f5..df3ba37f3 100644
--- a/src/theory/arith/nl/cad/cdcac.cpp
+++ b/src/theory/arith/nl/cad/cdcac.cpp
@@ -39,11 +39,15 @@ namespace arith {
 namespace nl {
 namespace cad {
 
-CDCAC::CDCAC() {}
-
-CDCAC::CDCAC(const std::vector<poly::Variable>& ordering)
+CDCAC::CDCAC(context::Context* ctx,
+             ProofNodeManager* pnm,
+             const std::vector<poly::Variable>& ordering)
     : d_variableOrdering(ordering)
 {
+  if (pnm != nullptr)
+  {
+    d_proof.reset(new CADProofGenerator(ctx, pnm));
+  }
 }
 
 void CDCAC::reset()
@@ -120,6 +124,17 @@ std::vector<CACInterval> CDCAC::getUnsatIntervals(std::size_t cur_variable)
       if (!is_minus_infinity(get_lower(i))) l = m;
       if (!is_plus_infinity(get_upper(i))) u = m;
       res.emplace_back(CACInterval{i, l, u, m, d, {n}});
+      if (isProofEnabled())
+      {
+        d_proof->addDirect(
+            d_constraints.varMapper()(d_variableOrdering[cur_variable]),
+            d_constraints.varMapper(),
+            p,
+            d_assignment,
+            sc,
+            i,
+            n);
+      }
     }
   }
   pruneRedundantIntervals(res);
@@ -335,6 +350,10 @@ CACInterval CDCAC::intervalFromCharacterization(
 std::vector<CACInterval> CDCAC::getUnsatCover(std::size_t curVariable,
                                               bool returnFirstInterval)
 {
+  if (isProofEnabled())
+  {
+    d_proof->startRecursive();
+  }
   Trace("cdcac") << "Looking for unsat cover for "
                  << d_variableOrdering[curVariable] << std::endl;
   std::vector<CACInterval> intervals = getUnsatIntervals(curVariable);
@@ -373,6 +392,10 @@ std::vector<CACInterval> CDCAC::getUnsatCover(std::size_t curVariable,
       Trace("cdcac") << "Found full assignment: " << d_assignment << std::endl;
       return {};
     }
+    if (isProofEnabled())
+    {
+      d_proof->startScope();
+    }
     // Recurse to next variable
     auto cov = getUnsatCover(curVariable + 1);
     if (cov.empty())
@@ -391,6 +414,16 @@ std::vector<CACInterval> CDCAC::getUnsatCover(std::size_t curVariable,
         intervalFromCharacterization(characterization, curVariable, sample);
     newInterval.d_origins = collectConstraints(cov);
     intervals.emplace_back(newInterval);
+    if (isProofEnabled())
+    {
+      auto cell = d_proof->constructCell(
+          d_constraints.varMapper()(d_variableOrdering[curVariable]),
+          newInterval,
+          d_assignment,
+          sample,
+          d_constraints.varMapper());
+      d_proof->endScope(cell);
+    }
 
     if (returnFirstInterval)
     {
@@ -421,9 +454,31 @@ std::vector<CACInterval> CDCAC::getUnsatCover(std::size_t curVariable,
       Trace("cdcac") << "-> " << i.d_interval << std::endl;
     }
   }
+  if (isProofEnabled())
+  {
+    d_proof->endRecursive();
+  }
   return intervals;
 }
 
+void CDCAC::startNewProof()
+{
+  if (isProofEnabled())
+  {
+    d_proof->startNewProof();
+  }
+}
+
+ProofGenerator* CDCAC::closeProof(const std::vector<Node>& assertions)
+{
+  if (isProofEnabled())
+  {
+    d_proof->endScope(assertions);
+    return d_proof->getProofGenerator();
+  }
+  return nullptr;
+}
+
 bool CDCAC::checkIntegrality(std::size_t cur_variable, const poly::Value& value)
 {
   Node var = d_constraints.varMapper()(d_variableOrdering[cur_variable]);
@@ -470,7 +525,19 @@ bool CDCAC::hasRootBelow(const poly::Polynomial& p,
 
 void CDCAC::pruneRedundantIntervals(std::vector<CACInterval>& intervals)
 {
-  cleanIntervals(intervals);
+  if (isProofEnabled())
+  {
+    std::vector<CACInterval> allIntervals = intervals;
+    cleanIntervals(intervals);
+    d_proof->pruneChildren([&allIntervals, &intervals](std::size_t i) {
+      return std::find(intervals.begin(), intervals.end(), allIntervals[i])
+             != intervals.end();
+    });
+  }
+  else
+  {
+    cleanIntervals(intervals);
+  }
 }
 
 }  // namespace cad
diff --git a/src/theory/arith/nl/cad/cdcac.h b/src/theory/arith/nl/cad/cdcac.h
index f9294fdf1..4511d0c55 100644
--- a/src/theory/arith/nl/cad/cdcac.h
+++ b/src/theory/arith/nl/cad/cdcac.h
@@ -30,6 +30,7 @@
 
 #include "theory/arith/nl/cad/cdcac_utils.h"
 #include "theory/arith/nl/cad/constraints.h"
+#include "theory/arith/nl/cad/proof_generator.h"
 #include "theory/arith/nl/cad/variable_ordering.h"
 #include "theory/arith/nl/nl_model.h"
 
@@ -46,10 +47,10 @@ namespace cad {
 class CDCAC
 {
  public:
-  /** Initialize without a variable ordering. */
-  CDCAC();
   /** Initialize this method with the given variable ordering. */
-  CDCAC(const std::vector<poly::Variable>& ordering);
+  CDCAC(context::Context* ctx,
+        ProofNodeManager* pnm,
+        const std::vector<poly::Variable>& ordering = {});
 
   /** Reset this instance. */
   void reset();
@@ -138,7 +139,20 @@ class CDCAC
   std::vector<CACInterval> getUnsatCover(std::size_t curVariable = 0,
                                          bool returnFirstInterval = false);
 
+  void startNewProof();
+  /**
+   * Finish the generated proof (if proofs are enabled) with a scope over the
+   * given assertions.
+   */
+  ProofGenerator* closeProof(const std::vector<Node>& assertions);
+
+  /** Get the proof generator */
+  CADProofGenerator* getProof() { return d_proof.get(); }
+
  private:
+  /** Check whether proofs are enabled */
+  bool isProofEnabled() const { return d_proof != nullptr; }
+
   /**
    * Check whether the current sample satisfies the integrality condition of the
    * current variable. Returns true if the variable is not integral or the
@@ -187,6 +201,9 @@ class CDCAC
 
   /** The linear assignment used as an initial guess. */
   std::vector<poly::Value> d_initialAssignment;
+
+  /** The proof generator */
+  std::unique_ptr<CADProofGenerator> d_proof;
 };
 
 }  // namespace cad