[mlir-c] Add Region iterators matching Block & Operation ones

[jpienaar]

Summary

Enables using the same iterator interface to these even though underlying storage is different.

mlir/include/mlir-c/IR.h

 
 /// Takes a block owned by the caller and inserts it before the (non-owned)
 /// reference block in the given region. The reference block must belong to the
 /// region. If the reference block is null, appends the block to the region.
 MLIR_CAPI_EXPORTED void mlirRegionInsertOwnedBlockBefore(MlirRegion region,
                                                          MlirBlock reference,
                                                          MlirBlock block);
 
+/// Returns first region attached to the operation.
+MLIR_CAPI_EXPORTED MlirRegion mlirOperationGetFirstRegion(MlirOperation op);
+
+/// Returns the region immediately following the given region in its parent
+/// operation.
+MLIR_CAPI_EXPORTED MlirRegion mlirRegionGetNextInOperation(MlirRegion region);
+
 //===----------------------------------------------------------------------===//
 // Block API.
 //===----------------------------------------------------------------------===//
 
 /// Creates a new empty block with the given argument types and transfers
 /// ownership to the caller.
 MLIR_CAPI_EXPORTED MlirBlock mlirBlockCreate(intptr_t nArgs,
                                              MlirType const *args);

mlir/lib/CAPI/IR/IR.cpp

 intptr_t mlirOperationGetNumRegions(MlirOperation op) {
   return static_cast<intptr_t>(unwrap(op)->getNumRegions());
 }
 
 MlirRegion mlirOperationGetRegion(MlirOperation op, intptr_t pos) {
   return wrap(&unwrap(op)->getRegion(static_cast<unsigned>(pos)));
 }
 
+MlirRegion mlirOperationGetFirstRegion(MlirOperation op) {
+  Operation *cppOp = unwrap(op);
+  if (cppOp->getNumRegions() == 0)
+    return wrap(static_cast<Region *>(nullptr));
+  return wrap(&cppOp->getRegion(0));
+}
+
+MlirRegion mlirRegionGetNextInOperation(MlirRegion region) {
+  Region *cppRegion = unwrap(region);
+  Operation *parent = cppRegion->getParentOp();
+  intptr_t next = cppRegion->getRegionNumber() + 1;
+  if (parent->getNumRegions() > next)
+    return wrap(&parent->getRegion(next));
+  return wrap(static_cast<Region *>(nullptr));
+}
+
 MlirOperation mlirOperationGetNextInBlock(MlirOperation op) {
   return wrap(unwrap(op)->getNextNode());
 }
 
 intptr_t mlirOperationGetNumOperands(MlirOperation op) {
   return static_cast<intptr_t>(unwrap(op)->getNumOperands());
 }
 

mlir/test/CAPI/ir.c

if (mlirValueIsABlockArgument(result))
       return 2;
     if (!mlirOperationEqual(operation, mlirOpResultGetOwner(result)))
       return 3;
     if (i != mlirOpResultGetResultNumber(result))
       return 4;
     ++stats->numOpResults;
   }
 
-  for (unsigned i = 0; i < numRegions; ++i) {
-    MlirRegion region = mlirOperationGetRegion(operation, i);
+  MlirRegion region = mlirOperationGetFirstRegion(operation);
+  while (!mlirRegionIsNull(region)) {
     for (MlirBlock block = mlirRegionGetFirstBlock(region);
          !mlirBlockIsNull(block); block = mlirBlockGetNextInRegion(block)) {
       ++stats->numBlocks;
       intptr_t numArgs = mlirBlockGetNumArguments(block);
       stats->numValues += numArgs;
       for (intptr_t j = 0; j < numArgs; ++j) {
         MlirValue arg = mlirBlockGetArgument(block, j);
         if (!mlirValueIsABlockArgument(arg))
for (MlirBlock block = mlirRegionGetFirstBlock(region);
            !mlirOperationIsNull(child);
            child = mlirOperationGetNextInBlock(child)) {
         OpListNode *node = malloc(sizeof(OpListNode));
         node->op = child;
         node->next = head->next;
         head->next = node;
       }
     }
+    region = mlirRegionGetNextInOperation(region);
   }
   return 0;
 }
 
 int collectStats(MlirOperation operation) {
   OpListNode *head = malloc(sizeof(OpListNode));
   head->op = operation;
   head->next = NULL;