[MLIR][Presburger] fix vector update in coalesce

[webmiche]

Summary

When addCoalescedPolyhedron was called with j == n - 1,
the polyhedrons-vector was not properly updated (the
IntegerPolyhedron at position n - 2 was "lost"). This patch adds
special handling to that case and a regression testcase.

[Groverkss]

LGTM.

mlir/lib/Analysis/Presburger/PresburgerSet.cpp

412–439

This comment would better fit inside the if condition.

[webmiche]

addressed comment.

[arjunp]

mlir/lib/Analysis/Presburger/PresburgerSet.cpp

425–433

Please add a comment here explaining why it's correct for other possible edge cases with i or j being n - 1 or n - 2.

[webmiche]

address comment.

mlir/lib/Analysis/Presburger/PresburgerSet.cpp

 /// equal.
 static void
 addCoalescedPolyhedron(SmallVectorImpl<IntegerPolyhedron> &polyhedrons,
                        unsigned i, unsigned j, const IntegerPolyhedron &poly,
                        SmallVectorImpl<Simplex> &simplices) {
   assert(i != j && "The indices must refer to different polyhedra");
 
   unsigned n = polyhedrons.size();
+  if (j == n - 1) {
+    // This case needs special handling since position `n` - 1 is removed from
+    // the vector, hence the `IntegerPolyhedron` at position `n` - 2 is lost
+    // otherwise.
+    polyhedrons[i] = polyhedrons[n - 2];
+    polyhedrons.pop_back();
+    polyhedrons[n - 2] = poly;
+
+    simplices[i] = simplices[n - 2];
+    simplices.pop_back();
+    simplices[n - 2] = Simplex(poly);
+
+  } else {
+    // Other possible edge cases are correct since for `j` or `i` == `n` - 2,
+    // the `IntegerPolyhedron` at position `n` - 2 should be lost. The case
+    // `i` == `n` - 1 makes the first following statement a noop. Hence, in this
+    // case the same thing is done as above, but with `j` rather than `i`.
     polyhedrons[i] = polyhedrons[n - 1];
     polyhedrons[j] = polyhedrons[n - 2];
     polyhedrons.pop_back();
     polyhedrons[n - 2] = poly;
 
     simplices[i] = simplices[n - 1];
     simplices[j] = simplices[n - 2];
     simplices.pop_back();
     simplices[n - 2] = Simplex(poly);
   }
+}
 
 /// Given two polyhedra `a` and `b` at positions `i` and `j` in `polyhedrons`
 /// and `redundantIneqsA` being the inequalities of `a` that are redundant for
 /// `b` (similarly for `cuttingIneqsA`, `redundantIneqsB`, and `cuttingIneqsB`),
 /// checks whether the facets of all cutting inequalites of `a` are contained in
 /// `b`. If so, a new polyhedron consisting of all redundant inequalites of `a`
 /// and `b` and all equalities of both is created.
 ///

mlir/unittests/Analysis/Presburger/PresburgerSetTest.cpp

PresburgerSet set = parsePresburgerSetFromPolyStrings(
              "(x) : (x == 0)",
              "(x) : (x - 2 == 0)",
              "(x) : (x + 2 == 0)",
              "(x) : (x - 2 >= 0, -x + 3 >= 0)",
          });
   expectCoalesce(3, set);
 }
 
+TEST(SetTest, coalesceLastCoalesced) {
+  PresburgerSet set = parsePresburgerSetFromPolyStrings(
+      1, {
+             "(x) : (x == 0)",
+             "(x) : (x - 1 >= 0, -x + 3 >= 0)",
+             "(x) : (x + 2 == 0)",
+             "(x) : (x - 2 >= 0, -x + 4 >= 0)",
+         });
+  expectCoalesce(3, set);
+}
+
 TEST(SetTest, coalesceDiv) {
   PresburgerSet set =
       parsePresburgerSetFromPolyStrings(1, {
                                                "(x) : (x floordiv 2 == 0)",
                                                "(x) : (x floordiv 2 - 1 == 0)",
                                            });
   expectCoalesce(2, set);
 }