[Mlir-commits] [mlir] [MLIR] Support interrupting AffineExpr walks (PR #74792)

[Uday Bondhugula]

================
@@ -123,8 +124,19 @@ class AffineExpr {
   /// Return true if the affine expression involves AffineSymbolExpr `position`.
   bool isFunctionOfSymbol(unsigned position) const;
 
-  /// Walk all of the AffineExpr's in this expression in postorder.
-  void walk(std::function<void(AffineExpr)> callback) const;
+  /// Walk all of the AffineExpr's in this expression in postorder. This allows
+  /// a lambda walk function that can either return `void` or a WalkResult. With
+  /// a WalkResult, interrupting is supported.
+  template <typename FnT, typename RetT = detail::walkResultType<FnT>>
+  std::enable_if_t<std::is_same<RetT, void>::value, RetT>
+  walk(FnT &&callback) const {
+    return walk<void>(*this, callback);
+  }
+  template <typename FnT, typename RetT = detail::walkResultType<FnT>>
+  std::enable_if_t<std::is_same<RetT, WalkResult>::value, RetT>
+  walk(FnT &&callback) const {
+    return walk<WalkResult>(*this, callback);
+  }
----------------
bondhugula wrote:

Thanks - I overlooked that when I created that trampoline. std::enable_if isn't needed here.

https://github.com/llvm/llvm-project/pull/74792