[MLIR][Presburger] PresburgerSet::subtract: automatically restore state on return

[Groverkss]

LGTM.

[arjunp]

rebase

mlir/lib/Analysis/Presburger/PresburgerSet.cpp

 //===- Set.cpp - MLIR PresburgerSet Class ---------------------------------===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 
 #include "mlir/Analysis/Presburger/PresburgerSet.h"
 #include "mlir/Analysis/Presburger/Simplex.h"
 #include "mlir/Analysis/Presburger/Utils.h"
 #include "llvm/ADT/STLExtras.h"
+#include "llvm/ADT/ScopeExit.h"
 #include "llvm/ADT/SmallBitVector.h"
 
 using namespace mlir;
 using namespace presburger_utils;
 
 PresburgerSet::PresburgerSet(const IntegerPolyhedron &poly)
     : PresburgerSpace(poly) {
   unionPolyInPlace(poly);
static void subtractRecursively(IntegerPolyhedron &b, Simplex &simplex,
   // removing all constraints beyond the original number of inequalities
   // and equalities, so we store these counts first.
   const unsigned bInitNumIneqs = b.getNumInequalities();
   const unsigned bInitNumEqs = b.getNumEqualities();
   const unsigned bInitNumLocals = b.getNumLocalIds();
   // Similarly, we also want to rollback simplex to its original state.
   const unsigned initialSnapshot = simplex.getSnapshot();
 
-  // Automatically restore the original state when we return.
   auto restoreState = [&]() {
     b.removeIdRange(IntegerPolyhedron::IdKind::Local, bInitNumLocals,
                     b.getNumLocalIds());
     b.removeInequalityRange(bInitNumIneqs, b.getNumInequalities());
     b.removeEqualityRange(bInitNumEqs, b.getNumEqualities());
     simplex.rollback(initialSnapshot);
   };
 
+  // Automatically restore the original state when we return.
+  auto stateRestorer = llvm::make_scope_exit(restoreState);
+
   // Find out which inequalities of sI correspond to division inequalities for
   // the local variables of sI.
   std::vector<MaybeLocalRepr> repr(sI.getNumLocalIds());
   sI.getLocalReprs(repr);
 
   // Add sI's locals to b, after b's locals. Also add b's locals to sI, before
   // sI's locals.
   b.mergeLocalIds(sI);
static void subtractRecursively(IntegerPolyhedron &b, Simplex &simplex,
   unsigned offset = simplex.getNumConstraints();
   unsigned numLocalsAdded = b.getNumLocalIds() - bInitNumLocals;
   simplex.appendVariable(numLocalsAdded);
 
   unsigned snapshotBeforeIntersect = simplex.getSnapshot();
   simplex.intersectIntegerPolyhedron(sI);
 
   if (simplex.isEmpty()) {
-    /// b ^ s_i is empty, so b \ s_i = b. We move directly to i + 1.
-    /// We are ignoring level i completely, so we restore the state
-    /// *before* going to level i + 1.
+    // b ^ s_i is empty, so b \ s_i = b. We move directly to i + 1.
+    // We are ignoring level i completely, so we restore the state
+    // *before* going to level i + 1.
     restoreState();
     subtractRecursively(b, simplex, s, i + 1, result);
+
+    // We already restored the state above and the recursive call should have
+    // restored to the same state before returning, so we don't need to restore
+    // the state again.
+    stateRestorer.release();
     return;
   }
 
   simplex.detectRedundant();
 
   // Equalities are added to simplex as a pair of inequalities.
   unsigned totalNewSimplexInequalities =
       2 * sI.getNumEqualities() + sI.getNumInequalities();
for (unsigned j = 0, e = sI.getNumEqualities(); j < e; ++j) {
     // skip it as above to make the result simpler. Divisions are always
     // represented in terms of inequalities and not equalities, so we do not
     // check for division inequalities here.
     if (!isMarkedRedundant[offset + 2 * j])
       processInequality(coeffs);
     if (!isMarkedRedundant[offset + 2 * j + 1])
       processInequality(getNegatedCoeffs(coeffs));
   }
-
-  restoreState();
 }
 
 /// Return the set difference poly \ set.
 ///
 /// The Poly here is modified in subtractRecursively, so it cannot be a const
 /// reference even though it is restored to its original state before returning
 /// from that function.
 PresburgerSet PresburgerSet::getSetDifference(IntegerPolyhedron poly,