RequirementMachine: Compare weight before length in Term/MutableTerm::compare()

Author: slavapestov

lib/AST/RequirementMachine/Term.cpp

@@ -122,19 +122,42 @@ bool Term::containsNameSymbols() const {
   return false;
 }
 
-/// Shortlex order on symbol ranges.
+/// Weighted shortlex order on symbol ranges, used for implementing
+/// Term::compare() and MutableTerm::compare().
 ///
-/// First we compare length, then perform a lexicographic comparison
-/// on symbols if the two ranges have the same length.
+/// We first compute a weight vector for both terms and compare the
+/// vectors lexicographically:
+/// - Weight of generic param symbols
+/// - Number of name symbols
+/// - Number of element symbols
 ///
-/// This is used to implement Term::compare() and MutableTerm::compare()
-/// below.
-static std::optional<int> shortlexCompare(const Symbol *lhsBegin,
-                                          const Symbol *lhsEnd,
-                                          const Symbol *rhsBegin,
-                                          const Symbol *rhsEnd,
-                                          RewriteContext &ctx) {
-  // First, compare the number of name and pack element symbols.
+/// If the terms have the same weight, we compare length.
+///
+/// If the terms have the same weight and length, we perform a
+/// lexicographic comparison on symbols.
+///
+static std::optional<int> compareImpl(const Symbol *lhsBegin,
+                                      const Symbol *lhsEnd,
+                                      const Symbol *rhsBegin,
+                                      const Symbol *rhsEnd,
+                                      RewriteContext &ctx) {
+  ASSERT(lhsBegin != lhsEnd);
+  ASSERT(rhsBegin != rhsEnd);
+
+  // First compare weights on generic parameters. The implicit
+  // assumption here is we don't form terms with generic parameter
+  // symbols in the middle, which is true. Otherwise, we'd need
+  // to add up their weights like we do below for name symbols,
+  // of course.
+  if (lhsBegin->getKind() == Symbol::Kind::GenericParam &&
+      rhsBegin->getKind() == Symbol::Kind::GenericParam) {
+    unsigned lhsWeight = lhsBegin->getGenericParam()->getWeight();
+    unsigned rhsWeight = rhsBegin->getGenericParam()->getWeight();
+    if (lhsWeight != rhsWeight)
+      return lhsWeight > rhsWeight ? 1 : -1;
+  }
+
+  // Compare the number of name and pack element symbols.
   unsigned lhsNameCount = 0;
   unsigned lhsPackElementCount = 0;
   for (auto *iter = lhsBegin; iter != lhsEnd; ++iter) {
@@ -192,17 +215,17 @@ static std::optional<int> shortlexCompare(const Symbol *lhsBegin,
   return 0;
 }
 
-/// Shortlex order on terms. Returns None if the terms are identical except
+/// Reduction order on terms. Returns None if the terms are identical except
 /// for an incomparable superclass or concrete type symbol at the end.
 std::optional<int> Term::compare(Term other, RewriteContext &ctx) const {
-  return shortlexCompare(begin(), end(), other.begin(), other.end(), ctx);
+  return compareImpl(begin(), end(), other.begin(), other.end(), ctx);
 }
 
-/// Shortlex order on mutable terms. Returns None if the terms are identical
+/// Reduction order on mutable terms. Returns None if the terms are identical
 /// except for an incomparable superclass or concrete type symbol at the end.
 std::optional<int> MutableTerm::compare(const MutableTerm &other,
                                         RewriteContext &ctx) const {
-  return shortlexCompare(begin(), end(), other.begin(), other.end(), ctx);
+  return compareImpl(begin(), end(), other.begin(), other.end(), ctx);
 }
 
 /// Replace the subterm in the range [from,to) of this term with \p rhs.