move duplicate exponent check to fromMonomials

Author: j2kun

mlir/include/mlir/Dialect/Polynomial/IR/Polynomial.h

@@ -9,8 +9,7 @@
 #ifndef MLIR_DIALECT_POLYNOMIAL_IR_POLYNOMIAL_H_
 #define MLIR_DIALECT_POLYNOMIAL_IR_POLYNOMIAL_H_
 
-#include <utility>
-
+#include "mlir/Support/LogicalResult.h"
 #include "mlir/Support/LLVM.h"
 #include "llvm/ADT/APInt.h"
 #include "llvm/ADT/ArrayRef.h"
@@ -76,7 +75,9 @@ class Polynomial {
 
   explicit Polynomial(ArrayRef<Monomial> terms) : terms(terms) {};
 
-  static Polynomial fromMonomials(ArrayRef<Monomial> monomials);
+  // Returns a Polynomial from a list of monomials.
+  // Fails if two monomials have the same exponent.
+  static FailureOr<Polynomial> fromMonomials(ArrayRef<Monomial> monomials);
 
   /// Returns a polynomial with coefficients given by `coeffs`. The value
   /// coeffs[i] is converted to a monomial with exponent i.

mlir/lib/Dialect/Polynomial/IR/Polynomial.cpp

@@ -8,7 +8,7 @@
 
 #include "mlir/Dialect/Polynomial/IR/Polynomial.h"
 
-#include "mlir/IR/MLIRContext.h"
+#include "mlir/Support/LogicalResult.h"
 #include "llvm/ADT/APInt.h"
 #include "llvm/ADT/SmallString.h"
 #include "llvm/ADT/SmallVector.h"
@@ -18,10 +18,20 @@
 namespace mlir {
 namespace polynomial {
 
-Polynomial Polynomial::fromMonomials(ArrayRef<Monomial> monomials) {
+FailureOr<Polynomial> Polynomial::fromMonomials(ArrayRef<Monomial> monomials) {
   // A polynomial's terms are canonically stored in order of increasing degree.
   auto monomialsCopy = llvm::SmallVector<Monomial>(monomials);
   std::sort(monomialsCopy.begin(), monomialsCopy.end());
+
+  // Ensure non-unique exponents are not present. Since we sorted the list by
+  // exponent, a linear scan of adjancent monomials suffices.
+  if (std::adjacent_find(monomialsCopy.begin(), monomialsCopy.end(),
+                         [](const Monomial &lhs, const Monomial &rhs) {
+                           return lhs.exponent == rhs.exponent;
+                         }) != monomialsCopy.end()) {
+    return failure();
+  }
+
   return Polynomial(monomialsCopy);
 }
 
@@ -32,7 +42,11 @@ Polynomial Polynomial::fromCoefficients(ArrayRef<int64_t> coeffs) {
   for (size_t i = 0; i < size; i++) {
     monomials.emplace_back(coeffs[i], i);
   }
-  return Polynomial::fromMonomials(monomials);
+  auto result = Polynomial::fromMonomials(monomials);
+  // Construction guarantees unique exponents, so the failure mode of
+  // fromMonomials can be bypassed.
+  assert(succeeded(result));
+  return result.value();
 }
 
 void Polynomial::print(raw_ostream &os, ::llvm::StringRef separator,

mlir/lib/Dialect/Polynomial/IR/PolynomialAttributes.cpp

@@ -97,7 +97,6 @@ Attribute PolynomialAttr::parse(AsmParser &parser, Type type) {
 
   llvm::SmallVector<Monomial> monomials;
   llvm::StringSet<> variables;
-  llvm::DenseSet<APInt> exponents;
 
   while (true) {
     Monomial parsedMonomial;
@@ -116,16 +115,6 @@ Attribute PolynomialAttr::parse(AsmParser &parser, Type type) {
     }
     monomials.push_back(parsedMonomial);
 
-    if (llvm::is_contained(exponents, parsedMonomial.exponent)) {
-      llvm::SmallString<16> coeffString;
-      parsedMonomial.exponent.toStringSigned(coeffString);
-      parser.emitError(parser.getCurrentLocation())
-          << "at most one monomial may have exponent " << coeffString
-          << ", but found multiple";
-      return {};
-    }
-    exponents.insert(parsedMonomial.exponent);
-
     if (shouldParseMore)
       continue;
 
@@ -146,8 +135,13 @@ Attribute PolynomialAttr::parse(AsmParser &parser, Type type) {
             vars);
   }
 
-  Polynomial poly = Polynomial::fromMonomials(monomials);
-  return PolynomialAttr::get(parser.getContext(), poly);
+  auto result = Polynomial::fromMonomials(monomials);
+  if (failed(result)) {
+      parser.emitError(parser.getCurrentLocation())
+          << "parsed polynomial must have unique exponents among monomials";
+      return {};
+  }
+  return PolynomialAttr::get(parser.getContext(), result.value());
 }
 
 void RingAttr::print(AsmPrinter &p) const {

mlir/test/Dialect/Polynomial/attributes.mlir

@@ -14,8 +14,8 @@
 
 // -----
 
-// expected-error@below {{at most one monomial may have exponent 2, but found multiple}}
 #my_poly = #polynomial.polynomial<5 + x**2 + 3x**2>
+// expected-error@below {{parsed polynomial must have unique exponents among monomials}}
 #ring1 = #polynomial.ring<coefficientType=i32, coefficientModulus=2837465, polynomialModulus=#my_poly>
 
 // -----