[mlir] Use a type for representing branch points in `RegionBranchOpInterface`

[zero9178]

Summary

The current implementation is not very ergonomic or descriptive: It uses std::optional<unsigned> where std::nullopt represents the parent op and unsigned is the region number.
This doesn't give us any useful methods specific to region control flow and makes the code fragile to changes due to now taking the region number into account.

This patch introduces a new type called RegionBranchPoint, replacing all uses of std::optional<unsigned> in the interface. It can be implicitly constructed from a region or a RegionSuccessor, can be compared with a region to check whether the branch point is branching from the parent, adds isParent to check whether we are coming from a parent op and adds RegionSuccessor::parent as a descriptive way to indicate branching from the parent.

[Mogball]

Glad someone finally went and did this :P

mlir/include/mlir/Interfaces/ControlFlowInterfaces.h

200

FWIW I think it's more idiomatic to have RegionBranchPointer::parent() instead of accessing the singleton directly. It mirrors WalkResult::advance(), LogicalResult::success(), etc.

205

Tagging /*implicit*/ is not usual for the MLIR coding style, or is it? I don't see it very often.

[zero9178]

Address review comments

[zero9178]

mlir/include/mlir/Interfaces/ControlFlowInterfaces.h

205

Seems so! A quick grep and wc showed only 18 instances. I must have gotten the wrong impression

mlir/include/mlir/Analysis/DataFlow/DenseAnalysis.h

protected:
   /// Propagate the dense lattice backwards along the control flow edge from
   /// `regionFrom` to `regionTo` regions of the `branch` operation. `nullopt`
   /// values correspond to control flow branches originating at or targeting the
   /// `branch` operation itself. Default implementation just meets the states,
   /// meaning that operations implementing `RegionBranchOpInterface` don't have
   /// any effect on the lattice that isn't already expressed by the interface
   /// itself.
   virtual void visitRegionBranchControlFlowTransfer(
-      RegionBranchOpInterface branch, std::optional<unsigned> regionFrom,
-      std::optional<unsigned> regionTo, const AbstractDenseLattice &after,
+      RegionBranchOpInterface branch, RegionBranchPoint regionFrom,
+      RegionBranchPoint regionTo, const AbstractDenseLattice &after,
       AbstractDenseLattice *before) {
     meet(before, after);
   }
 
   /// Propagate the dense lattice backwards along the call control flow edge,
   /// which can be either entering or exiting the callee. Default implementation
   /// just meets the states, meaning that operations implementing
   /// `CallOpInterface` don't have any effect on hte lattice that isn't already
private:
   void visitBlock(Block *block);
 
   /// Visit a program point within a region branch operation with successors
   /// (from which the state is propagated) in or after it. `regionNo` indicates
   /// the region that contains the successor, `nullopt` indicating the successor
   /// of the branch operation itself.
   void visitRegionBranchOperation(ProgramPoint point,
                                   RegionBranchOpInterface branch,
-                                  std::optional<unsigned> regionNo,
+                                  RegionBranchPoint branchPoint,
                                   AbstractDenseLattice *before);
 
   /// Visit an operation for which the data flow is described by the
   /// `CallOpInterface`. Performs inter-procedural data flow as follows:
   ///
   ///   - find the callable (resolve via the symbol table),
   ///   - get the entry block of the callable region,
   ///   - take the state before the first operation if present or at block end
public:
   /// Concrete analyses can override this behavior or delegate to the parent
   /// call for the default behavior. Specifically, if the `branch` op may affect
   /// the lattice before entering any region, the custom behavior can be added
   /// for `regionFrom == nullopt`. If the `branch` op may affect the lattice
   /// after all terminated, the custom behavior can be added for `regionTo ==
   /// nullptr`. The behavior can be further refined for specific pairs of "from"
   /// and "to" regions.
   virtual void visitRegionBranchControlFlowTransfer(
-      RegionBranchOpInterface branch, std::optional<unsigned> regionFrom,
-      std::optional<unsigned> regionTo, const LatticeT &after,
-      LatticeT *before) {
+      RegionBranchOpInterface branch, RegionBranchPoint regionFrom,
+      RegionBranchPoint regionTo, const LatticeT &after, LatticeT *before) {
     AbstractDenseBackwardDataFlowAnalysis::visitRegionBranchControlFlowTransfer(
         branch, regionFrom, regionTo, after, before);
   }
 
 protected:
   /// Get the dense lattice at the given program point.
   LatticeT *getLattice(ProgramPoint point) override {
     return getOrCreate<LatticeT>(point);
void visitCallControlFlowTransfer(CallOpInterface call,
                                     CallControlFlowAction action,
                                     const AbstractDenseLattice &after,
                                     AbstractDenseLattice *before) final {
     visitCallControlFlowTransfer(call, action,
                                  static_cast<const LatticeT &>(after),
                                  static_cast<LatticeT *>(before));
   }
   void visitRegionBranchControlFlowTransfer(
-      RegionBranchOpInterface branch, std::optional<unsigned> regionForm,
-      std::optional<unsigned> regionTo, const AbstractDenseLattice &after,
+      RegionBranchOpInterface branch, RegionBranchPoint regionForm,
+      RegionBranchPoint regionTo, const AbstractDenseLattice &after,
       AbstractDenseLattice *before) final {
     visitRegionBranchControlFlowTransfer(branch, regionForm, regionTo,
                                          static_cast<const LatticeT &>(after),
                                          static_cast<LatticeT *>(before));
   }
 };
 
 } // end namespace dataflow
 } // end namespace mlir
 
 #endif // MLIR_ANALYSIS_DENSEDATAFLOWANALYSIS_H

mlir/include/mlir/Analysis/DataFlow/SparseAnalysis.h

private:
   /// determined from block predecessors.
   void visitBlock(Block *block);
 
   /// Visit a program point `point` with predecessors within a region branch
   /// operation `branch`, which can either be the entry block of one of the
   /// regions or the parent operation itself, and set either the argument or
   /// parent result lattices.
   void visitRegionSuccessors(ProgramPoint point, RegionBranchOpInterface branch,
-                             std::optional<unsigned> successorIndex,
+                             RegionBranchPoint successor,
                              ArrayRef<AbstractSparseLattice *> lattices);
 };
 
 //===----------------------------------------------------------------------===//
 // SparseForwardDataFlowAnalysis
 //===----------------------------------------------------------------------===//
 
 /// A sparse forward data-flow analysis for propagating SSA value lattices

mlir/include/mlir/Interfaces/ControlFlowInterfaces.h

public:
   /// the current region.
   ValueRange getSuccessorInputs() const { return inputs; }
 
 private:
   Region *region{nullptr};
   ValueRange inputs;
 };
 
