[MLIR][Analysis] Fix incorrect RegionBranchOpInterface API usage in SliceWalk

Project: llvm/llvm-project | Language: C++ | Concepts: api-design, error-handling, testing-strategy | Difficulty: mid
Link: https://github.com/llvm/llvm-project/pull/188758

The problem

The getControlFlowPredecessors function in MLIR's analysis library was incorrectly passing the result or block argument index directly to getPredecessorValues. However, getPredecessorValues expects an index into the list of successor inputs (as returned by getSuccessorInputs), not the raw result/block argument list. If an operation has results or block arguments that are not successor inputs (for example, an extra loop counter result in an scf.for-like op), this could lead to out-of-bounds accesses or incorrect results.

Why it's non-trivial

A naive fix would simply use the result or argument number as the index, but this assumes a one-to-one correspondence between results/block arguments and successor inputs. In reality, some results or block arguments may not be successor inputs at all, so the mapping is not always direct. Failing to account for this can cause subtle bugs or crashes.

The approach

1. Properly map the value to its corresponding successor input index (if any):

Instead of passing the result or argument number directly, the code now looks up the value in the successor inputs and uses its position as the index. If the value is not a successor input, it returns std::nullopt.

+    ValueRange successorInputs = regionOp.getSuccessorInputs(parentSuccessor);
+    auto it = llvm::find(successorInputs, opResult);
+    if (it == successorInputs.end())
+      return std::nullopt;

This snippet finds the correct index of the result value (opResult) in the list of successor inputs. If the value is not found, it returns std::nullopt, indicating there are no control-flow predecessors.

2. Use the correct index when calling getPredecessorValues:

-    regionOp.getPredecessorValues(region, opResult.getResultNumber(),
-                                  predecessorOperands);
+    regionOp.getPredecessorValues(parentSuccessor,
+                                  std::distance(successorInputs.begin(), it),
+                                  predecessorOperands);

Now, the index passed to getPredecessorValues is computed as the position of the value in the successor inputs, not just its result number.

3. Apply the same fix for block arguments:

+      ValueRange successorInputs =
+          regionBranchOp.getSuccessorInputs(regionSuccessor);
+      auto it = llvm::find(successorInputs, blockArg);
+      if (it == successorInputs.end())
+        return std::nullopt;

This snippet does the same for block arguments, ensuring only block arguments that are actually successor inputs are considered.

4. Testing with new cases that would previously fail or return wrong results:

A new test file was added to cover scenarios where values are not successor inputs:

// Test case 1: scf.for with no iter_args.
// CHECK-LABEL: Control flow predecessors for '"test_scf_for_no_iter_args"'
// CHECK:   'scf.for': region #0 block arg #0: no predecessors

// Test case 2: test.loop_with_extra_result.
// CHECK-LABEL: Control flow predecessors for '"test_loop_with_extra_result"'
// CHECK:   'test.loop_with_extra_result': result #0: no predecessors
// CHECK:   'test.loop_with_extra_result': result #1: 1 predecessor(s)
// CHECK:   'test.loop_with_extra_result': region #0 block arg #0: no predecessors
// CHECK:   'test.loop_with_extra_result': region #0 block arg #1: 2 predecessor(s)

These tests ensure that only actual successor inputs are considered to have control-flow predecessors, and that the fix prevents out-of-bounds errors.

From the review thread

No notable review discussion on this PR.

Takeaways

• APIs that use indices must clearly define what the index refers to; never assume that ordinal positions in different lists (e.g., results vs. successor inputs) are interchangeable.
• When mapping between related but non-identical collections (like results and successor inputs), always search for the actual value rather than assuming positional correspondence.
• Comprehensive tests should cover edge cases where values are intentionally excluded from certain relationships (e.g., not all results are successor inputs), to prevent silent errors or crashes.