[mlir] Expose printer flags in AsmState

[sgrechanik]

Summary

This change exposes printer flags in AsmState and AsmStateImpl. All functions
receiving AsmState as a parameter now use the flags from the AsmState instead of
taking an additional OpPrintingFlags parameter.

[rriddle]

Nice!

mlir/include/mlir/IR/AsmState.h

public:
 
   /// Initialize the asm state at the level of the given operation. A location
   /// map may optionally be provided to be populated when printing.
   AsmState(Operation *op,
            const OpPrintingFlags &printerFlags = OpPrintingFlags(),
            LocationMap *locationMap = nullptr);
   ~AsmState();
 
+  /// Get the printer flags.
+  const OpPrintingFlags &getPrinterFlags() const;
+
   /// Return an instance of the internal implementation. Returns nullptr if the
   /// state has not been initialized.
   detail::AsmStateImpl &getImpl() { return *impl; }
 
 private:
   AsmState() = delete;
 
   /// A pointer to allocated storage for the impl state.

mlir/include/mlir/IR/OpDefinition.h

public:
 
   /// Return the context this operation belongs to.
   MLIRContext *getContext() { return getOperation()->getContext(); }
 
   /// Print the operation to the given stream.
   void print(raw_ostream &os, OpPrintingFlags flags = llvm::None) {
     state->print(os, flags);
   }
-  void print(raw_ostream &os, AsmState &asmState,
-             OpPrintingFlags flags = llvm::None) {
-    state->print(os, asmState, flags);
+  void print(raw_ostream &os, AsmState &asmState) {
+    state->print(os, asmState);
   }
 
   /// Dump this operation.
   void dump() { state->dump(); }
 
   /// The source location the operation was defined or derived from.
   Location getLoc() { return state->getLoc(); }
 

mlir/include/mlir/IR/Operation.h

public:
   /// Given an operation 'other' that is within the same parent block, return
   /// whether the current operation is before 'other' in the operation list
   /// of the parent block.
   /// Note: This function has an average complexity of O(1), but worst case may
   /// take O(N) where N is the number of operations within the parent block.
   bool isBeforeInBlock(Operation *other);
 
   void print(raw_ostream &os, const OpPrintingFlags &flags = llvm::None);
-  void print(raw_ostream &os, AsmState &state,
-             const OpPrintingFlags &flags = llvm::None);
+  void print(raw_ostream &os, AsmState &state);
   void dump();
 
   //===--------------------------------------------------------------------===//
   // Operands
   //===--------------------------------------------------------------------===//
 
   /// Replace the current operands of this operation with the ones provided in
   /// 'operands'.

mlir/lib/IR/AsmPrinter.cpp

public:
   }
 
   /// Get the state used for aliases.
   AliasState &getAliasState() { return aliasState; }
 
   /// Get the state used for SSA names.
   SSANameState &getSSANameState() { return nameState; }
 
+  /// Get the printer flags.
+  const OpPrintingFlags &getPrinterFlags() const { return printerFlags; }
+
   /// Register the location, line and column, within the buffer that the given
   /// operation was printed at.
   void registerOperationLocation(Operation *op, unsigned line, unsigned col) {
     if (locationMap)
       (*locationMap)[op] = std::make_pair(line, col);
   }
 
 private:

 } // namespace detail
 } // namespace mlir
 
 AsmState::AsmState(Operation *op, const OpPrintingFlags &printerFlags,
                    LocationMap *locationMap)
     : impl(std::make_unique<AsmStateImpl>(op, printerFlags, locationMap)) {}
 AsmState::~AsmState() = default;
 
