From fc620343734a9a2415792182a91f7fd273d9c9c1 Mon Sep 17 00:00:00 2001
From: Gereon Kremer <gkremer@stanford.edu>
Date: Mon, 31 Jan 2022 18:19:58 -0800
Subject: [PATCH] Consider RANs in variable ordering (#7964)

The variable ordering in the arithmetic normal form should also consider real algebraic numbers. This PR also adds new comparators (essentially doing the same) in the rewriter that will replace the normal form utility classes.
---
 src/theory/arith/arith_rewriter.cpp | 82 +++++++++++++++++++++++++++++
 src/theory/arith/normal_form.h      | 10 +++-
 2 files changed, 91 insertions(+), 1 deletion(-)

diff --git a/src/theory/arith/arith_rewriter.cpp b/src/theory/arith/arith_rewriter.cpp
index e157484af..74d0b0d22 100644
--- a/src/theory/arith/arith_rewriter.cpp
+++ b/src/theory/arith/arith_rewriter.cpp
@@ -44,6 +44,88 @@ namespace arith {
 
 namespace {
 
+/**
+ * Implements an ordering on arithmetic leaf nodes, excluding rationals. As this
+ * comparator is meant to be used on children of Kind::NONLINEAR_MULT, we expect
+ * rationals to be handled separately. Furthermore, we expect there to be only a
+ * single real algebraic number.
+ * It broadly categorizes leaf nodes into real algebraic numbers, integers,
+ * variables, and the rest. The ordering is built as follows:
+ * - real algebraic numbers come first
+ * - real terms come before integer terms
+ * - variables come before non-variable terms
+ * - finally, fall back to node ordering
+ */
+struct LeafNodeComparator
+{
+  /** Implements operator<(a, b) as described above */
+  bool operator()(TNode a, TNode b)
+  {
+    if (a == b) return false;
+
+    bool aIsRAN = a.getKind() == Kind::REAL_ALGEBRAIC_NUMBER;
+    bool bIsRAN = b.getKind() == Kind::REAL_ALGEBRAIC_NUMBER;
+    if (aIsRAN != bIsRAN) return aIsRAN;
+    Assert(!aIsRAN && !bIsRAN) << "real algebraic numbers should be combined";
+
+    bool aIsInt = a.getType().isInteger();
+    bool bIsInt = b.getType().isInteger();
+    if (aIsInt != bIsInt) return !aIsInt;
+
+    bool aIsVar = a.isVar();
+    bool bIsVar = b.isVar();
+    if (aIsVar != bIsVar) return aIsVar;
+
+    return a < b;
+  }
+};
+
+/**
+ * Implements an ordering on arithmetic nonlinear multiplications. As we assume
+ * rationals to be handled separately, we only consider Kind::NONLINEAR_MULT as
+ * multiplication terms. For individual factors of the product, we rely on the
+ * ordering from LeafNodeComparator. Furthermore, we expect products to be
+ * sorted according to LeafNodeComparator. The ordering is built as follows:
+ * - single factors come first (everything that is not NONLINEAR_MULT)
+ * - multiplications with less factors come first
+ * - multiplications are compared lexicographically
+ */
+struct ProductNodeComparator
+{
+  /** Implements operator<(a, b) as described above */
+  bool operator()(TNode a, TNode b)
+  {
+    if (a == b) return false;
+
+    Assert(a.getKind() != Kind::MULT);
+    Assert(b.getKind() != Kind::MULT);
+
+    bool aIsMult = a.getKind() == Kind::NONLINEAR_MULT;
+    bool bIsMult = b.getKind() == Kind::NONLINEAR_MULT;
+    if (aIsMult != bIsMult) return !aIsMult;
+
+    if (!aIsMult)
+    {
+      return LeafNodeComparator()(a, b);
+    }
+
+    size_t aLen = a.getNumChildren();
+    size_t bLen = b.getNumChildren();
+    if (aLen != bLen) return aLen < bLen;
+
+    for (size_t i = 0; i < aLen; ++i)
+    {
+      if (a[i] != b[i])
+      {
+        return LeafNodeComparator()(a[i], b[i]);
+      }
+    }
+    Unreachable() << "Nodes are different, but have the same content";
+    return false;
+  }
+};
+
+
 template <typename L, typename R>
 bool evaluateRelation(Kind rel, const L& l, const R& r)
 {
diff --git a/src/theory/arith/normal_form.h b/src/theory/arith/normal_form.h
index 577bd052d..b96c51a9f 100644
--- a/src/theory/arith/normal_form.h
+++ b/src/theory/arith/normal_form.h
@@ -292,7 +292,15 @@ public:
     static inline int cmp(const Node& n, const Node& m) {
       if ( n == m ) { return 0; }
 
-      // this is now slightly off of the old variable order.
+      // RAN < real var < int var < non-variable
+
+      bool nIsRAN = n.getKind() == Kind::REAL_ALGEBRAIC_NUMBER;
+      bool mIsRAN = m.getKind() == Kind::REAL_ALGEBRAIC_NUMBER;
+
+      if (mIsRAN != nIsRAN)
+      {
+        return nIsRAN ? -1 : 1;
+      }
 
       bool nIsInteger = n.getType().isInteger();
       bool mIsInteger = m.getType().isInteger();
-- 
2.30.2