NFC: Add several utilities to OpState.

* 'walk' functionality.
* Method to set the location.
* 'dump'/'print' methods.
* Method to set the attributes.

PiperOrigin-RevId: 256097960

mlir/include/mlir/IR/OpDefinition.h

@@ -71,11 +71,18 @@ public:
   /// Return the operation that this refers to.
   Operation *getOperation() { return state; }
 
-  ///  Return the context this operation belongs to.
+  /// Return the context this operation belongs to.
   MLIRContext *getContext() { return getOperation()->getContext(); }
 
+  /// Print the operation to the given stream.
+  void print(raw_ostream &os) { state->print(os); }
+
+  /// Dump this operation.
+  void dump() { state->dump(); }
+
   /// The source location the operation was defined or derived from.
   Location getLoc() { return state->getLoc(); }
+  void setLoc(Location loc) { state->setLoc(loc); }
 
   /// Return all of the attributes on this operation.
   ArrayRef<NamedAttribute> getAttrs() { return state->getAttrs(); }
@@ -97,6 +104,12 @@ public:
     setAttr(Identifier::get(name, getContext()), value);
   }
 
+  /// Set the attributes held by this operation.
+  void setAttrs(ArrayRef<NamedAttribute> attributes) {
+    state->setAttrs(attributes);
+  }
+  void setAttrs(NamedAttributeList newAttrs) { state->setAttrs(newAttrs); }
+
   /// Remove the attribute with the specified name if it exists.  The return
   /// value indicates whether the attribute was present or not.
   NamedAttributeList::RemoveResult removeAttr(Identifier name) {
@@ -128,6 +141,20 @@ public:
   /// handlers that may be listening.
   InFlightDiagnostic emitRemark(const Twine &message = {});
 
+  /// Walk the operation in postorder, calling the callback for each nested
+  /// operation(including this one).
+  void walk(const std::function<void(Operation *)> &callback) {
+    state->walk(callback);
+  }
+
+  /// Specialization of walk to only visit operations of 'OpTy'.
+  template <typename OpTy> void walk(std::function<void(OpTy)> callback) {
+    walk([&](Operation *opInst) {
+      if (auto op = dyn_cast<OpTy>(opInst))
+        callback(op);
+    });
+  }
+
   // These are default implementations of customization hooks.
 public:
   /// This hook returns any canonicalization pattern rewrites that the operation
@@ -797,8 +824,6 @@ public:
     return BaseProperties<Traits<ConcreteType>...>::getTraitProperties();
   }
 
-  // TODO: Provide a dump() method.
-
   /// Expose the type we are instantiated on to template machinery that may want
   /// to introspect traits on this operation.
   using ConcreteOpType = ConcreteType;

mlir/include/mlir/IR/Operation.h

@@ -266,7 +266,7 @@ public:
   /// Set the attribute list on this operation.
   /// Using a NamedAttributeList is more efficient as it does not require new
   /// uniquing in the MLIRContext.
-  void setAttrList(NamedAttributeList newAttrs) { attrs = newAttrs; }
+  void setAttrs(NamedAttributeList newAttrs) { attrs = newAttrs; }
 
   /// Return the specified attribute if present, null otherwise.
   Attribute getAttr(Identifier name) { return attrs.get(name); }