add iterVarKind for convenient iterating over variables

[j2kun]

I find myself doing this alot

for (unsigned varIndex = rel.getVarKindOffset(VarKind::Domain);
     varIndex < rel.getVarKindEnd(VarKind::Domain); ++varIndex) {
  ...
}

Adding this convenience method so I can instead do

for (unsigned varIndex : rel.iterVarKind(VarKind::Domain)) {
  ...
}

[llvmbot]

@llvm/pr-subscribers-mlir

@llvm/pr-subscribers-mlir-presburger

Author: Jeremy Kun (j2kun)

Changes

I find myself doing this alot

for (unsigned varIndex = rel.getVarKindOffset(VarKind::Domain);
     varIndex < rel.getVarKindEnd(VarKind::Domain); ++varIndex) {
  ...
}

Adding this convenience method so I can instead do

for (unsigned varIndex : rel.iterVarKind(VarKind::Domain)) {
  ...
}

---

Full diff: https://github.com/llvm/llvm-project/pull/152091.diff

2 Files Affected:

• (modified) mlir/include/mlir/Analysis/Presburger/IntegerRelation.h (+7)
• (modified) mlir/unittests/Analysis/Presburger/IntegerRelationTest.cpp (+11)

diff --git a/mlir/include/mlir/Analysis/Presburger/IntegerRelation.h b/mlir/include/mlir/Analysis/Presburger/IntegerRelation.h
index ee401cca8f552..fb895b1a6fb78 100644
--- a/mlir/include/mlir/Analysis/Presburger/IntegerRelation.h
+++ b/mlir/include/mlir/Analysis/Presburger/IntegerRelation.h
@@ -20,6 +20,7 @@
 #include "mlir/Analysis/Presburger/PresburgerSpace.h"
 #include "mlir/Analysis/Presburger/Utils.h"
 #include "llvm/ADT/DynamicAPInt.h"
+#include "llvm/ADT/Sequence.h"
 #include "llvm/ADT/SmallVector.h"
 #include "llvm/Support/LogicalResult.h"
 #include <optional>
@@ -268,6 +269,12 @@ class IntegerRelation {
     return space.getVarKindEnd(kind);
   }
 
+  /// Return an interator over the variables of the specified kind
+  /// starting at the relevant offset.
+  auto iterVarKind(VarKind kind) {
+    return llvm::seq(getVarKindOffset(kind), getVarKindEnd(kind));
+  }
+
   /// Get the number of elements of the specified kind in the range
   /// [varStart, varLimit).
   unsigned getVarKindOverlap(VarKind kind, unsigned varStart,
diff --git a/mlir/unittests/Analysis/Presburger/IntegerRelationTest.cpp b/mlir/unittests/Analysis/Presburger/IntegerRelationTest.cpp
index dd0b09f7f05d2..818dcfd578e48 100644
--- a/mlir/unittests/Analysis/Presburger/IntegerRelationTest.cpp
+++ b/mlir/unittests/Analysis/Presburger/IntegerRelationTest.cpp
@@ -16,6 +16,7 @@
 
 using namespace mlir;
 using namespace presburger;
+using ::testing::ElementsAre;
 
 TEST(IntegerRelationTest, getDomainAndRangeSet) {
   IntegerRelation rel = parseRelationFromSet(
@@ -702,3 +703,13 @@ TEST(IntegerRelationTest, rangeProductSymbols) {
 
   EXPECT_TRUE(expected.isEqual(rangeProd));
 }
+
+TEST(IntegerRelationTest, getVarKindRange) {
+  IntegerRelation r1 = parseRelationFromSet(
+      "(i1, i2, i3, i4, i5) : (i1 >= 0, i2 >= 0, i3 >= 0, i4 >= 0, i5 >= 0)", 2);
+  SmallVector<unsigned> actual;
+  for (unsigned var : r1.iterVarKind(VarKind::Range)) {
+    actual.push_back(var);
+  }
+  EXPECT_THAT(actual, ElementsAre(2, 3, 4));
+}

[github-actions]

⚠️ We detected that you are using a GitHub private e-mail address to contribute to the repo.
Please turn off Keep my email addresses private setting in your account.
See LLVM Developer Policy and LLVM Discourse for more information.

[github-actions]

⚠️ C/C++ code formatter, clang-format found issues in your code. ⚠️

You can test this locally with the following command:

git-clang-format --diff HEAD~1 HEAD --extensions cpp,h -- mlir/include/mlir/Analysis/Presburger/IntegerRelation.h mlir/unittests/Analysis/Presburger/IntegerRelationTest.cpp

View the diff from clang-format here.

diff --git a/mlir/unittests/Analysis/Presburger/IntegerRelationTest.cpp b/mlir/unittests/Analysis/Presburger/IntegerRelationTest.cpp
index 818dcfd57..a6ed5c5b2 100644
--- a/mlir/unittests/Analysis/Presburger/IntegerRelationTest.cpp
+++ b/mlir/unittests/Analysis/Presburger/IntegerRelationTest.cpp
@@ -706,7 +706,8 @@ TEST(IntegerRelationTest, rangeProductSymbols) {
 
 TEST(IntegerRelationTest, getVarKindRange) {
   IntegerRelation r1 = parseRelationFromSet(
-      "(i1, i2, i3, i4, i5) : (i1 >= 0, i2 >= 0, i3 >= 0, i4 >= 0, i5 >= 0)", 2);
+      "(i1, i2, i3, i4, i5) : (i1 >= 0, i2 >= 0, i3 >= 0, i4 >= 0, i5 >= 0)",
+      2);
   SmallVector<unsigned> actual;
   for (unsigned var : r1.iterVarKind(VarKind::Range)) {
     actual.push_back(var);