Fix reduction of `sqrt` (#3478)

5 files changed, 101 insertions, 25 deletions

diff --git a/src/theory/arith/theory_arith_private.cpp b/src/theory/arith/theory_arith_private.cpp
index a6a47c73c..bc4d7db02 100644
--- a/src/theory/arith/theory_arith_private.cpp
+++ b/src/theory/arith/theory_arith_private.cpp
@@ -4973,7 +4973,7 @@ Node TheoryArithPrivate::expandDefinition(LogicRequest &logicRequest, Node node)
       if (!den.isConst() || den.getConst<Rational>().sgn() == 0)
       {
         Node divByZeroNum =
-            getArithSkolemApp(logicRequest, num, arith_skolem_div_by_zero);
+            getArithSkolemApp(logicRequest, num, ArithSkolemId::DIV_BY_ZERO);
         Node denEq0 = nm->mkNode(kind::EQUAL, den, nm->mkConst(Rational(0)));
         ret = nm->mkNode(kind::ITE, denEq0, divByZeroNum, ret);
       }
@@ -4988,8 +4988,8 @@ Node TheoryArithPrivate::expandDefinition(LogicRequest &logicRequest, Node node)
       Node ret = nm->mkNode(kind::INTS_DIVISION_TOTAL, num, den);
       if (!den.isConst() || den.getConst<Rational>().sgn() == 0)
       {
-        Node intDivByZeroNum =
-            getArithSkolemApp(logicRequest, num, arith_skolem_int_div_by_zero);
+        Node intDivByZeroNum = getArithSkolemApp(
+            logicRequest, num, ArithSkolemId::INT_DIV_BY_ZERO);
         Node denEq0 = nm->mkNode(kind::EQUAL, den, nm->mkConst(Rational(0)));
         ret = nm->mkNode(kind::ITE, denEq0, intDivByZeroNum, ret);
       }
@@ -5005,7 +5005,7 @@ Node TheoryArithPrivate::expandDefinition(LogicRequest &logicRequest, Node node)
       if (!den.isConst() || den.getConst<Rational>().sgn() == 0)
       {
         Node modZeroNum =
-            getArithSkolemApp(logicRequest, num, arith_skolem_mod_by_zero);
+            getArithSkolemApp(logicRequest, num, ArithSkolemId::MOD_BY_ZERO);
         Node denEq0 = nm->mkNode(kind::EQUAL, den, nm->mkConst(Rational(0)));
         ret = nm->mkNode(kind::ITE, denEq0, modZeroNum, ret);
       }
@@ -5037,7 +5037,25 @@ Node TheoryArithPrivate::expandDefinition(LogicRequest &logicRequest, Node node)
         Node lem;
         if (k == kind::SQRT)
         {
-          lem = nm->mkNode(kind::MULT, var, var).eqNode(node[0]);
+          Node skolemApp =
+              getArithSkolemApp(logicRequest, node[0], ArithSkolemId::SQRT);
+          Node uf = skolemApp.eqNode(var);
+          Node nonNeg = nm->mkNode(
+              kind::AND, nm->mkNode(kind::MULT, var, var).eqNode(node[0]), uf);
+
+          // sqrt(x) reduces to:
+          // choice y. ite(x >= 0.0, y * y = x ^ Uf(x), Uf(x))
+          //
+          // Uf(x) makes sure that the reduction still behaves like a function,
+          // otherwise the reduction of (x = 1) ^ (sqrt(x) != sqrt(1)) would be
+          // satisfiable. On the original formula, this would require that we
+          // simultaneously interpret sqrt(1) as 1 and -1, which is not a valid
+          // model.
+          lem = nm->mkNode(
+              kind::ITE,
+              nm->mkNode(kind::GEQ, node[0], nm->mkConst(Rational(0))),
+              nonNeg,
+              uf);
         }
         else
         {
@@ -5102,23 +5120,29 @@ Node TheoryArithPrivate::getArithSkolem(LogicRequest& logicRequest,
     TypeNode tn;
     std::string name;
     std::string desc;
-    if (asi == arith_skolem_div_by_zero)
+    switch (asi)
     {
-      tn = nm->realType();
-      name = std::string("divByZero");
-      desc = std::string("partial real division");
-    }
-    else if (asi == arith_skolem_int_div_by_zero)
-    {
-      tn = nm->integerType();
-      name = std::string("intDivByZero");
-      desc = std::string("partial int division");
-    }
-    else if (asi == arith_skolem_mod_by_zero)
-    {
-      tn = nm->integerType();
-      name = std::string("modZero");
-      desc = std::string("partial modulus");
+      case ArithSkolemId::DIV_BY_ZERO:
+        tn = nm->realType();
+        name = std::string("divByZero");
+        desc = std::string("partial real division");
+        break;
+      case ArithSkolemId::INT_DIV_BY_ZERO:
+        tn = nm->integerType();
+        name = std::string("intDivByZero");
+        desc = std::string("partial int division");
+        break;
+      case ArithSkolemId::MOD_BY_ZERO:
+        tn = nm->integerType();
+        name = std::string("modZero");
+        desc = std::string("partial modulus");
+        break;
+      case ArithSkolemId::SQRT:
+        tn = nm->realType();
+        name = std::string("sqrtUf");
+        desc = std::string("partial sqrt");
+        break;
+      default: Unhandled();
     }
 
     Node skolem;
diff --git a/src/theory/arith/theory_arith_private.h b/src/theory/arith/theory_arith_private.h
index 03cb81785..3f46ddd59 100644
--- a/src/theory/arith/theory_arith_private.h
+++ b/src/theory/arith/theory_arith_private.h
@@ -828,11 +828,16 @@ private:
 
   Statistics d_statistics;
 
-  enum ArithSkolemId
+  enum class ArithSkolemId
   {
-    arith_skolem_div_by_zero,
-    arith_skolem_int_div_by_zero,
-    arith_skolem_mod_by_zero,
+    /* an uninterpreted function f s.t. f(x) = x / 0.0 (real division) */
+    DIV_BY_ZERO,
+    /* an uninterpreted function f s.t. f(x) = x / 0 (integer division) */
+    INT_DIV_BY_ZERO,
+    /* an uninterpreted function f s.t. f(x) = x mod 0 */
+    MOD_BY_ZERO,
+    /* an uninterpreted function f s.t. f(x) = sqrt(x) */
+    SQRT,
   };
 
   /**
diff --git a/test/regress/CMakeLists.txt b/test/regress/CMakeLists.txt
index 0312c13b7..3c060a304 100644
--- a/test/regress/CMakeLists.txt
+++ b/test/regress/CMakeLists.txt
@@ -537,6 +537,7 @@ set(regress_0_tests
   regress0/nl/coeff-sat.smt2
   regress0/nl/ext-rew-aggr-test.smt2
   regress0/nl/issue3407.smt2
+  regress0/nl/issue3475.smt2
   regress0/nl/magnitude-wrong-1020-m.smt2
   regress0/nl/mult-po.smt2
   regress0/nl/nia-wrong-tl.smt2
@@ -552,6 +553,7 @@ set(regress_0_tests
   regress0/nl/nta/tan-rewrite.smt2
   regress0/nl/real-as-int.smt2
   regress0/nl/real-div-ufnra.smt2
+  regress0/nl/sqrt.smt2
   regress0/nl/sqrt2-value.smt2
   regress0/nl/subs0-unsat-confirm.smt2
   regress0/nl/very-easy-sat.smt2
diff --git a/test/regress/regress0/nl/issue3475.smt2 b/test/regress/regress0/nl/issue3475.smt2
new file mode 100644
index 000000000..128d08a18
--- /dev/null
+++ b/test/regress/regress0/nl/issue3475.smt2
@@ -0,0 +1,6 @@
+(set-logic ALL)
+(declare-fun x () Real)
+(assert (< x 0))
+(assert (not (= (/ (sqrt x) (sqrt x)) x)))
+(set-info :status sat)
+(check-sat)
diff --git a/test/regress/regress0/nl/sqrt.smt2 b/test/regress/regress0/nl/sqrt.smt2
new file mode 100644
index 000000000..fdcec3d62
--- /dev/null
+++ b/test/regress/regress0/nl/sqrt.smt2
@@ -0,0 +1,39 @@
+; EXPECT: sat
+; EXPECT: sat
+; EXPECT: unsat
+; EXPECT: sat
+; EXPECT: unsat
+(set-option :incremental true)
+(set-logic ALL)
+(declare-fun x () Real)
+(declare-fun y () Real)
+(declare-fun z () Real)
+
+(push)
+(assert (= (sqrt 1.0) 1.0))
+(check-sat)
+(pop)
+
+(push)
+(assert (= (sqrt 1.0) (- 1.0)))
+(check-sat)
+(pop)
+
+(push)
+(assert (= x 1.0))
+(assert (not (= (sqrt 1.0) (sqrt x))))
+(check-sat)
+(pop)
+
+(push)
+(assert (< x 0))
+(assert (= (sqrt 1.0) (sqrt x)))
+(check-sat)
+(pop)
+
+(push)
+(assert (= (sqrt y) z))
+(assert (= (sqrt x) (sqrt y)))
+(assert (not (= (sqrt x) z)))
+(check-sat)
+(pop)