[MLIR][Presburger] Introduce SimplexRollbackScopeExit to rollback on scope exit

[arjunp]

Summary

This simplifies many places where we just want to do something in a "transient context"
and return some value.

[Groverkss]

LGTM.

mlir/include/mlir/Analysis/Presburger/Simplex.h

private:
   /// orientation.
   void markRowRedundant(Unknown &u);
 
   /// Reduce the given basis, starting at the specified level, using general
   /// basis reduction.
   void reduceBasis(Matrix &basis, unsigned level);
 };
 
+/// Takes a snapshot of the simplex state on construction and rolls back to the
+/// snapshot on destruction.
+///
+/// Useful for performing operations in a "transient context", all changes from
+/// which get rolled back on scope exit.
+class SimplexRollbackScopeExit {
+public:
+  SimplexRollbackScopeExit(Simplex &simplex) : simplex(simplex) {
+    snapshot = simplex.getSnapshot();
+  };
+  ~SimplexRollbackScopeExit() { simplex.rollback(snapshot); }
+
+private:
+  SimplexBase &simplex;
+  unsigned snapshot;
+};
+
 } // namespace presburger
 } // namespace mlir
 
 #endif // MLIR_ANALYSIS_PRESBURGER_SIMPLEX_H

mlir/lib/Analysis/Presburger/PresburgerRelation.cpp

static void subtractRecursively(IntegerRelation &b, Simplex &simplex,
   simplex.rollback(snapshotBeforeIntersect);
 
   // Recurse with the part b ^ ~ineq. Note that b is modified throughout
   // subtractRecursively. At the time this function is called, the current b is
   // actually equal to b ^ s_i1 ^ s_i2 ^ ... ^ s_ij, and ineq is the next
   // inequality, s_{i,j+1}. This function recurses into the next level i + 1
   // with the part b ^ s_i1 ^ s_i2 ^ ... ^ s_ij ^ ~s_{i,j+1}.
   auto recurseWithInequality = [&, i](ArrayRef<int64_t> ineq) {
-    size_t snapshot = simplex.getSnapshot();
+    SimplexRollbackScopeExit scopeExit(simplex);
     b.addInequality(ineq);
     simplex.addInequality(ineq);
     subtractRecursively(b, simplex, s, i + 1, result);
     b.removeInequality(b.getNumInequalities() - 1);
-    simplex.rollback(snapshot);
   };
 
   // For each inequality ineq, we first recurse with the part where ineq
   // is not satisfied, and then add the ineq to b and simplex because
   // ineq must be satisfied by all later parts.
   auto processInequality = [&](ArrayRef<int64_t> ineq) {
     recurseWithInequality(getComplementIneq(ineq));
     b.addInequality(ineq);
PresburgerRelation SetCoalescer::coalesce() {
 
   return newSet;
 }
 
 /// Given a Simplex `simp` and one of its inequalities `ineq`, check
 /// that all inequalities of `cuttingIneqsB` are redundant for the facet of
 /// `simp` where `ineq` holds as an equality is contained within `a`.
 bool SetCoalescer::isFacetContained(ArrayRef<int64_t> ineq, Simplex &simp) {
-  unsigned snapshot = simp.getSnapshot();
+  SimplexRollbackScopeExit scopeExit(simp);
   simp.addEquality(ineq);
-  if (llvm::any_of(cuttingIneqsB, [&simp](ArrayRef<int64_t> curr) {
-        return !simp.isRedundantInequality(curr);
-      })) {
-    simp.rollback(snapshot);
-    return false;
-  }
-  simp.rollback(snapshot);
-  return true;
+  return llvm::all_of(cuttingIneqsB, [&simp](ArrayRef<int64_t> curr) {
+    return simp.isRedundantInequality(curr);
+  });
 }
 
 void SetCoalescer::addCoalescedDisjunct(unsigned i, unsigned j,
                                         const IntegerRelation &disjunct) {
   assert(i != j && "The indices must refer to different disjuncts");
   unsigned n = disjuncts.size();
   if (j == n - 1) {
     // This case needs special handling since position `n` - 1 is removed

mlir/lib/Analysis/Presburger/Simplex.cpp

 }
 
 /// Compute the optimum of the specified expression in the specified direction,
 /// or None if it is unbounded.
 MaybeOptimum<Fraction> Simplex::computeOptimum(Direction direction,
                                                ArrayRef<int64_t> coeffs) {
   if (empty)
     return OptimumKind::Empty;
-  unsigned snapshot = getSnapshot();
+
+  SimplexRollbackScopeExit scopeExit(*this);
   unsigned conIndex = addRow(coeffs);
   unsigned row = con[conIndex].pos;
   MaybeOptimum<Fraction> optimum = computeRowOptimum(direction, row);
-  rollback(snapshot);
   return optimum;
 }
 
 MaybeOptimum<Fraction> Simplex::computeOptimum(Direction direction,
                                                Unknown &u) {
   if (empty)
     return OptimumKind::Empty;
   if (u.orientation == Orientation::Column) {
Fraction computeWidthAndDuals(ArrayRef<int64_t> dir,
                                 SmallVectorImpl<int64_t> &dual,
                                 int64_t &dualDenom) {
     // We can't just call into computeWidth or computeOptimum since we need to
     // access the state of the tableau after computing the optimum, and these
     // functions rollback the insertion of the objective function into the
     // tableau before returning. We instead add a row for the objective function
     // ourselves, call into computeOptimum, compute the duals from the tableau
     // state, and finally rollback the addition of the row before returning.
-    unsigned snap = simplex.getSnapshot();
+    SimplexRollbackScopeExit scopeExit(simplex);
     unsigned conIndex = simplex.addRow(getCoeffsForDirection(dir));
     unsigned row = simplex.con[conIndex].pos;
     MaybeOptimum<Fraction> maybeWidth =
         simplex.computeRowOptimum(Simplex::Direction::Up, row);
     assert(maybeWidth.isBounded() && "Width should be bounded!");
     dualDenom = simplex.tableau(row, 0);
     dual.clear();
 
for (unsigned i = simplexConstraintOffset; i < conIndex; i += 2) {
              "orientation!");
       if (simplex.con[i].orientation == Orientation::Column)
         dual.push_back(-simplex.tableau(row, simplex.con[i].pos));
       else if (simplex.con[i + 1].orientation == Orientation::Column)
         dual.push_back(simplex.tableau(row, simplex.con[i + 1].pos));
       else
         dual.push_back(0);
     }
-    simplex.rollback(snap);
     return *maybeWidth;
   }
 
   /// Remove the last equality that was added through addEqualityForDirection.
   ///
   /// We do this by rolling back to the snapshot at the top of the stack, which
   /// should be a snapshot taken just before the last equality was added.
   void removeLastEquality() {