From e55e6b3bd4384f0b80e04d7acc6b6d613a244dd4 Mon Sep 17 00:00:00 2001
From: Andrew Reynolds <andrew.j.reynolds@gmail.com>
Date: Wed, 12 Jan 2022 11:52:07 -0600
Subject: [PATCH] Always use partial function for sqrt (#7926)

Fixes the behavior reported in #7924.
---
 src/theory/arith/operator_elim.cpp            | 20 ++++++++++++-------
 src/theory/arith/operator_elim.h              |  3 +++
 test/regress/CMakeLists.txt                   |  1 +
 .../regress1/nl/issue7924-sqrt-partial.smt2   |  5 +++++
 4 files changed, 22 insertions(+), 7 deletions(-)
 create mode 100644 test/regress/regress1/nl/issue7924-sqrt-partial.smt2

diff --git a/src/theory/arith/operator_elim.cpp b/src/theory/arith/operator_elim.cpp
index 05a83c81c..f346d7596 100644
--- a/src/theory/arith/operator_elim.cpp
+++ b/src/theory/arith/operator_elim.cpp
@@ -432,16 +432,16 @@ Node OperatorElim::getArithSkolem(SkolemFunId id)
     }
     Node skolem;
     SkolemManager* sm = nm->getSkolemManager();
-    if (options().arith.arithNoPartialFun)
+    if (usePartialFunction(id))
     {
-      // partial function: division, where we treat the skolem function as
-      // a constant
-      skolem = sm->mkSkolemFunction(id, tn);
+      // partial function: division
+      skolem = sm->mkSkolemFunction(id, nm->mkFunctionType(tn, tn));
     }
     else
     {
-      // partial function: division
-      skolem = sm->mkSkolemFunction(id, nm->mkFunctionType(tn, tn));
+      // partial function: division, where we treat the skolem function as
+      // a constant
+      skolem = sm->mkSkolemFunction(id, tn);
     }
     // cache it
     d_arithSkolem[id] = skolem;
@@ -453,13 +453,19 @@ Node OperatorElim::getArithSkolem(SkolemFunId id)
 Node OperatorElim::getArithSkolemApp(Node n, SkolemFunId id)
 {
   Node skolem = getArithSkolem(id);
-  if (!options().arith.arithNoPartialFun)
+  if (usePartialFunction(id))
   {
     skolem = NodeManager::currentNM()->mkNode(APPLY_UF, skolem, n);
   }
   return skolem;
 }
 
+bool OperatorElim::usePartialFunction(SkolemFunId id) const
+{
+  // always use partial function for sqrt
+  return !options().arith.arithNoPartialFun || id == SkolemFunId::SQRT;
+}
+
 Node OperatorElim::mkWitnessTerm(Node v,
                                  Node pred,
                                  const std::string& prefix,
diff --git a/src/theory/arith/operator_elim.h b/src/theory/arith/operator_elim.h
index e2c0ab297..10329f93c 100644
--- a/src/theory/arith/operator_elim.h
+++ b/src/theory/arith/operator_elim.h
@@ -128,6 +128,9 @@ class OperatorElim : protected EnvObj, public EagerProofGenerator
    * if the logic is linear.
    */
   void checkNonLinearLogic(Node term);
+
+  /** Whether we should use a partial function for the given id */
+  bool usePartialFunction(SkolemFunId id) const;
 };
 
 }  // namespace arith
diff --git a/test/regress/CMakeLists.txt b/test/regress/CMakeLists.txt
index 96c34638f..8cd30b38d 100644
--- a/test/regress/CMakeLists.txt
+++ b/test/regress/CMakeLists.txt
@@ -1845,6 +1845,7 @@ set(regress_1_tests
   regress1/nl/issue5660-mb-success.smt2
   regress1/nl/issue5662-nl-tc.smt2
   regress1/nl/issue5662-nl-tc-min.smt2
+  regress1/nl/issue7924-sqrt-partial.smt2
   regress1/nl/metitarski-1025.smt2
   regress1/nl/metitarski-3-4.smt2
   regress1/nl/metitarski_3_4_2e.smt2
diff --git a/test/regress/regress1/nl/issue7924-sqrt-partial.smt2 b/test/regress/regress1/nl/issue7924-sqrt-partial.smt2
new file mode 100644
index 000000000..e142d50fd
--- /dev/null
+++ b/test/regress/regress1/nl/issue7924-sqrt-partial.smt2
@@ -0,0 +1,5 @@
+; COMMAND-LINE: --arith-no-partial-fun -q
+; EXPECT: sat
+(set-logic ALL)
+(assert (exists ((V Real)) (distinct (sqrt 1.0) (sqrt 0.0))))
+(check-sat)
-- 
2.30.2