[llvm-branch-commits] [mlir] [mlir] Add `OpFoldResults` definition (PR #229215)

[Mehdi Amini via llvm-branch-commits]

================
@@ -54,6 +61,101 @@ struct CastInfo<To, const mlir::OpFoldResult>
 namespace mlir {
 /// Allow printing to a stream.
 raw_ostream &operator<<(raw_ostream &os, OpFoldResult ofr);
+
+/// The result of a fold of any op. Replacement i is an Attribute (replace
+/// result i with a constant), a Value (replace result i with that value), or
+/// null / op->getResult(i) (keep result i). A separate bit records an in-place
+/// change of the op.
+class [[nodiscard]] OpFoldResults {
+public:
+  /// Failure: the fold did not apply and the IR is unchanged.
+  OpFoldResults() = default;
+  /// Failure, the same as the default constructor.
+  OpFoldResults(std::nullptr_t);
+  /// success(): the op changed in place. failure(): failure.
+  OpFoldResults(LogicalResult status);
+  /// One replacement. Valid only when the op has exactly one result at run
+  /// time. A null replacement, or the op's own result, keeps the result, so the
+  /// fold fails. Unlike `OpFoldResult fold`, the op's own result does not mean
+  /// in place; use success() or setModifiedInPlace() for an in-place change.
+  OpFoldResults(OpFoldResult replacement);
+  OpFoldResults(Value replacement);
+  OpFoldResults(Attribute replacement);
+  /// One replacement per result.
+  OpFoldResults(std::initializer_list<OpFoldResult> replacements);
+  /// One replacement per range element. An empty range is a failure.
+  template <typename RangeT,
+            typename = std::enable_if_t<
+                !std::is_convertible_v<RangeT, Attribute> &&
+                !std::is_convertible_v<RangeT, Value> &&
+                std::is_convertible_v<llvm::detail::ValueOfRange<RangeT>,
+                                      OpFoldResult>>>
+  OpFoldResults(RangeT &&range) {
+    llvm::append_range(replacements, range);
+  }
+  /// Incremental form: every result is kept and the op is not changed in
+  /// place.
+  explicit OpFoldResults(Operation *op);
+
+  /// Set the replacement of `result`, which must be a result of the op. If
+  /// another op owns `result`, or if normalize() already ran, the behavior is
+  /// undefined. The object must already hold one replacement per result, for
+  /// example from OpFoldResults(op). A null replacement, or
+  /// `replacement == result`, keeps the result. The last write wins.
+  void replace(Value result, OpFoldResult replacement);
+  /// Same as above, with the index of the result. normalize() maps a
+  /// replacement that is the op's own result to "keep".
+  void replace(unsigned resultIndex, OpFoldResult replacement);
+  /// Set whether the fold changed the op in place (operands, attributes,
+  /// properties, or regions).
+  void setModifiedInPlace(bool modified = true);
+
+  // Queries for drivers. They are valid after normalize().
+
+  /// Return true if the fold changed the op in place or replaces a result.
+  bool succeeded() const;
+  /// Return true if the fold did not apply.
+  bool failed() const;
+  /// Return true if the fold changed the op in place.
+  bool modifiedInPlace() const;
+  /// Return true if the fold replaces at least one result.
+  bool replacesAny() const;
+  /// Return true if the operation has at least one result and all the results
+  /// are replaced.
+  bool replacesAll() const;
+  /// Return the number of replacements: 0 if no replacement, or one per result
+  /// of the op.
+  unsigned size() const;
+  /// Return the replacement of result `i`; null means keep. If size() == 0,
+  /// every index reads as keep; otherwise `i` must be less than size().
+  OpFoldResult operator[](unsigned i) const;
+  /// Return the replacements. The returned range points into this object.
+  ArrayRef<OpFoldResult> getReplacements() const LLVM_LIFETIME_BOUND;
+
+  /// If replacement i matches op's result i, set replacement i to null. Clear
+  /// replacements if all end up null. Idempotent.
+  void normalize(Operation *op);
----------------
joker-eph wrote:

I don't understand the need for this API.  Right now, this class looks a bit complex, with invariants that seem hard to track.

Can we simplify this all to make it easier to use?

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