+const OpPrintingFlags &AsmState::getPrinterFlags() const {
+  return impl->getPrinterFlags();
+}
+
 //===----------------------------------------------------------------------===//
 // AsmPrinter::Impl
 //===----------------------------------------------------------------------===//
 
 namespace mlir {
 class AsmPrinter::Impl {
 public:
   Impl(raw_ostream &os, OpPrintingFlags flags = llvm::None,

 
 namespace {
 /// This class contains the logic for printing operations, regions, and blocks.
 class OperationPrinter : public AsmPrinter::Impl, private OpAsmPrinter {
 public:
   using Impl = AsmPrinter::Impl;
   using Impl::printType;
 
-  explicit OperationPrinter(raw_ostream &os, OpPrintingFlags flags,
-                            AsmStateImpl &state)
-      : Impl(os, flags, &state), OpAsmPrinter(static_cast<Impl &>(*this)) {}
+  explicit OperationPrinter(raw_ostream &os, AsmStateImpl &state)
+      : Impl(os, state.getPrinterFlags(), &state),
+        OpAsmPrinter(static_cast<Impl &>(*this)) {}
 
   /// Print the given top-level operation.
   void printTopLevelOperation(Operation *op);
 
   /// Print the given operation with its indent and location.
   void print(Operation *op);
   /// Print the bare location, not including indentation/location/etc.
   void printOperation(Operation *op);
state.getImpl().getSSANameState().printValueID(*this, /*printResultNo=*/true,
                                                  os);
 }
 
 void Operation::print(raw_ostream &os, const OpPrintingFlags &printerFlags) {
   // If this is a top level operation, we also print aliases.
   if (!getParent() && !printerFlags.shouldUseLocalScope()) {
     AsmState state(this, printerFlags);
     state.getImpl().initializeAliases(this);
-    print(os, state, printerFlags);
+    print(os, state);
     return;
   }
 
   // Find the operation to number from based upon the provided flags.
   Operation *op = this;
   bool shouldUseLocalScope = printerFlags.shouldUseLocalScope();
   do {
     // If we are printing local scope, stop at the first operation that is
     // isolated from above.
     if (shouldUseLocalScope && op->hasTrait<OpTrait::IsIsolatedFromAbove>())
       break;
 
     // Otherwise, traverse up to the next parent.
     Operation *parentOp = op->getParentOp();
     if (!parentOp)
       break;
     op = parentOp;
   } while (true);
 
   AsmState state(op, printerFlags);
-  print(os, state, printerFlags);
+  print(os, state);
 }
-void Operation::print(raw_ostream &os, AsmState &state,
-                      const OpPrintingFlags &flags) {
-  OperationPrinter printer(os, flags, state.getImpl());
-  if (!getParent() && !flags.shouldUseLocalScope())
+void Operation::print(raw_ostream &os, AsmState &state) {
+  OperationPrinter printer(os, state.getImpl());
+  if (!getParent() && !state.getPrinterFlags().shouldUseLocalScope())
     printer.printTopLevelOperation(this);
   else
     printer.print(this);
 }
 
 void Operation::dump() {
   print(llvm::errs(), OpPrintingFlags().useLocalScope());
   llvm::errs() << "\n";
 }
 
 void Block::print(raw_ostream &os) {
   Operation *parentOp = getParentOp();
   if (!parentOp) {
     os << "<<UNLINKED BLOCK>>\n";
     return;
   }
   // Get the top-level op.
   while (auto *nextOp = parentOp->getParentOp())
     parentOp = nextOp;
 
   AsmState state(parentOp);
   print(os, state);
 }
 void Block::print(raw_ostream &os, AsmState &state) {
-  OperationPrinter(os, /*flags=*/llvm::None, state.getImpl()).print(this);
+  OperationPrinter(os, state.getImpl()).print(this);
 }
 
 void Block::dump() { print(llvm::errs()); }
 
 /// Print out the name of the block without printing its body.
 void Block::printAsOperand(raw_ostream &os, bool printType) {
   Operation *parentOp = getParentOp();
   if (!parentOp) {
     os << "<<UNLINKED BLOCK>>\n";
     return;
   }
   AsmState state(parentOp);
   printAsOperand(os, state);
 }
 void Block::printAsOperand(raw_ostream &os, AsmState &state) {
-  OperationPrinter printer(os, /*flags=*/llvm::None, state.getImpl());
+  OperationPrinter printer(os, state.getImpl());
   printer.printBlockName(this);
 }

mlir/lib/Transforms/LocationSnapshot.cpp

 /// NameLoc with the given tag as the name, and then fused with the existing
 /// locations. Otherwise, the existing locations are replaced.
 static void generateLocationsFromIR(raw_ostream &os, StringRef fileName,
                                     Operation *op, const OpPrintingFlags &flags,
                                     StringRef tag) {
   // Print the IR to the stream, and collect the raw line+column information.
   AsmState::LocationMap opToLineCol;
   AsmState state(op, flags, &opToLineCol);
-  op->print(os, state, flags);
+  op->print(os, state);
 
   Builder builder(op->getContext());
   Optional<StringAttr> tagIdentifier;
   if (!tag.empty())
     tagIdentifier = builder.getStringAttr(tag);
 
   // Walk and generate new locations for each of the operations.
   StringAttr file = builder.getStringAttr(fileName);