[MLIR][Presburger] add iterVarKind for convenient iterating over variables (#152091)

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)) {
  ...
}
```

---------

Co-authored-by: Jeremy Kun <j2kun@users.noreply.github.com>
This commit is contained in:
Jeremy Kun 2025-08-05 13:18:59 -07:00 committed by GitHub
parent d478502a42
commit 97dee32445
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
2 changed files with 20 additions and 0 deletions

View File

@ -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,13 @@ public:
return space.getVarKindEnd(kind);
}
/// Return an interator over the variables of the specified kind
/// starting at the relevant offset. The return type is auto in
/// keeping with the convention for iterators.
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,

View File

@ -16,6 +16,7 @@
using namespace mlir;
using namespace presburger;
using ::testing::ElementsAre;
TEST(IntegerRelationTest, getDomainAndRangeSet) {
IntegerRelation rel = parseRelationFromSet(
@ -702,3 +703,14 @@ 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));
}