+/// This class represents a point being branched from in the methods of the
+/// `RegionBranchOpInterface`.
+/// One can branch from one of two kinds of places:
+/// * The parent operation (aka the `RegionBranchOpInterface` implementation)
+/// * A region within the parent operation.
+class RegionBranchPoint {
+public:
+  /// Returns an instance of `RegionBranchPoint` representing the parent
+  /// operation.
+  static constexpr RegionBranchPoint parent() { return RegionBranchPoint(); }
+
+  /// Creates a `RegionBranchPoint` that branches from the given region.
+  /// The pointer must not be null.
+  RegionBranchPoint(Region *region) : maybeRegion(region) {
+    assert(region && "Region must not be null");
+  }
+
+  RegionBranchPoint(Region &region) : RegionBranchPoint(&region) {}
+
+  /// Explicitly stops users from constructing with `nullptr`.
+  RegionBranchPoint(std::nullptr_t) = delete;
+
+  /// Constructs a `RegionBranchPoint` from the the target of a
+  /// `RegionSuccessor` instance.
+  RegionBranchPoint(RegionSuccessor successor) {
+    if (successor.isParent())
+      maybeRegion = nullptr;
+    else
+      maybeRegion = successor.getSuccessor();
+  }
+
+  /// Assigns a region being branched from.
+  RegionBranchPoint &operator=(Region &region) {
+    maybeRegion = &region;
+    return *this;
+  }
+
+  /// Returns true if branching from the parent op.
+  bool isParent() const { return maybeRegion == nullptr; }
+
+  /// Returns the region if branching from a region.
+  /// A null pointer otherwise.
+  Region *getRegionOrNull() const { return maybeRegion; }
+
+  /// Returns true if the two branch points are equal.
+  friend bool operator==(RegionBranchPoint lhs, RegionBranchPoint rhs) {
+    return lhs.maybeRegion == rhs.maybeRegion;
+  }
+
+private:
+  // Private constructor to encourage the use of `RegionBranchPoint::parent`.
+  constexpr RegionBranchPoint() : maybeRegion(nullptr) {}
+
+  /// Internal encoding. Uses nullptr for representing branching from the parent
+  /// op and the region being branched from otherwise.
+  Region *maybeRegion;
+};
+
+inline bool operator!=(RegionBranchPoint lhs, RegionBranchPoint rhs) {
+  return !(lhs == rhs);
+}
+
 /// This class represents upper and lower bounds on the number of times a region
 /// of a `RegionBranchOpInterface` can be invoked. The lower bound is at least
 /// zero, but the upper bound may not be known.
 class InvocationBounds {
 public:
   /// Create invocation bounds. The lower bound must be at least 0 and only the
   /// upper bound can be unknown.
   InvocationBounds(unsigned lb, std::optional<unsigned> ub)

mlir/include/mlir/Interfaces/ControlFlowInterfaces.td

let description = [{
     `areTypesCompatible` can be implemented to compare types along control-flow
     edges. By default, type equality is used.
   }];
   let cppNamespace = "::mlir";
 
   let methods = [
     InterfaceMethod<[{
         Returns the operands of this operation used as the entry arguments when
-        entering the region at `index`, which was specified as a successor of
+        branching from `point`, which was specified as a successor of
         this operation by `getEntrySuccessorRegions`, or the operands forwarded
         to the operation's results when it branches back to itself. These operands
         should correspond 1-1 with the successor inputs specified in
         `getEntrySuccessorRegions`.
       }],
       "::mlir::OperandRange", "getEntrySuccessorOperands",
-      (ins "::std::optional<unsigned>":$index), [{}],
+      (ins "::mlir::RegionBranchPoint":$point), [{}],
       /*defaultImplementation=*/[{
         auto operandEnd = this->getOperation()->operand_end();
         return ::mlir::OperandRange(operandEnd, operandEnd);
       }]
     >,
     InterfaceMethod<[{
         Returns the viable region successors that are branched to when first
         executing the op.
         Unlike `getSuccessorRegions`, this method also passes along the
         constant operands of this op. Based on these, different region
         successors can be determined.
         `operands` contains an entry for every operand of the implementing
         op with a null attribute if the operand has no constant value or
         the corresponding attribute if it is a constant.
 
         By default, simply dispatches to `getSuccessorRegions`.
       }],
       "void", "getEntrySuccessorRegions",
       (ins "::llvm::ArrayRef<::mlir::Attribute>":$operands,
            "::llvm::SmallVectorImpl<::mlir::RegionSuccessor> &":$regions),
            [{}], [{
-        $_op.getSuccessorRegions(std::nullopt, regions);
+        $_op.getSuccessorRegions(mlir::RegionBranchPoint::parent(), regions);
       }]
     >,
     InterfaceMethod<[{
-        Returns the viable successors of a region at `index`, or the possible
-        successors when branching from the parent op if `index` is None. These
-        are the regions that may be selected during the flow of control. The
-        parent operation, i.e. a null `index`, may specify itself as successor,
-        which indicates that the control flow may not enter any region at all.
-        This method allows for describing which regions may be executed when
-        entering an operation, and which regions are executed after having
-        executed another region of the parent op. The successor region must be
-        non-empty.
+        Returns the viable successors of `point`. These are the regions that may
+        be selected during the flow of control. The parent operation, may
+        specify itself as successor, which indicates that the control flow may
+        not enter any region at all. This method allows for describing which
+        regions may be executed when entering an operation, and which regions
+        are executed after having executed another region of the parent op. The
+        successor region must be non-empty.
       }],
       "void", "getSuccessorRegions",
-      (ins "::std::optional<unsigned>":$index,
+      (ins "::mlir::RegionBranchPoint":$point,
            "::llvm::SmallVectorImpl<::mlir::RegionSuccessor> &":$regions)
     >,
     InterfaceMethod<[{
         Populates `invocationBounds` with the minimum and maximum number of
         times this operation will invoke the attached regions (assuming the
         regions yield normally, i.e. do not abort or invoke an infinite loop).
         The minimum number of invocations is at least 0. If the maximum number
         of invocations cannot be statically determined, then it will not have a
let description = [{
     in the presence of a parent RegionBranchOpInterface implementation. It
     specifies which operands are passed to which successor region.
   }];
   let cppNamespace = "::mlir";
 
   let methods = [
     InterfaceMethod<[{
         Returns a mutable range of operands that are semantically "returned" by
-        passing them to the region successor given by `index`.  If `index` is
-        None, this function returns the operands that are passed as a result to
-        the parent operation.
+        passing them to the region successor given by `point`.
       }],
       "::mlir::MutableOperandRange", "getMutableSuccessorOperands",
-      (ins "::std::optional<unsigned>":$index)
+      (ins "::mlir::RegionBranchPoint":$point)
     >,
     InterfaceMethod<[{
         Returns the viable region successors that are branched to after this
         terminator based on the given constant operands.
 
         `operands` contains an entry for every operand of the
         implementing op with a null attribute if the operand has no constant
         value or the corresponding attribute if it is a constant.
 
         Default implementation simply dispatches to the parent
         `RegionBranchOpInterface`'s `getSuccessorRegions` implementation.
       }],
       "void", "getSuccessorRegions",
       (ins "::llvm::ArrayRef<::mlir::Attribute>":$operands,
            "::llvm::SmallVectorImpl<::mlir::RegionSuccessor> &":$regions), [{}],
       [{
         ::mlir::Operation *op = $_op;
         ::llvm::cast<::mlir::RegionBranchOpInterface>(op->getParentOp())
-          .getSuccessorRegions(op->getParentRegion()->getRegionNumber(),
-            regions);
+          .getSuccessorRegions(op->getParentRegion(), regions);
       }]
     >,
   ];
 
   let verify = [{
     static_assert(ConcreteOp::template hasTrait<OpTrait::IsTerminator>(),
                   "expected operation to be a terminator");
     static_assert(ConcreteOp::template hasTrait<OpTrait::ZeroResults>(),
                   "expected operation to have zero results");
     static_assert(ConcreteOp::template hasTrait<OpTrait::ZeroSuccessors>(),
                   "expected operation to have zero successors");
     return success();
   }];
 
   let extraClassDeclaration = [{
     // Returns a range of operands that are semantically "returned" by passing
     // them to the region successor given by `index`.  If `index` is None, this
     // function returns the operands that are passed as a result to the parent
     // operation.
-    ::mlir::OperandRange getSuccessorOperands(std::optional<unsigned> index) {
-      return getMutableSuccessorOperands(index);
+    ::mlir::OperandRange getSuccessorOperands(::mlir::RegionBranchPoint point) {
+      return getMutableSuccessorOperands(point);
     }
   }];
 }
 
 //===----------------------------------------------------------------------===//
 // ControlFlow Traits
 //===----------------------------------------------------------------------===//
 
 // Op is "return-like".
 def ReturnLike : TraitList<[
     DeclareOpInterfaceMethods<RegionBranchTerminatorOpInterface>,
     NativeOpTrait<
         /*name=*/"ReturnLike",
         /*traits=*/[],
         /*extraOpDeclaration=*/"",
         /*extraOpDefinition=*/[{
           ::mlir::MutableOperandRange $cppClass::getMutableSuccessorOperands(
-            ::std::optional<unsigned> index) {
+            ::mlir::RegionBranchPoint point) {
             return ::mlir::MutableOperandRange(*this);
           }
         }]
     >
 ]>;
 
 #endif // MLIR_INTERFACES_CONTROLFLOWINTERFACES

mlir/lib/Analysis/AliasAnalysis/LocalAliasAnalysis.cpp

static void collectUnderlyingAddressValues(RegionBranchOpInterface branch,
                                            unsigned maxDepth,
                                            DenseSet<Value> &visited,
                                            SmallVectorImpl<Value> &output) {
   // Given the index of a region of the branch (`predIndex`), or std::nullopt to
   // represent the parent operation, try to return the index into the outputs of
   // this region predecessor that correspond to the input values of `region`. If
   // an index could not be found, std::nullopt is returned instead.
   auto getOperandIndexIfPred =
-      [&](std::optional<unsigned> predIndex) -> std::optional<unsigned> {
+      [&](RegionBranchPoint pred) -> std::optional<unsigned> {
     SmallVector<RegionSuccessor, 2> successors;
-    branch.getSuccessorRegions(predIndex, successors);
+    branch.getSuccessorRegions(pred, successors);
     for (RegionSuccessor &successor : successors) {
       if (successor.getSuccessor() != region)
         continue;
       // Check that the successor inputs map to the given input value.
       ValueRange inputs = successor.getSuccessorInputs();
       if (inputs.empty()) {
         output.push_back(inputValue);
         break;
for (RegionSuccessor &successor : successors) {
         break;
       }
       return inputIndex - firstInputIndex;
     }
     return std::nullopt;
   };
 
   // Check branches from the parent operation.
-  std::optional<unsigned> regionIndex;
-  if (region) {
-    // Determine the actual region number from the passed region.
-    regionIndex = region->getRegionNumber();
-  }
+  auto branchPoint = RegionBranchPoint::parent();
+  if (region)
+    branchPoint = region;
+
   if (std::optional<unsigned> operandIndex =
-          getOperandIndexIfPred(/*predIndex=*/std::nullopt)) {
+          getOperandIndexIfPred(/*predIndex=*/RegionBranchPoint::parent())) {
     collectUnderlyingAddressValues(
-        branch.getEntrySuccessorOperands(regionIndex)[*operandIndex], maxDepth,
+        branch.getEntrySuccessorOperands(branchPoint)[*operandIndex], maxDepth,
         visited, output);
   }
   // Check branches from each child region.
   Operation *op = branch.getOperation();
-  for (int i = 0, e = op->getNumRegions(); i != e; ++i) {
-    if (std::optional<unsigned> operandIndex = getOperandIndexIfPred(i)) {
-      for (Block &block : op->getRegion(i)) {
+  for (Region &region : op->getRegions()) {
+    if (std::optional<unsigned> operandIndex = getOperandIndexIfPred(region)) {
+      for (Block &block : region) {
         // Try to determine possible region-branch successor operands for the
         // current region.
         if (auto term = dyn_cast<RegionBranchTerminatorOpInterface>(
                 block.getTerminator())) {
           collectUnderlyingAddressValues(
-              term.getSuccessorOperands(regionIndex)[*operandIndex], maxDepth,
+              term.getSuccessorOperands(branchPoint)[*operandIndex], maxDepth,
               visited, output);
         } else if (block.getNumSuccessors()) {
           // Otherwise, if this terminator may exit the region we can't make
           // any assumptions about which values get passed.
           output.push_back(inputValue);
           return;
         }
       }

mlir/lib/Analysis/DataFlow/DenseAnalysis.cpp

void AbstractDenseBackwardDataFlowAnalysis::processOperation(Operation *op) {
   const AbstractDenseLattice *after;
   if (Operation *next = op->getNextNode())
     after = getLatticeFor(op, next);
   else
     after = getLatticeFor(op, op->getBlock());
 
   // Special cases where control flow may dictate data flow.
   if (auto branch = dyn_cast<RegionBranchOpInterface>(op))
-    return visitRegionBranchOperation(op, branch, std::nullopt, before);
+    return visitRegionBranchOperation(op, branch, RegionBranchPoint::parent(),
+                                      before);
   if (auto call = dyn_cast<CallOpInterface>(op))
     return visitCallOperation(call, before);
 
   // Invoke the operation transfer function.
   visitOperationImpl(op, *after, before);
 }
 
 void AbstractDenseBackwardDataFlowAnalysis::visitBlock(Block *block) {
if (callable && callable.getCallableRegion() == block->getParent()) {
                                      before);
       }
       return;
     }
 
     // If this block is exiting from an operation with region-based control
     // flow, propagate the lattice back along the control flow edge.
     if (auto branch = dyn_cast<RegionBranchOpInterface>(block->getParentOp())) {
-      visitRegionBranchOperation(block, branch,
-                                 block->getParent()->getRegionNumber(), before);
+      visitRegionBranchOperation(block, branch, block->getParent(), before);
       return;
     }
 
     // Cannot reason about successors of an exit block, set the pessimistic
     // fixpoint.
     return setToExitState(before);
   }
 
if (successor->empty())
       meet(before, *getLatticeFor(block, successor));
     else
       meet(before, *getLatticeFor(block, &successor->front()));
   }
 }
 
 void AbstractDenseBackwardDataFlowAnalysis::visitRegionBranchOperation(
     ProgramPoint point, RegionBranchOpInterface branch,
-    std::optional<unsigned> regionNo, AbstractDenseLattice *before) {
+    RegionBranchPoint branchPoint, AbstractDenseLattice *before) {
 
   // The successors of the operation may be either the first operation of the
   // entry block of each possible successor region, or the next operation when
   // the branch is a successor of itself.
   SmallVector<RegionSuccessor> successors;
-  branch.getSuccessorRegions(regionNo, successors);
+  branch.getSuccessorRegions(branchPoint, successors);
   for (const RegionSuccessor &successor : successors) {
     const AbstractDenseLattice *after;
     if (successor.isParent() || successor.getSuccessor()->empty()) {
       if (Operation *next = branch->getNextNode())
         after = getLatticeFor(point, next);
       else
         after = getLatticeFor(point, branch->getBlock());
     } else {
       Region *successorRegion = successor.getSuccessor();
       assert(!successorRegion->empty() && "unexpected empty successor region");
       Block *successorBlock = &successorRegion->front();
 
       if (!getOrCreateFor<Executable>(point, successorBlock)->isLive())
         continue;
 
       if (successorBlock->empty())
         after = getLatticeFor(point, successorBlock);
       else
         after = getLatticeFor(point, &successorBlock->front());
     }
-    std::optional<unsigned> successorNo =
-        successor.isParent() ? std::optional<unsigned>()
-                             : successor.getSuccessor()->getRegionNumber();
-    visitRegionBranchControlFlowTransfer(branch, regionNo, successorNo, *after,
+
+    visitRegionBranchControlFlowTransfer(branch, branchPoint, successor, *after,
                                          before);
   }
 }
 
 const AbstractDenseLattice *
 AbstractDenseBackwardDataFlowAnalysis::getLatticeFor(ProgramPoint dependent,
                                                      ProgramPoint point) {
   AbstractDenseLattice *state = getLattice(point);
   addDependency(state, dependent);
   return state;
 }

mlir/lib/Analysis/DataFlow/SparseAnalysis.cpp

void AbstractSparseForwardDataFlowAnalysis::visitOperation(Operation *op) {
   for (Value result : op->getResults()) {
     AbstractSparseLattice *resultLattice = getLatticeElement(result);
     resultLattices.push_back(resultLattice);
   }
 
   // The results of a region branch operation are determined by control-flow.
   if (auto branch = dyn_cast<RegionBranchOpInterface>(op)) {
     return visitRegionSuccessors({branch}, branch,
-                                 /*successorIndex=*/std::nullopt,
+                                 /*successor=*/RegionBranchPoint::parent(),
                                  resultLattices);
   }
 
   // The results of a call operation are determined by the callgraph.
   if (auto call = dyn_cast<CallOpInterface>(op)) {
     const auto *predecessors = getOrCreateFor<PredecessorState>(op, call);
     // If not all return sites are known, then conservatively assume we can't
     // reason about the data-flow.
if (callable && callable.getCallableRegion() == block->getParent()) {
         for (auto it : llvm::zip(call.getArgOperands(), argLattices))
           join(std::get<1>(it), *getLatticeElementFor(block, std::get<0>(it)));
       }
       return;
     }
 
     // Check if the lattices can be determined from region control flow.
     if (auto branch = dyn_cast<RegionBranchOpInterface>(block->getParentOp())) {
-      return visitRegionSuccessors(
-          block, branch, block->getParent()->getRegionNumber(), argLattices);
+      return visitRegionSuccessors(block, branch, block->getParent(),
+                                   argLattices);
     }
 
     // Otherwise, we can't reason about the data-flow.
     return visitNonControlFlowArgumentsImpl(block->getParentOp(),
                                             RegionSuccessor(block->getParent()),
                                             argLattices, /*firstIndex=*/0);
   }
 
for (Block::pred_iterator it = block->pred_begin(), e = block->pred_end();
     } else {
       return setAllToEntryStates(argLattices);
     }
   }
 }
 
 void AbstractSparseForwardDataFlowAnalysis::visitRegionSuccessors(
     ProgramPoint point, RegionBranchOpInterface branch,
-    std::optional<unsigned> successorIndex,
-    ArrayRef<AbstractSparseLattice *> lattices) {
+    RegionBranchPoint successor, ArrayRef<AbstractSparseLattice *> lattices) {
   const auto *predecessors = getOrCreateFor<PredecessorState>(point, point);
   assert(predecessors->allPredecessorsKnown() &&
          "unexpected unresolved region successors");
 
   for (Operation *op : predecessors->getKnownPredecessors()) {
     // Get the incoming successor operands.
     std::optional<OperandRange> operands;
 
     // Check if the predecessor is the parent op.
     if (op == branch) {
-      operands = branch.getEntrySuccessorOperands(successorIndex);
+      operands = branch.getEntrySuccessorOperands(successor);
       // Otherwise, try to deduce the operands from a region return-like op.
     } else if (auto regionTerminator =
                    dyn_cast<RegionBranchTerminatorOpInterface>(op)) {
-      operands = regionTerminator.getSuccessorOperands(successorIndex);
+      operands = regionTerminator.getSuccessorOperands(successor);
     }
 
     if (!operands) {
       // We can't reason about the data-flow.
       return setAllToEntryStates(lattices);
     }
 
     ValueRange inputs = predecessors->getSuccessorInputs(op);
void AbstractSparseBackwardDataFlowAnalysis::visitRegionSuccessors(
   SmallVector<Attribute> operands(op->getNumOperands(), nullptr);
   branch.getEntrySuccessorRegions(operands, successors);
 
   // All operands not forwarded to any successor. This set can be non-contiguous
   // in the presence of multiple successors.
   BitVector unaccounted(op->getNumOperands(), true);
 
   for (RegionSuccessor &successor : successors) {
-    Region *region = successor.getSuccessor();
-    OperandRange operands =
-        region ? branch.getEntrySuccessorOperands(region->getRegionNumber())
-               : branch.getEntrySuccessorOperands({});
+    OperandRange operands = branch.getEntrySuccessorOperands(successor);
     MutableArrayRef<OpOperand> opoperands = operandsToOpOperands(operands);
     ValueRange inputs = successor.getSuccessorInputs();
     for (auto [operand, input] : llvm::zip(opoperands, inputs)) {
       meet(getLatticeElement(operand.get()), *getLatticeElementFor(op, input));
       unaccounted.reset(operand.getOperandNumber());
     }
   }
   // All operands not forwarded to regions are typically parameters of the
void AbstractSparseBackwardDataFlowAnalysis::
   SmallVector<RegionSuccessor> successors;
   terminator.getSuccessorRegions(operandAttributes, successors);
   // All operands not forwarded to any successor. This set can be
   // non-contiguous in the presence of multiple successors.
   BitVector unaccounted(terminator->getNumOperands(), true);
 
   for (const RegionSuccessor &successor : successors) {
     ValueRange inputs = successor.getSuccessorInputs();
-    Region *region = successor.getSuccessor();
-    OperandRange operands = terminator.getSuccessorOperands(
-        region ? region->getRegionNumber() : std::optional<unsigned>{});
+    OperandRange operands = terminator.getSuccessorOperands(successor);
     MutableArrayRef<OpOperand> opOperands = operandsToOpOperands(operands);
     for (auto [opOperand, input] : llvm::zip(opOperands, inputs)) {
       meet(getLatticeElement(opOperand.get()),
            *getLatticeElementFor(terminator, input));
       unaccounted.reset(const_cast<OpOperand &>(opOperand).getOperandNumber());
     }
   }
   // Visit operands of the branch op not forwarded to the next region.

mlir/lib/Dialect/Affine/IR/AffineOps.cpp

void AffineForOp::getCanonicalizationPatterns(RewritePatternSet &results,
                                               MLIRContext *context) {
   results.add<AffineForEmptyLoopFolder>(context);
 }
 
 /// Return operands used when entering the region at 'index'. These operands
 /// correspond to the loop iterator operands, i.e., those excluding the
 /// induction variable. AffineForOp only has one region, so zero is the only
 /// valid value for `index`.
-OperandRange
-AffineForOp::getEntrySuccessorOperands(std::optional<unsigned> index) {
-  assert((!index || *index == 0) && "invalid region index");
+OperandRange AffineForOp::getEntrySuccessorOperands(RegionBranchPoint point) {
+  assert((point.isParent() || point == getLoopBody()) &&
+         "invalid region point");
 
   // The initial operands map to the loop arguments after the induction
   // variable or are forwarded to the results when the trip count is zero.
   return getIterOperands();
 }
 
 /// Given the region at `index`, or the parent operation if `index` is None,
 /// return the successor regions. These are the regions that may be selected
 /// during the flow of control. `operands` is a set of optional attributes that
 /// correspond to a constant value for each operand, or null if that operand is
 /// not a constant.
 void AffineForOp::getSuccessorRegions(
-    std::optional<unsigned> index, SmallVectorImpl<RegionSuccessor> &regions) {
-  assert((!index.has_value() || index.value() == 0) && "expected loop region");
+    RegionBranchPoint point, SmallVectorImpl<RegionSuccessor> &regions) {
+  assert((point.isParent() || point == getLoopBody()) &&
+         "expected loop region");
   // The loop may typically branch back to its body or to the parent operation.
   // If the predecessor is the parent op and the trip count is known to be at
   // least one, branch into the body using the iterator arguments. And in cases
   // we know the trip count is zero, it can only branch back to its parent.
   std::optional<uint64_t> tripCount = getTrivialConstantTripCount(*this);
-  if (!index.has_value() && tripCount.has_value()) {
+  if (point.isParent() && tripCount.has_value()) {
     if (tripCount.value() > 0) {
       regions.push_back(RegionSuccessor(&getLoopBody(), getRegionIterArgs()));
       return;
     }
     if (tripCount.value() == 0) {
       regions.push_back(RegionSuccessor(getResults()));
       return;
     }
   }
 
   // From the loop body, if the trip count is one, we can only branch back to
   // the parent.
-  if (index && tripCount && *tripCount == 1) {
+  if (!point.isParent() && tripCount && *tripCount == 1) {
     regions.push_back(RegionSuccessor(getResults()));
     return;
   }
 
   // In all other cases, the loop may branch back to itself or the parent
   // operation.
   regions.push_back(RegionSuccessor(&getLoopBody(), getRegionIterArgs()));
   regions.push_back(RegionSuccessor(getResults()));
LogicalResult matchAndRewrite(AffineIfOp op,
     return success();
   }
 };
 } // namespace
 
 /// AffineIfOp has two regions -- `then` and `else`. The flow of data should be
 /// as follows: AffineIfOp -> `then`/`else` -> AffineIfOp
 void AffineIfOp::getSuccessorRegions(
-    std::optional<unsigned> index, SmallVectorImpl<RegionSuccessor> &regions) {
+    RegionBranchPoint point, SmallVectorImpl<RegionSuccessor> &regions) {
   // If the predecessor is an AffineIfOp, then branching into both `then` and
   // `else` region is valid.
-  if (!index.has_value()) {
+  if (point.isParent()) {
     regions.reserve(2);
     regions.push_back(
         RegionSuccessor(&getThenRegion(), getThenRegion().getArguments()));
     // If the "else" region is empty, branch bach into parent.
     if (getElseRegion().empty()) {
       regions.push_back(getResults());
     } else {
       regions.push_back(

mlir/lib/Dialect/Async/IR/Async.cpp

 }
 
 //===----------------------------------------------------------------------===//
 /// ExecuteOp
 //===----------------------------------------------------------------------===//
 
 constexpr char kOperandSegmentSizesAttr[] = "operandSegmentSizes";
 
-OperandRange
-ExecuteOp::getEntrySuccessorOperands(std::optional<unsigned> index) {
-  assert(index && *index == 0 && "invalid region index");
+OperandRange ExecuteOp::getEntrySuccessorOperands(RegionBranchPoint point) {
+  assert(point == getBodyRegion() && "invalid region index");
   return getBodyOperands();
 }
 
 bool ExecuteOp::areTypesCompatible(Type lhs, Type rhs) {
   const auto getValueOrTokenType = [](Type type) {
     if (auto value = llvm::dyn_cast<ValueType>(type))
       return value.getValueType();
     return type;
   };
   return getValueOrTokenType(lhs) == getValueOrTokenType(rhs);
 }
 
-void ExecuteOp::getSuccessorRegions(std::optional<unsigned> index,
+void ExecuteOp::getSuccessorRegions(RegionBranchPoint point,
                                     SmallVectorImpl<RegionSuccessor> &regions) {
   // The `body` region branch back to the parent operation.
-  if (index) {
-    assert(*index == 0 && "invalid region index");
+  if (point == getBodyRegion()) {
     regions.push_back(RegionSuccessor(getBodyResults()));
     return;
   }
 
   // Otherwise the successor is the body region.
   regions.push_back(
       RegionSuccessor(&getBodyRegion(), getBodyRegion().getArguments()));
 }

mlir/lib/Dialect/Bufferization/Transforms/BufferDeallocation.cpp

if (failed(introduceClonesForRegionSuccessors(
               return successorRegion.getSuccessor() == argRegion;
             })))
       return failure();
 
     // Check whether the block argument belongs to an entry region of the
     // parent operation. In this case, we have to introduce an additional clone
     // for buffer that is passed to the argument.
     SmallVector<RegionSuccessor, 2> successorRegions;
-    regionInterface.getSuccessorRegions(/*index=*/std::nullopt,
+    regionInterface.getSuccessorRegions(/*point=*/RegionBranchPoint::parent(),
                                         successorRegions);
     auto *it =
         llvm::find_if(successorRegions, [&](RegionSuccessor &successorRegion) {
           return successorRegion.getSuccessor() == argRegion;
         });
     if (it == successorRegions.end())
       return success();
 
     // Determine the actual operand to introduce a clone for and rewire the
     // operand to point to the clone instead.
-    auto operands =
-        regionInterface.getEntrySuccessorOperands(argRegion->getRegionNumber());
+    auto operands = regionInterface.getEntrySuccessorOperands(argRegion);
     size_t operandIndex =
         llvm::find(it->getSuccessorInputs(), blockArg).getIndex() +
         operands.getBeginOperandIndex();
     Value operand = parentOp->getOperand(operandIndex);
     assert(operand ==
                operands[operandIndex - operands.getBeginOperandIndex()] &&
            "region interface operands don't match parentOp operands");
     auto clone = introduceCloneBuffers(operand, parentOp);
private:
   template <typename TPredicate>
   LogicalResult introduceClonesForRegionSuccessors(
       RegionBranchOpInterface regionInterface, MutableArrayRef<Region> regions,
       Value argValue, const TPredicate &regionPredicate) {
     for (Region &region : regions) {
       // Query the regionInterface to get all successor regions of the current
       // one.
       SmallVector<RegionSuccessor, 2> successorRegions;
-      regionInterface.getSuccessorRegions(region.getRegionNumber(),
-                                          successorRegions);
+      regionInterface.getSuccessorRegions(region, successorRegions);
       // Try to find a matching region successor.
       RegionSuccessor *regionSuccessor =
           llvm::find_if(successorRegions, regionPredicate);
       if (regionSuccessor == successorRegions.end())
         continue;
       // Get the operand index in the context of the current successor input
       // bindings.
       size_t operandIndex =
           llvm::find(regionSuccessor->getSuccessorInputs(), argValue)
               .getIndex();
 
-      std::optional<unsigned> successorRegionNumber;
-      if (Region *successorRegion = regionSuccessor->getSuccessor())
-        successorRegionNumber = successorRegion->getRegionNumber();
-
       // Iterate over all immediate terminator operations to introduce
       // new buffer allocations. Thereby, the appropriate terminator operand
       // will be adjusted to point to the newly allocated buffer instead.
       if (failed(walkReturnOperations(
               &region, [&](RegionBranchTerminatorOpInterface terminator) {
                 // Get the actual mutable operands for this terminator op.
                 auto terminatorOperands =
-                    terminator.getMutableSuccessorOperands(
-                        successorRegionNumber);
+                    terminator.getMutableSuccessorOperands(*regionSuccessor);
                 // Extract the source value from the current terminator.
                 // This conversion needs to exist on a separate line due to a
                 // bug in GCC conversion analysis.
                 OperandRange immutableTerminatorOperands = terminatorOperands;
                 Value sourceValue = immutableTerminatorOperands[operandIndex];
                 // Create a new clone at the current location of the terminator.
                 auto clone = introduceCloneBuffers(sourceValue, terminator);
                 if (failed(clone))

mlir/lib/Dialect/Bufferization/Transforms/BufferUtils.cpp

bool BufferPlacementTransformationBase::isLoop(Operation *op) {
   std::function<bool(Region *)> recurse = [&](Region *current) {
     if (!current)
       return false;
     // If we have found a back edge, the parent operation induces a loop.
     if (!visitedRegions.insert(current).second)
       return true;
     // Recurses into all region successors.
     SmallVector<RegionSuccessor, 2> successors;
-    regionInterface.getSuccessorRegions(current->getRegionNumber(), successors);
+    regionInterface.getSuccessorRegions(current, successors);
     for (RegionSuccessor &regionEntry : successors)
       if (recurse(regionEntry.getSuccessor()))
         return true;
     return false;
   };
 
   // Start with all entry regions and test whether they induce a loop.
   SmallVector<RegionSuccessor, 2> successorRegions;
-  regionInterface.getSuccessorRegions(/*index=*/std::nullopt, successorRegions);
+  regionInterface.getSuccessorRegions(/*point=*/RegionBranchPoint::parent(),
+                                      successorRegions);
   for (RegionSuccessor &regionEntry : successorRegions) {
     if (recurse(regionEntry.getSuccessor()))
       return true;
     visitedRegions.clear();
   }
 
   return false;
 }

mlir/lib/Dialect/Bufferization/Transforms/BufferViewFlowAnalysis.cpp

if (auto branchInterface = dyn_cast<BranchOpInterface>(op)) {
       }
       return WalkResult::advance();
     }
 
     if (auto regionInterface = dyn_cast<RegionBranchOpInterface>(op)) {
       // Query the RegionBranchOpInterface to find potential successor regions.
       // Extract all entry regions and wire all initial entry successor inputs.
       SmallVector<RegionSuccessor, 2> entrySuccessors;
-      regionInterface.getSuccessorRegions(/*index=*/std::nullopt,
+      regionInterface.getSuccessorRegions(/*point=*/RegionBranchPoint::parent(),
                                           entrySuccessors);
       for (RegionSuccessor &entrySuccessor : entrySuccessors) {
         // Wire the entry region's successor arguments with the initial
         // successor inputs.
         registerDependencies(
-            regionInterface.getEntrySuccessorOperands(
-                entrySuccessor.isParent()
-                    ? std::optional<unsigned>()
-                    : entrySuccessor.getSuccessor()->getRegionNumber()),
+            regionInterface.getEntrySuccessorOperands(entrySuccessor),
             entrySuccessor.getSuccessorInputs());
       }
 
       // Wire flow between regions and from region exits.
       for (Region &region : regionInterface->getRegions()) {
         // Iterate over all successor region entries that are reachable from the
         // current region.
         SmallVector<RegionSuccessor, 2> successorRegions;
-        regionInterface.getSuccessorRegions(region.getRegionNumber(),
-                                            successorRegions);
+        regionInterface.getSuccessorRegions(region, successorRegions);
         for (RegionSuccessor &successorRegion : successorRegions) {
-          // Determine the current region index (if any).
-          std::optional<unsigned> regionIndex;
-          Region *regionSuccessor = successorRegion.getSuccessor();
-          if (regionSuccessor)
-            regionIndex = regionSuccessor->getRegionNumber();
           // Iterate over all immediate terminator operations and wire the
           // successor inputs with the successor operands of each terminator.
           for (Block &block : region)
             if (auto terminator = dyn_cast<RegionBranchTerminatorOpInterface>(
                     block.getTerminator()))
-              registerDependencies(terminator.getSuccessorOperands(regionIndex),
-                                   successorRegion.getSuccessorInputs());
+              registerDependencies(
+                  terminator.getSuccessorOperands(successorRegion),
         }
       }
 
       return WalkResult::advance();
     }
 
     // Unknown op: Assume that all operands alias with all results.
     for (Value operand : op->getOperands()) {

mlir/lib/Dialect/MemRef/IR/MemRefOps.cpp

ParseResult AllocaScopeOp::parse(OpAsmParser &parser, OperationState &result) {
   // Parse the optional attribute list.
   if (parser.parseOptionalAttrDict(result.attributes))
     return failure();
 
   return success();
 }
 
 void AllocaScopeOp::getSuccessorRegions(
-    std::optional<unsigned> index, SmallVectorImpl<RegionSuccessor> &regions) {
-  if (index) {
+    RegionBranchPoint point, SmallVectorImpl<RegionSuccessor> &regions) {
+  if (!point.isParent()) {
     regions.push_back(RegionSuccessor(getResults()));
     return;
   }
 
   regions.push_back(RegionSuccessor(&getBodyRegion()));
 }
 
 /// Given an operation, return whether this op is guaranteed to

mlir/lib/Dialect/SCF/IR/SCF.cpp

 }
 
 /// Given the region at `index`, or the parent operation if `index` is None,
 /// return the successor regions. These are the regions that may be selected
 /// during the flow of control. `operands` is a set of optional attributes that
 /// correspond to a constant value for each operand, or null if that operand is
 /// not a constant.
 void ExecuteRegionOp::getSuccessorRegions(
-    std::optional<unsigned> index, SmallVectorImpl<RegionSuccessor> &regions) {
+    RegionBranchPoint point, SmallVectorImpl<RegionSuccessor> &regions) {
   // If the predecessor is the ExecuteRegionOp, branch into the body.
-  if (!index) {
+  if (point.isParent()) {
     regions.push_back(RegionSuccessor(&getRegion()));
     return;
   }
 
   // Otherwise, the region branches back to the parent operation.
   regions.push_back(RegionSuccessor(getResults()));
 }
 
 //===----------------------------------------------------------------------===//
 // ConditionOp
 //===----------------------------------------------------------------------===//
 
 MutableOperandRange
-ConditionOp::getMutableSuccessorOperands(std::optional<unsigned> index) {
-  assert((!index || index == getParentOp().getAfter().getRegionNumber()) &&
+ConditionOp::getMutableSuccessorOperands(RegionBranchPoint point) {
+  assert((point.isParent() || point == getParentOp().getAfter()) &&
          "condition op can only exit the loop or branch to the after"
          "region");
   // Pass all operands except the condition to the successor region.
   return getArgsMutable();
 }
 
 void ConditionOp::getSuccessorRegions(
     ArrayRef<Attribute> operands, SmallVectorImpl<RegionSuccessor> &regions) {
ForOp mlir::scf::getForInductionVarOwner(Value val) {
   assert(ivArg.getOwner() && "unlinked block argument");
   auto *containingOp = ivArg.getOwner()->getParentOp();
   return dyn_cast_or_null<ForOp>(containingOp);
 }
 
 /// Return operands used when entering the region at 'index'. These operands
 /// correspond to the loop iterator operands, i.e., those excluding the
 /// induction variable.
-OperandRange ForOp::getEntrySuccessorOperands(std::optional<unsigned> index) {
+OperandRange ForOp::getEntrySuccessorOperands(RegionBranchPoint point) {
   return getInitArgs();
 }
 
 /// Given the region at `index`, or the parent operation if `index` is None,
 /// return the successor regions. These are the regions that may be selected
 /// during the flow of control. `operands` is a set of optional attributes that
 /// correspond to a constant value for each operand, or null if that operand is
 /// not a constant.
-void ForOp::getSuccessorRegions(std::optional<unsigned> index,
+void ForOp::getSuccessorRegions(RegionBranchPoint point,
                                 SmallVectorImpl<RegionSuccessor> &regions) {
   // Both the operation itself and the region may be branching into the body or
   // back into the operation itself. It is possible for loop not to enter the
   // body.
   regions.push_back(RegionSuccessor(&getLoopBody(), getRegionIterArgs()));
   regions.push_back(RegionSuccessor(getResults()));
 }
 
results.add<DimOfForallOp, FoldTensorCastOfOutputIntoForallOp,
               ForallOpSingleOrZeroIterationDimsFolder>(context);
 }
 
 /// Given the region at `index`, or the parent operation if `index` is None,
 /// return the successor regions. These are the regions that may be selected
 /// during the flow of control. `operands` is a set of optional attributes that
 /// correspond to a constant value for each operand, or null if that operand is
 /// not a constant.
-void ForallOp::getSuccessorRegions(std::optional<unsigned> index,
+void ForallOp::getSuccessorRegions(RegionBranchPoint point,
                                    SmallVectorImpl<RegionSuccessor> &regions) {
   // Both the operation itself and the region may be branching into the body or
   // back into the operation itself. It is possible for loop not to enter the
   // body.
   regions.push_back(RegionSuccessor(&getRegion()));
   regions.push_back(RegionSuccessor());
 }
 
void IfOp::print(OpAsmPrinter &p) {
   p.printOptionalAttrDict((*this)->getAttrs());
 }
 
 /// Given the region at `index`, or the parent operation if `index` is None,
 /// return the successor regions. These are the regions that may be selected
 /// during the flow of control. `operands` is a set of optional attributes that
 /// correspond to a constant value for each operand, or null if that operand is
 /// not a constant.
-void IfOp::getSuccessorRegions(std::optional<unsigned> index,
+void IfOp::getSuccessorRegions(RegionBranchPoint point,
                                SmallVectorImpl<RegionSuccessor> &regions) {
   // The `then` and the `else` region branch back to the parent operation.
-  if (index) {
+  if (!point.isParent()) {
     regions.push_back(RegionSuccessor(getResults()));
     return;
   }
 
   regions.push_back(RegionSuccessor(&getThenRegion()));
 
   // Don't consider the else region if it is empty.
   Region *elseRegion = &this->getElseRegion();

 }
 
 /// Given the region at `index`, or the parent operation if `index` is None,
 /// return the successor regions. These are the regions that may be selected
 /// during the flow of control. `operands` is a set of optional attributes that
 /// correspond to a constant value for each operand, or null if that operand is
 /// not a constant.
 void ParallelOp::getSuccessorRegions(
-    std::optional<unsigned> index, SmallVectorImpl<RegionSuccessor> &regions) {
+    RegionBranchPoint point, SmallVectorImpl<RegionSuccessor> &regions) {
   // Both the operation itself and the region may be branching into the body or
   // back into the operation itself. It is possible for loop not to enter the
   // body.
   regions.push_back(RegionSuccessor(&getRegion()));
   regions.push_back(RegionSuccessor());
 }
 
 //===----------------------------------------------------------------------===//
void WhileOp::build(::mlir::OpBuilder &odsBuilder,
   Region *afterRegion = odsState.addRegion();
   Block *afterBlock = odsBuilder.createBlock(afterRegion, /*insertPt=*/{},
                                              resultTypes, afterArgLocs);
 
   if (afterBuilder)
     afterBuilder(odsBuilder, odsState.location, afterBlock->getArguments());
 }
 
-OperandRange WhileOp::getEntrySuccessorOperands(std::optional<unsigned> index) {
-  assert(index && *index == 0 &&
+OperandRange WhileOp::getEntrySuccessorOperands(RegionBranchPoint point) {
+  assert(point == getBefore() &&
          "WhileOp is expected to branch only to the first region");
 
   return getInits();
 }
 
 ConditionOp WhileOp::getConditionOp() {
   return cast<ConditionOp>(getBeforeBody()->getTerminator());
 }
 
 YieldOp WhileOp::getYieldOp() {
   return cast<YieldOp>(getAfterBody()->getTerminator());
 }
 
 Block::BlockArgListType WhileOp::getBeforeArguments() {
   return getBeforeBody()->getArguments();
 }
 
 Block::BlockArgListType WhileOp::getAfterArguments() {
   return getAfterBody()->getArguments();
 }
 
-void WhileOp::getSuccessorRegions(std::optional<unsigned> index,
+void WhileOp::getSuccessorRegions(RegionBranchPoint point,
                                   SmallVectorImpl<RegionSuccessor> &regions) {
   // The parent op always branches to the condition region.
-  if (!index) {
+  if (point.isParent()) {
     regions.emplace_back(&getBefore(), getBefore().getArguments());
     return;
   }
 
-  assert(*index < 2 && "there are only two regions in a WhileOp");
+  assert(llvm::is_contained({&getAfter(), &getBefore()}, point) &&
+         "there are only two regions in a WhileOp");
   // The body region always branches back to the condition region.
-  if (*index == 1) {
+  if (point == getAfter()) {
     regions.emplace_back(&getBefore(), getBefore().getArguments());
     return;
   }
 
   regions.emplace_back(getResults());
   regions.emplace_back(&getAfter(), getAfter().getArguments());
 }
 

 }
 
 Block &scf::IndexSwitchOp::getCaseBlock(unsigned idx) {
   assert(idx < getNumCases() && "case index out-of-bounds");
   return getCaseRegions()[idx].front();
 }
 
 void IndexSwitchOp::getSuccessorRegions(
-    std::optional<unsigned> index,
-    SmallVectorImpl<RegionSuccessor> &successors) {
+    RegionBranchPoint point, SmallVectorImpl<RegionSuccessor> &successors) {
   // All regions branch back to the parent op.
-  if (index) {
+  if (!point.isParent()) {
     successors.emplace_back(getResults());
     return;
   }
 
   llvm::copy(getRegions(), std::back_inserter(successors));
 }
 
 void IndexSwitchOp::getEntrySuccessorRegions(

mlir/lib/Dialect/Shape/IR/Shape.cpp

 
 void AssumingOp::getCanonicalizationPatterns(RewritePatternSet &patterns,
                                              MLIRContext *context) {
   patterns.add<AssumingOpRemoveUnusedResults, AssumingWithTrue>(context);
 }
 
 // See RegionBranchOpInterface in Interfaces/ControlFlowInterfaces.td
 void AssumingOp::getSuccessorRegions(
-    std::optional<unsigned> index, SmallVectorImpl<RegionSuccessor> &regions) {
+    RegionBranchPoint point, SmallVectorImpl<RegionSuccessor> &regions) {
   // AssumingOp has unconditional control flow into the region and back to the
   // parent, so return the correct RegionSuccessor purely based on the index
   // being None or 0.
-  if (index) {
+  if (!point.isParent()) {
     regions.push_back(RegionSuccessor(getResults()));
     return;
   }
 
   regions.push_back(RegionSuccessor(&getDoRegion()));
 }
 
 void AssumingOp::inlineRegionIntoParent(AssumingOp &op,

mlir/lib/Dialect/Transform/IR/TransformOps.cpp

 
 #define GET_OP_CLASSES
 #include "mlir/Dialect/Transform/IR/TransformOps.cpp.inc"
 
 //===----------------------------------------------------------------------===//
 // AlternativesOp
 //===----------------------------------------------------------------------===//
 
-OperandRange transform::AlternativesOp::getEntrySuccessorOperands(
-    std::optional<unsigned> index) {
-  if (index && getOperation()->getNumOperands() == 1)
+OperandRange
+transform::AlternativesOp::getEntrySuccessorOperands(RegionBranchPoint point) {
+  if (!point.isParent() && getOperation()->getNumOperands() == 1)
     return getOperation()->getOperands();
   return OperandRange(getOperation()->operand_end(),
                       getOperation()->operand_end());
 }
 
 void transform::AlternativesOp::getSuccessorRegions(
-    std::optional<unsigned> index, SmallVectorImpl<RegionSuccessor> &regions) {
+    RegionBranchPoint point, SmallVectorImpl<RegionSuccessor> &regions) {
   for (Region &alternative : llvm::drop_begin(
-           getAlternatives(), index.has_value() ? *index + 1 : 0)) {
+           getAlternatives(),
+           point.isParent() ? 0
+                            : point.getRegionOrNull()->getRegionNumber() + 1)) {
     regions.emplace_back(&alternative, !getOperands().empty()
                                            ? alternative.getArguments()
                                            : Block::BlockArgListType());
   }
-  if (index.has_value())
+  if (!point.isParent())
     regions.emplace_back(getOperation()->getResults());
 }
 
 void transform::AlternativesOp::getRegionInvocationBounds(
     ArrayRef<Attribute> operands, SmallVectorImpl<InvocationBounds> &bounds) {
   (void)operands;
   // The region corresponding to the first alternative is always executed, the
   // remaining may or may not be executed.
if (any_of(getBody().front().without_terminator(), [&](Operation &op) {
     onlyReadsHandle(getTarget(), effects);
   }
 
   for (Value result : getResults())
     producesHandle(result, effects);
 }
 
 void transform::ForeachOp::getSuccessorRegions(
-    std::optional<unsigned> index, SmallVectorImpl<RegionSuccessor> &regions) {
+    RegionBranchPoint point, SmallVectorImpl<RegionSuccessor> &regions) {
   Region *bodyRegion = &getBody();
-  if (!index) {
+  if (point.isParent()) {
     regions.emplace_back(bodyRegion, bodyRegion->getArguments());
     return;
   }
 
   // Branch back to the region or the parent.
-  assert(*index == 0 && "unexpected region index");
+  assert(point == getBody() && "unexpected region index");
   regions.emplace_back(bodyRegion, bodyRegion->getArguments());
   regions.emplace_back();
 }
 
 OperandRange
-transform::ForeachOp::getEntrySuccessorOperands(std::optional<unsigned> index) {
+transform::ForeachOp::getEntrySuccessorOperands(RegionBranchPoint point) {
   // The iteration variable op handle is mapped to a subset (one op to be
   // precise) of the payload ops of the ForeachOp operand.
-  assert(index && *index == 0 && "unexpected region index");
+  assert(point == getBody() && "unexpected region index");
   return getOperation()->getOperands();
 }
 
 transform::YieldOp transform::ForeachOp::getYieldOp() {
   return cast<transform::YieldOp>(getBody().front().getTerminator());
 }
 
 LogicalResult transform::ForeachOp::verify() {
LogicalResult transform::SequenceOp::verify() {
   return success();
 }
 
 void transform::SequenceOp::getEffects(
     SmallVectorImpl<MemoryEffects::EffectInstance> &effects) {
   getPotentialTopLevelEffects(effects);
 }
 
-OperandRange transform::SequenceOp::getEntrySuccessorOperands(
-    std::optional<unsigned> index) {
-  assert(index && *index == 0 && "unexpected region index");
+OperandRange
+transform::SequenceOp::getEntrySuccessorOperands(RegionBranchPoint point) {
+  assert(point == getBody() && "unexpected region index");
   if (getOperation()->getNumOperands() > 0)
     return getOperation()->getOperands();
   return OperandRange(getOperation()->operand_end(),
                       getOperation()->operand_end());
 }
 
 void transform::SequenceOp::getSuccessorRegions(
-    std::optional<unsigned> index, SmallVectorImpl<RegionSuccessor> &regions) {
-  if (!index) {
+    RegionBranchPoint point, SmallVectorImpl<RegionSuccessor> &regions) {
+  if (point.isParent()) {
     Region *bodyRegion = &getBody();
     regions.emplace_back(bodyRegion, getNumOperands() != 0
                                          ? bodyRegion->getArguments()
                                          : Block::BlockArgListType());
     return;
   }
 
-  assert(*index == 0 && "unexpected region index");
+  assert(point == getBody() && "unexpected region index");
   regions.emplace_back(getOperation()->getResults());
 }
 
 void transform::SequenceOp::getRegionInvocationBounds(
     ArrayRef<Attribute> operands, SmallVectorImpl<InvocationBounds> &bounds) {
   (void)operands;
   bounds.emplace_back(1, 1);
 }

mlir/lib/Dialect/Vector/IR/VectorOps.cpp

ParseResult WarpExecuteOnLane0Op::parse(OpAsmParser &parser,
 
   // Parse the optional attribute list.
   if (parser.parseOptionalAttrDict(result.attributes))
     return failure();
   return success();
 }
 
 void WarpExecuteOnLane0Op::getSuccessorRegions(
-    std::optional<unsigned> index, SmallVectorImpl<RegionSuccessor> &regions) {
-  if (index) {
+    RegionBranchPoint point, SmallVectorImpl<RegionSuccessor> &regions) {
+  if (!point.isParent()) {
     regions.push_back(RegionSuccessor(getResults()));
     return;
   }
 
   // The warp region is always executed
   regions.push_back(RegionSuccessor(&getWarpRegion()));
 }
 

mlir/lib/Interfaces/ControlFlowInterfaces.cpp

detail::verifyBranchSuccessorOperands(Operation *op, unsigned succNo,
   }
   return success();
 }
 
 //===----------------------------------------------------------------------===//
 // RegionBranchOpInterface
 //===----------------------------------------------------------------------===//
 
-static InFlightDiagnostic &
-printRegionEdgeName(InFlightDiagnostic &diag, std::optional<unsigned> sourceNo,
-                    std::optional<unsigned> succRegionNo) {
+static InFlightDiagnostic &printRegionEdgeName(InFlightDiagnostic &diag,
+                                               RegionBranchPoint sourceNo,
+                                               RegionBranchPoint succRegionNo) {
   diag << "from ";
-  if (sourceNo)
-    diag << "Region #" << sourceNo.value();
+  if (Region *region = sourceNo.getRegionOrNull())
+    diag << "Region #" << region->getRegionNumber();
   else
     diag << "parent operands";
 
   diag << " to ";
-  if (succRegionNo)
-    diag << "Region #" << succRegionNo.value();
+  if (Region *region = succRegionNo.getRegionOrNull())
+    diag << "Region #" << region->getRegionNumber();
   else
     diag << "parent results";
   return diag;
 }
 
 /// Verify that types match along all region control flow edges originating from
 /// `sourceNo` (region # if source is a region, std::nullopt if source is parent
 /// op). `getInputsTypesForRegion` is a function that returns the types of the
 /// inputs that flow from `sourceIndex' to the given region, or std::nullopt if
 /// the exact type match verification is not necessary (e.g., if the Op verifies
 /// the match itself).
-static LogicalResult verifyTypesAlongAllEdges(
-    Operation *op, std::optional<unsigned> sourceNo,
-    function_ref<FailureOr<TypeRange>(std::optional<unsigned>)>
+static LogicalResult
+verifyTypesAlongAllEdges(Operation *op, RegionBranchPoint sourcePoint,
+                         function_ref<FailureOr<TypeRange>(RegionBranchPoint)>
                              getInputsTypesForRegion) {
   auto regionInterface = cast<RegionBranchOpInterface>(op);
 
   SmallVector<RegionSuccessor, 2> successors;
-  regionInterface.getSuccessorRegions(sourceNo, successors);
+  regionInterface.getSuccessorRegions(sourcePoint, successors);
 
   for (RegionSuccessor &succ : successors) {
-    std::optional<unsigned> succRegionNo;
-    if (!succ.isParent())
-      succRegionNo = succ.getSuccessor()->getRegionNumber();
-
-    FailureOr<TypeRange> sourceTypes = getInputsTypesForRegion(succRegionNo);
+    FailureOr<TypeRange> sourceTypes = getInputsTypesForRegion(succ);
     if (failed(sourceTypes))
       return failure();
 
     TypeRange succInputsTypes = succ.getSuccessorInputs().getTypes();
     if (sourceTypes->size() != succInputsTypes.size()) {
       InFlightDiagnostic diag = op->emitOpError(" region control flow edge ");
-      return printRegionEdgeName(diag, sourceNo, succRegionNo)
+      return printRegionEdgeName(diag, sourcePoint, succ)
              << ": source has " << sourceTypes->size()
              << " operands, but target successor needs "
              << succInputsTypes.size();
     }
 
     for (const auto &typesIdx :
          llvm::enumerate(llvm::zip(*sourceTypes, succInputsTypes))) {
       Type sourceType = std::get<0>(typesIdx.value());
       Type inputType = std::get<1>(typesIdx.value());
       if (!regionInterface.areTypesCompatible(sourceType, inputType)) {
         InFlightDiagnostic diag = op->emitOpError(" along control flow edge ");
-        return printRegionEdgeName(diag, sourceNo, succRegionNo)
+        return printRegionEdgeName(diag, sourcePoint, succ)
                << ": source type #" << typesIdx.index() << " " << sourceType
                << " should match input type #" << typesIdx.index() << " "
                << inputType;
       }
     }
   }
   return success();
 }
 
 /// Verify that types match along control flow edges described the given op.
 LogicalResult detail::verifyTypesAlongControlFlowEdges(Operation *op) {
   auto regionInterface = cast<RegionBranchOpInterface>(op);
 
-  auto inputTypesFromParent =
-      [&](std::optional<unsigned> regionNo) -> TypeRange {
+  auto inputTypesFromParent = [&](RegionBranchPoint regionNo) -> TypeRange {
     return regionInterface.getEntrySuccessorOperands(regionNo).getTypes();
   };
 
   // Verify types along control flow edges originating from the parent.
-  if (failed(verifyTypesAlongAllEdges(op, std::nullopt, inputTypesFromParent)))
+  if (failed(verifyTypesAlongAllEdges(op, RegionBranchPoint::parent(),
+                                      inputTypesFromParent)))
     return failure();
 
   auto areTypesCompatible = [&](TypeRange lhs, TypeRange rhs) {
     if (lhs.size() != rhs.size())
       return false;
     for (auto types : llvm::zip(lhs, rhs)) {
       if (!regionInterface.areTypesCompatible(std::get<0>(types),
                                               std::get<1>(types))) {
         return false;
       }
     }
     return true;
   };
 
   // Verify types along control flow edges originating from each region.
-  for (unsigned regionNo : llvm::seq(0U, op->getNumRegions())) {
-    Region &region = op->getRegion(regionNo);
+  for (Region &region : op->getRegions()) {
 
     // Since there can be multiple terminators implementing the
     // `RegionBranchTerminatorOpInterface`, all should have the same operand
     // types when passing them to the same region.
 
     SmallVector<RegionBranchTerminatorOpInterface> regionReturnOps;
     for (Block &block : region)
       if (auto terminator = dyn_cast<RegionBranchTerminatorOpInterface>(
               block.getTerminator()))
         regionReturnOps.push_back(terminator);
 
     // If there is no return-like terminator, the op itself should verify
     // type consistency.
     if (regionReturnOps.empty())
       continue;
 
     auto inputTypesForRegion =
-        [&](std::optional<unsigned> succRegionNo) -> FailureOr<TypeRange> {
+        [&](RegionBranchPoint succRegionNo) -> FailureOr<TypeRange> {
       std::optional<OperandRange> regionReturnOperands;
       for (RegionBranchTerminatorOpInterface regionReturnOp : regionReturnOps) {
         auto terminatorOperands =
             regionReturnOp.getSuccessorOperands(succRegionNo);
 
         if (!regionReturnOperands) {
           regionReturnOperands = terminatorOperands;
           continue;
         }
 
         // Found more than one ReturnLike terminator. Make sure the operand
         // types match with the first one.
         if (!areTypesCompatible(regionReturnOperands->getTypes(),
                                 terminatorOperands.getTypes())) {
           InFlightDiagnostic diag = op->emitOpError(" along control flow edge");
-          return printRegionEdgeName(diag, regionNo, succRegionNo)
+          return printRegionEdgeName(diag, region, succRegionNo)
                  << " operands mismatch between return-like terminators";
         }
       }
 
       // All successors get the same set of operand types.
       return TypeRange(regionReturnOperands->getTypes());
     };
 
-    if (failed(verifyTypesAlongAllEdges(op, regionNo, inputTypesForRegion)))
+    if (failed(verifyTypesAlongAllEdges(op, region, inputTypesForRegion)))
       return failure();
   }
 
   return success();
 }
 
 /// Return `true` if region `r` is reachable from region `begin` according to
 /// the RegionBranchOpInterface (by taking a branch).
 static bool isRegionReachable(Region *begin, Region *r) {
   assert(begin->getParentOp() == r->getParentOp() &&
          "expected that both regions belong to the same op");
   auto op = cast<RegionBranchOpInterface>(begin->getParentOp());
   SmallVector<bool> visited(op->getNumRegions(), false);
   visited[begin->getRegionNumber()] = true;
 
   // Retrieve all successors of the region and enqueue them in the worklist.
-  SmallVector<unsigned> worklist;
-  auto enqueueAllSuccessors = [&](unsigned index) {
+  SmallVector<Region *> worklist;
+  auto enqueueAllSuccessors = [&](Region *region) {
     SmallVector<RegionSuccessor> successors;
-    op.getSuccessorRegions(index, successors);
+    op.getSuccessorRegions(region, successors);
     for (RegionSuccessor successor : successors)
       if (!successor.isParent())
-        worklist.push_back(successor.getSuccessor()->getRegionNumber());
+        worklist.push_back(successor.getSuccessor());
   };
-  enqueueAllSuccessors(begin->getRegionNumber());
+  enqueueAllSuccessors(begin);
 
   // Process all regions in the worklist via DFS.
   while (!worklist.empty()) {
-    unsigned nextRegion = worklist.pop_back_val();
-    if (nextRegion == r->getRegionNumber())
+    Region *nextRegion = worklist.pop_back_val();
+    if (nextRegion == r)
       return true;
-    if (visited[nextRegion])
+    if (visited[nextRegion->getRegionNumber()])
       continue;
-    visited[nextRegion] = true;
+    visited[nextRegion->getRegionNumber()] = true;
     enqueueAllSuccessors(nextRegion);
   }
 
   return false;
 }
 
 /// Return `true` if `a` and `b` are in mutually exclusive regions.
 ///

mlir/lib/Transforms/RemoveDeadValues.cpp

for (Region &region : regionBranchOp->getRegions()) {
       BitVector regionLiveArgs = markLives(arguments, la);
       liveArgs[&region] = regionLiveArgs;
     }
   };
 
   // Return the successors of `region` if the latter is not null. Else return
   // the successors of `regionBranchOp`.
   auto getSuccessors = [&](Region *region = nullptr) {
-    std::optional<unsigned> index =
-        region ? std::optional(region->getRegionNumber()) : std::nullopt;
+    auto point = region ? region : RegionBranchPoint::parent();
     SmallVector<Attribute> operandAttributes(regionBranchOp->getNumOperands(),
                                              nullptr);
     SmallVector<RegionSuccessor> successors;
-    if (!index)
-      regionBranchOp.getEntrySuccessorRegions(operandAttributes, successors);
-    else
-      regionBranchOp.getSuccessorRegions(index, successors);
+    regionBranchOp.getSuccessorRegions(point, successors);
     return successors;
   };
 
   // Return the operands of `terminator` that are forwarded to `successor` if
   // the former is not null. Else return the operands of `regionBranchOp`
   // forwarded to `successor`.
   auto getForwardedOpOperands = [&](const RegionSuccessor &successor,
                                     Operation *terminator = nullptr) {
-    Region *successorRegion = successor.getSuccessor();
-    std::optional<unsigned> index =
-        successorRegion ? std::optional(successorRegion->getRegionNumber())
-                        : std::nullopt;
     OperandRange operands =
         terminator ? cast<RegionBranchTerminatorOpInterface>(terminator)
-                         .getSuccessorOperands(index)
-                   : regionBranchOp.getEntrySuccessorOperands(index);
+                         .getSuccessorOperands(successor)
+                   : regionBranchOp.getEntrySuccessorOperands(successor);
     SmallVector<OpOperand *> opOperands = operandsToOpOperands(operands);
     return opOperands;
   };
 
   // Mark the non-forwarded operands of `regionBranchOp` in
   // `nonForwardedOperands`.
   auto markNonForwardedOperands = [&](BitVector &nonForwardedOperands) {
     nonForwardedOperands.resize(regionBranchOp->getNumOperands(), true);

mlir/test/lib/Analysis/DataFlow/TestDenseBackwardDataFlowAnalysis.cpp

void visitOperation(Operation *op, const NextAccess &after,
                       NextAccess *before) override;
 
   void visitCallControlFlowTransfer(CallOpInterface call,
                                     CallControlFlowAction action,
                                     const NextAccess &after,
                                     NextAccess *before) override;
 
   void visitRegionBranchControlFlowTransfer(RegionBranchOpInterface branch,
-                                            std::optional<unsigned> regionFrom,
-                                            std::optional<unsigned> regionTo,
+                                            RegionBranchPoint regionFrom,
+                                            RegionBranchPoint regionTo,
                                             const NextAccess &after,
                                             NextAccess *before) override;
 
   // TODO: this isn't ideal for the analysis. When there is no next access, it
   // means "we don't know what the next access is" rather than "there is no next
   // access". But it's unclear how to differentiate the two cases...
   void setToExitState(NextAccess *lattice) override {
     propagateIfChanged(lattice, lattice->reset());
if (testCallAndStore && ((action == CallControlFlowAction::EnterCallee &&
     visitOperation(call, after, before);
   } else {
     AbstractDenseBackwardDataFlowAnalysis::visitCallControlFlowTransfer(
         call, action, after, before);
   }
 }
 
 void NextAccessAnalysis::visitRegionBranchControlFlowTransfer(
-    RegionBranchOpInterface branch, std::optional<unsigned> regionFrom,
-    std::optional<unsigned> regionTo, const NextAccess &after,
-    NextAccess *before) {
+    RegionBranchOpInterface branch, RegionBranchPoint regionFrom,
+    RegionBranchPoint regionTo, const NextAccess &after, NextAccess *before) {
   auto testStoreWithARegion =
       dyn_cast<::test::TestStoreWithARegion>(branch.getOperation());
 
   if (testStoreWithARegion &&
-      ((!regionTo && !testStoreWithARegion.getStoreBeforeRegion()) ||
-       (!regionFrom && testStoreWithARegion.getStoreBeforeRegion()))) {
+      ((regionTo.isParent() && !testStoreWithARegion.getStoreBeforeRegion()) ||
+       (regionFrom.isParent() &&
+        testStoreWithARegion.getStoreBeforeRegion()))) {
     visitOperation(branch, static_cast<const NextAccess &>(after),
                    static_cast<NextAccess *>(before));
   } else {
     propagateIfChanged(before, before->meet(after));
   }
 }
 
 namespace {
op->walk([&](Operation *op) {
                   makeNextAccessAttribute(op, solver, nextAccess));
 
       auto iface = dyn_cast<RegionBranchOpInterface>(op);
       if (!iface)
         return;
 
       SmallVector<Attribute> entryPointNextAccess;
       SmallVector<RegionSuccessor> regionSuccessors;
-      iface.getSuccessorRegions(std::nullopt, regionSuccessors);
+      iface.getSuccessorRegions(RegionBranchPoint::parent(), regionSuccessors);
       for (const RegionSuccessor &successor : regionSuccessors) {
         if (!successor.getSuccessor() || successor.getSuccessor()->empty())
           continue;
         Block &successorBlock = successor.getSuccessor()->front();
         ProgramPoint successorPoint = successorBlock.empty()
                                           ? ProgramPoint(&successorBlock)
                                           : &successorBlock.front();
         entryPointNextAccess.push_back(makeNextAccessAttribute(

mlir/test/lib/Dialect/Test/TestDialect.cpp

if (parser.parseKeyword("then") || parser.parseRegion(*thenRegion, {}, {}) ||
       parser.parseKeyword("else") || parser.parseRegion(*elseRegion, {}, {}) ||
       parser.parseKeyword("join") || parser.parseRegion(*joinRegion, {}, {}))
     return failure();
 
   return parser.resolveOperands(operandInfos, operandTypes,
                                 parser.getCurrentLocation(), result.operands);
 }
 
-OperandRange
-RegionIfOp::getEntrySuccessorOperands(std::optional<unsigned> index) {
-  assert(index && *index < 2 && "invalid region index");
+OperandRange RegionIfOp::getEntrySuccessorOperands(RegionBranchPoint point) {
+  assert(llvm::is_contained({&getThenRegion(), &getElseRegion()}, point) &&
+         "invalid region index");
   return getOperands();
 }
 
 void RegionIfOp::getSuccessorRegions(
-    std::optional<unsigned> index, SmallVectorImpl<RegionSuccessor> &regions) {
+    RegionBranchPoint point, SmallVectorImpl<RegionSuccessor> &regions) {
   // We always branch to the join region.
-  if (index.has_value()) {
-    if (index.value() < 2)
+  if (!point.isParent()) {
+    if (point != getJoinRegion())
       regions.push_back(RegionSuccessor(&getJoinRegion(), getJoinArgs()));
     else
       regions.push_back(RegionSuccessor(getResults()));
     return;
   }
 
   // The then and else regions are the entry regions of this op.
   regions.push_back(RegionSuccessor(&getThenRegion(), getThenArgs()));
   regions.push_back(RegionSuccessor(&getElseRegion(), getElseArgs()));
 }
 
 void RegionIfOp::getRegionInvocationBounds(
     ArrayRef<Attribute> operands,
     SmallVectorImpl<InvocationBounds> &invocationBounds) {
   // Each region is invoked at most once.
   invocationBounds.assign(/*NumElts=*/3, /*Elt=*/{0, 1});
 }
 
 //===----------------------------------------------------------------------===//
 // AnyCondOp
 //===----------------------------------------------------------------------===//
 
-void AnyCondOp::getSuccessorRegions(std::optional<unsigned> index,
+void AnyCondOp::getSuccessorRegions(RegionBranchPoint point,
                                     SmallVectorImpl<RegionSuccessor> &regions) {
   // The parent op branches into the only region, and the region branches back
   // to the parent op.
-  if (!index)
+  if (point.isParent())
     regions.emplace_back(&getRegion());
   else
     regions.emplace_back(getResults());
 }
 
 void AnyCondOp::getRegionInvocationBounds(
     ArrayRef<Attribute> operands,
     SmallVectorImpl<InvocationBounds> &invocationBounds) {
   invocationBounds.emplace_back(1, 1);
 }
 
 //===----------------------------------------------------------------------===//
 // LoopBlockOp
 //===----------------------------------------------------------------------===//
 
 void LoopBlockOp::getSuccessorRegions(
-    std::optional<unsigned> index, SmallVectorImpl<RegionSuccessor> &regions) {
+    RegionBranchPoint point, SmallVectorImpl<RegionSuccessor> &regions) {
   regions.emplace_back(&getBody(), getBody().getArguments());
-  if (!index)
+  if (point.isParent())
     return;
 
   regions.emplace_back((*this)->getResults());
 }
 
-OperandRange
-LoopBlockOp::getEntrySuccessorOperands(std::optional<unsigned> index) {
-  assert(index == 0);
+OperandRange LoopBlockOp::getEntrySuccessorOperands(RegionBranchPoint point) {
+  assert(point == getBody());
   return getInitMutable();
 }
 
 //===----------------------------------------------------------------------===//
 // LoopBlockTerminatorOp
 //===----------------------------------------------------------------------===//
 
-MutableOperandRange LoopBlockTerminatorOp::getMutableSuccessorOperands(
-    std::optional<unsigned> index) {
-  assert(!index || index == 0);
-  if (!index)
+MutableOperandRange
+LoopBlockTerminatorOp::getMutableSuccessorOperands(RegionBranchPoint point) {
+  if (point.isParent())
     return getExitArgMutable();
   return getNextIterArgMutable();
 }
 
 //===----------------------------------------------------------------------===//
 // SingleNoTerminatorCustomAsmOp
 //===----------------------------------------------------------------------===//
 
Operation::operand_range TestCallOnDeviceOp::getArgOperands() {
   return getForwardedOperands();
 }
 
 MutableOperandRange TestCallOnDeviceOp::getArgOperandsMutable() {
   return getForwardedOperandsMutable();
 }
 
 void TestStoreWithARegion::getSuccessorRegions(
-    std::optional<unsigned> index, SmallVectorImpl<RegionSuccessor> &regions) {
-  if (!index) {
+    RegionBranchPoint point, SmallVectorImpl<RegionSuccessor> &regions) {
+  if (point.isParent())
     regions.emplace_back(&getBody(), getBody().front().getArguments());
-  } else {
+  else
     regions.emplace_back();
-  }
-}
 
 LogicalResult
 TestVersionedOpA::readProperties(::mlir::DialectBytecodeReader &reader,
                                  ::mlir::OperationState &state) {
   auto &prop = state.getOrAddProperties<Properties>();
   if (::mlir::failed(reader.readAttribute(prop.dims)))
     return ::mlir::failure();
 

mlir/unittests/Interfaces/ControlFlowInterfacesTest.cpp

struct MutuallyExclusiveRegionsOp
   using Op::Op;
   static ArrayRef<StringRef> getAttributeNames() { return {}; }
 
   static StringRef getOperationName() {
     return "cftest.mutually_exclusive_regions_op";
   }
 
   // Regions have no successors.
-  void getSuccessorRegions(std::optional<unsigned> index,
+  void getSuccessorRegions(RegionBranchPoint point,
                            SmallVectorImpl<RegionSuccessor> &regions) {}
 };
 
 /// All regions of this op call each other in a large circle.
 struct LoopRegionsOp
     : public Op<LoopRegionsOp, RegionBranchOpInterface::Trait> {
   using Op::Op;
   static const unsigned kNumRegions = 3;
 
   static ArrayRef<StringRef> getAttributeNames() { return {}; }
 
   static StringRef getOperationName() { return "cftest.loop_regions_op"; }
 
-  void getSuccessorRegions(std::optional<unsigned> index,
+  void getSuccessorRegions(RegionBranchPoint point,
                            SmallVectorImpl<RegionSuccessor> &regions) {
-    if (index) {
-      if (*index == 1)
+    if (Region *region = point.getRegionOrNull()) {
+      if (point == (*this)->getRegion(1))
         // This region also branches back to the parent.
         regions.push_back(RegionSuccessor());
-      regions.push_back(
-          RegionSuccessor(&getOperation()->getRegion(*index % kNumRegions)));
+      regions.push_back(RegionSuccessor(region));
     }
   }
 };
 
 /// Each region branches back it itself or the parent.
 struct DoubleLoopRegionsOp
     : public Op<DoubleLoopRegionsOp, RegionBranchOpInterface::Trait> {
   using Op::Op;
 
   static ArrayRef<StringRef> getAttributeNames() { return {}; }
 
   static StringRef getOperationName() {
     return "cftest.double_loop_regions_op";
   }
 
-  void getSuccessorRegions(std::optional<unsigned> index,
+  void getSuccessorRegions(RegionBranchPoint point,
                            SmallVectorImpl<RegionSuccessor> &regions) {
-    if (index.has_value()) {
+    if (Region *region = point.getRegionOrNull()) {
       regions.push_back(RegionSuccessor());
-      regions.push_back(RegionSuccessor(&getOperation()->getRegion(*index)));
+      regions.push_back(RegionSuccessor(region));
     }
   }
 };
 
 /// Regions are executed sequentially.
 struct SequentialRegionsOp
     : public Op<SequentialRegionsOp, RegionBranchOpInterface::Trait> {
   using Op::Op;
   static ArrayRef<StringRef> getAttributeNames() { return {}; }
 
   static StringRef getOperationName() { return "cftest.sequential_regions_op"; }
 
   // Region 0 has Region 1 as a successor.
-  void getSuccessorRegions(std::optional<unsigned> index,
+  void getSuccessorRegions(RegionBranchPoint point,
                            SmallVectorImpl<RegionSuccessor> &regions) {
-    if (index == 0u) {
+    if (point == (*this)->getRegion(0)) {
       Operation *thisOp = this->getOperation();
       regions.push_back(RegionSuccessor(&thisOp->getRegion(1)));
     }
   }
 };
 
 /// A dialect putting all the above together.
 struct CFTestDialect : Dialect {