[MLIR][Presburger] IntegerRelation::truncate: fix bug when truncating equalities

[arjunp]

Summary

This was truncating inequalities instead of equalities.

[Groverkss]

LGTM.

[arjunp]

rebase

mlir/lib/Analysis/Presburger/IntegerRelation.cpp

 }
 
 void IntegerRelation::truncate(const CountsSnapshot &counts) {
   truncateIdKind(IdKind::Domain, counts);
   truncateIdKind(IdKind::Range, counts);
   truncateIdKind(IdKind::Symbol, counts);
   truncateIdKind(IdKind::Local, counts);
   removeInequalityRange(counts.getNumIneqs(), getNumInequalities());
-  removeInequalityRange(counts.getNumEqs(), getNumEqualities());
+  removeEqualityRange(counts.getNumEqs(), getNumEqualities());
 }
 
 unsigned IntegerRelation::insertId(IdKind kind, unsigned pos, unsigned num) {
   assert(pos <= getNumIdKind(kind));
 
   unsigned insertPos = PresburgerSpace::insertId(kind, pos, num);
   inequalities.insertColumns(insertPos, num);
   equalities.insertColumns(insertPos, num);

mlir/unittests/Analysis/Presburger/IntegerPolyhedronTest.cpp

IntegerPolyhedron poly2 = parsePoly(
       "(x) : (x - 2*(x floordiv 2) == 0, x - 4*(x floordiv 4) - 2 == 0)");
   EXPECT_TRUE(poly2.containsPointNoLocal({6}));
   EXPECT_FALSE(poly2.containsPointNoLocal({4}));
 
   IntegerPolyhedron poly3 = parsePoly("(x, y) : (2*x - y >= 0, y - 3*x >= 0)");
   EXPECT_TRUE(poly3.containsPointNoLocal({0, 0}));
   EXPECT_FALSE(poly3.containsPointNoLocal({1, 0}));
 }
+
+TEST(IntegerPolyhedronTest, truncateEqualityRegressionTest) {
+  // IntegerRelation::truncate was truncating inequalities to the number of
+  // equalities.
+  IntegerRelation set(1);
+  IntegerRelation::CountsSnapshot snapshot = set.getCounts();
+  set.addEquality({1, 0});
+  set.truncate(snapshot);
+  EXPECT_EQ(set.getNumEqualities(), 0u);
+}