1 files changed, 57 insertions, 2 deletions

diff --git a/src/theory/arith/theory_arith.cpp b/src/theory/arith/theory_arith.cpp
index 1436198a8..4884d8484 100644
--- a/src/theory/arith/theory_arith.cpp
+++ b/src/theory/arith/theory_arith.cpp
@@ -21,6 +21,7 @@
 #include "smt/smt_statistics_registry.h"
 #include "theory/arith/arith_rewriter.h"
 #include "theory/arith/infer_bounds.h"
+#include "theory/arith/nl/nonlinear_extension.h"
 #include "theory/arith/theory_arith_private.h"
 #include "theory/ext_theory.h"
 
@@ -42,7 +43,8 @@ TheoryArith::TheoryArith(context::Context* c,
           new TheoryArithPrivate(*this, c, u, out, valuation, logicInfo, pnm)),
       d_ppRewriteTimer("theory::arith::ppRewriteTimer"),
       d_astate(*d_internal, c, u, valuation),
-      d_inferenceManager(*this, d_astate, pnm)
+      d_inferenceManager(*this, d_astate, pnm),
+      d_nonlinearExtension(nullptr)
 {
   smtStatisticsRegistry()->registerStat(&d_ppRewriteTimer);
 
@@ -76,6 +78,13 @@ void TheoryArith::finishInit()
     d_valuation.setUnevaluatedKind(kind::SINE);
     d_valuation.setUnevaluatedKind(kind::PI);
   }
+  // only need to create nonlinear extension if non-linear logic
+  const LogicInfo& logicInfo = getLogicInfo();
+  if (logicInfo.isTheoryEnabled(THEORY_ARITH) && !logicInfo.isLinear())
+  {
+    d_nonlinearExtension.reset(
+        new nl::NonlinearExtension(*this, d_equalityEngine));
+  }
   // finish initialize internally
   d_internal->finishInit();
 }
@@ -123,7 +132,53 @@ void TheoryArith::propagate(Effort e) {
 }
 bool TheoryArith::collectModelInfo(TheoryModel* m)
 {
-  return d_internal->collectModelInfo(m);
+  std::set<Node> termSet;
+  // Work out which variables are needed
+  const std::set<Kind>& irrKinds = m->getIrrelevantKinds();
+  computeAssertedTerms(termSet, irrKinds);
+  // this overrides behavior to not assert equality engine
+  return collectModelValues(m, termSet);
+}
+
+bool TheoryArith::collectModelValues(TheoryModel* m,
+                                     const std::set<Node>& termSet)
+{
+  // get the model from the linear solver
+  std::map<Node, Node> arithModel;
+  d_internal->collectModelValues(termSet, arithModel);
+  // if non-linear is enabled, intercept the model, which may repair its values
+  if (d_nonlinearExtension != nullptr)
+  {
+    // Non-linear may repair values to satisfy non-linear constraints (see
+    // documentation for NonlinearExtension::interceptModel).
+    d_nonlinearExtension->interceptModel(arithModel);
+  }
+  // We are now ready to assert the model.
+  for (const std::pair<const Node, Node>& p : arithModel)
+  {
+    // maps to constant of comparable type
+    Assert(p.first.getType().isComparableTo(p.second.getType()));
+    Assert(p.second.isConst());
+    if (m->assertEquality(p.first, p.second, true))
+    {
+      continue;
+    }
+    // If we failed to assert an equality, it is likely due to theory
+    // combination, namely the repaired model for non-linear changed
+    // an equality status that was agreed upon by both (linear) arithmetic
+    // and another theory. In this case, we must add a lemma, or otherwise
+    // we would terminate with an invalid model. Thus, we add a splitting
+    // lemma of the form ( x = v V x != v ) where v is the model value
+    // assigned by the non-linear solver to x.
+    if (d_nonlinearExtension != nullptr)
+    {
+      Node eq = p.first.eqNode(p.second);
+      Node lem = NodeManager::currentNM()->mkNode(kind::OR, eq, eq.negate());
+      d_out->lemma(lem);
+    }
+    return false;
+  }
+  return true;
 }
 
 void TheoryArith::notifyRestart(){