[Mlir-commits] [mlir] 280d41d - [MLIR][Presburger] PresburgerSet::subtract: automatically restore state on return

[Arjun P]

Author: Arjun P
Date: 2022-02-23T19:20:44Z
New Revision: 280d41d9bbf00a3978d943bc3e56381eed1544ba

URL: https://github.com/llvm/llvm-project/commit/280d41d9bbf00a3978d943bc3e56381eed1544ba
DIFF: https://github.com/llvm/llvm-project/commit/280d41d9bbf00a3978d943bc3e56381eed1544ba.diff

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

Reviewed By: Groverkss

Differential Revision: https://reviews.llvm.org/D120339

Added: 
    

Modified: 
    mlir/lib/Analysis/Presburger/PresburgerSet.cpp

Removed: 
    


################################################################################
diff  --git a/mlir/lib/Analysis/Presburger/PresburgerSet.cpp b/mlir/lib/Analysis/Presburger/PresburgerSet.cpp
index b75cd1d51fddd..82946297c7e09 100644
--- a/mlir/lib/Analysis/Presburger/PresburgerSet.cpp
+++ b/mlir/lib/Analysis/Presburger/PresburgerSet.cpp
@@ -10,6 +10,7 @@
 #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;
@@ -198,7 +199,6 @@ static void subtractRecursively(IntegerPolyhedron &b, Simplex &simplex,
   // 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());
@@ -207,6 +207,9 @@ static void subtractRecursively(IntegerPolyhedron &b, Simplex &simplex,
     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());
@@ -243,11 +246,16 @@ static void subtractRecursively(IntegerPolyhedron &b, Simplex &simplex,
   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;
   }
 
@@ -309,8 +317,6 @@ static void subtractRecursively(IntegerPolyhedron &b, Simplex &simplex,
     if (!isMarkedRedundant[offset + 2 * j + 1])
       processInequality(getNegatedCoeffs(coeffs));
   }
-
-  restoreState();
 }
 
 /// Return the set 
diff erence poly \ set.