[mlir] Allow passing AsmState when printing Attributes and Types

This allows for extracting assembly information when printing an attribute
or type, such as the dialect resources referenced. This functionality is used in
a followup that adds resource support to the bytecode. This change also results
in a nice cleanup of AsmPrinter now that we don't need to awkwardly workaround
optional AsmStates.

Differential Revision: https://reviews.llvm.org/D132728

mlir/include/mlir/IR/AsmState.h

@@ -21,6 +21,7 @@
 
 namespace mlir {
 class AsmResourcePrinter;
+class AsmDialectResourceHandle;
 class Operation;
 
 namespace detail {
@@ -455,6 +456,9 @@ class AsmState {
   AsmState(Operation *op,
            const OpPrintingFlags &printerFlags = OpPrintingFlags(),
            LocationMap *locationMap = nullptr);
+  AsmState(MLIRContext *ctx,
+           const OpPrintingFlags &printerFlags = OpPrintingFlags(),
+           LocationMap *locationMap = nullptr);
   ~AsmState();
 
   /// Get the printer flags.
@@ -480,6 +484,11 @@ class AsmState {
         name, std::forward<CallableT>(printFn)));
   }
 
+  /// Returns a map of dialect resources that were referenced when using this
+  /// state to print IR.
+  DenseMap<Dialect *, SetVector<AsmDialectResourceHandle>> &
+  getDialectResources() const;
+
 private:
   AsmState() = delete;
 

mlir/include/mlir/IR/Attributes.h

@@ -13,6 +13,7 @@
 #include "llvm/Support/PointerLikeTypeTraits.h"
 
 namespace mlir {
+class AsmState;
 class StringAttr;
 
 /// Attributes are known-constant values of operations.
@@ -76,6 +77,7 @@ class Attribute {
 
   /// Print the attribute.
   void print(raw_ostream &os) const;
+  void print(raw_ostream &os, AsmState &state) const;
   void dump() const;
 
   /// Get an opaque pointer to the attribute.

mlir/include/mlir/IR/Types.h

@@ -15,6 +15,8 @@
 #include "llvm/Support/PointerLikeTypeTraits.h"
 
 namespace mlir {
+class AsmState;
+
 /// Instances of the Type class are uniqued, have an immutable identifier and an
 /// optional mutable component.  They wrap a pointer to the storage object owned
 /// by MLIRContext.  Therefore, instances of Type are passed around by value.
@@ -162,6 +164,7 @@ class Type {
 
   /// Print the current type.
   void print(raw_ostream &os) const;
+  void print(raw_ostream &os, AsmState &state) const;
   void dump() const;
 
   friend ::llvm::hash_code hash_value(Type arg);