1 //===- IndentedOstream.h - raw ostream wrapper to indent --------*- C++ -*-===//
2 //
3 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
4 // See https://llvm.org/LICENSE.txt for license information.
5 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
6 //
7 //===----------------------------------------------------------------------===//
8 //
9 // raw_ostream subclass that keeps track of indentation for textual output
10 // where indentation helps readability.
11 //
12 //===----------------------------------------------------------------------===//
13 
14 #ifndef MLIR_SUPPORT_INDENTEDOSTREAM_H_
15 #define MLIR_SUPPORT_INDENTEDOSTREAM_H_
16 
17 #include "mlir/Support/LLVM.h"
18 #include "llvm/Support/raw_ostream.h"
19 
20 namespace mlir {
21 
22 /// raw_ostream subclass that simplifies indention a sequence of code.
23 class raw_indented_ostream : public raw_ostream {
24 public:
raw_indented_ostream(llvm::raw_ostream & os)25   explicit raw_indented_ostream(llvm::raw_ostream &os) : os(os) {
26     SetUnbuffered();
27   }
28 
29   /// Simple RAII struct to use to indentation around entering/exiting region.
30   struct DelimitedScope {
31     explicit DelimitedScope(raw_indented_ostream &os, StringRef open = "",
32                             StringRef close = "")
osDelimitedScope33         : os(os), open(open), close(close) {
34       os << open;
35       os.indent();
36     }
~DelimitedScopeDelimitedScope37     ~DelimitedScope() {
38       os.unindent();
39       os << close;
40     }
41 
42     raw_indented_ostream &os;
43 
44   private:
45     llvm::StringRef open, close;
46   };
47 
48   /// Returns DelimitedScope.
49   DelimitedScope scope(StringRef open = "", StringRef close = "") {
50     return DelimitedScope(*this, open, close);
51   }
52 
53   /// Re-indents by removing the leading whitespace from the first non-empty
54   /// line from every line of the string, skipping over empty lines at the
55   /// start.
56   raw_indented_ostream &reindent(StringRef str);
57 
58   /// Increases the indent and returning this raw_indented_ostream.
indent()59   raw_indented_ostream &indent() {
60     currentIndent += indentSize;
61     return *this;
62   }
63 
64   /// Decreases the indent and returning this raw_indented_ostream.
unindent()65   raw_indented_ostream &unindent() {
66     currentIndent = std::max(0, currentIndent - indentSize);
67     return *this;
68   }
69 
70   /// Emits whitespace and sets the indendation for the stream.
indent(int with)71   raw_indented_ostream &indent(int with) {
72     os.indent(with);
73     atStartOfLine = false;
74     currentIndent = with;
75     return *this;
76   }
77 
78 private:
79   void write_impl(const char *ptr, size_t size) override;
80 
81   /// Return the current position within the stream, not counting the bytes
82   /// currently in the buffer.
current_pos()83   uint64_t current_pos() const override { return os.tell(); }
84 
85   /// Constant indent added/removed.
86   static constexpr int indentSize = 2;
87 
88   // Tracker for current indentation.
89   int currentIndent = 0;
90 
91   // The leading whitespace of the string being printed, if reindent is used.
92   int leadingWs = 0;
93 
94   // Tracks whether at start of line and so indent is required or not.
95   bool atStartOfLine = true;
96 
97   // The underlying raw_ostream.
98   raw_ostream &os;
99 };
100 
101 } // namespace mlir
102 #endif // MLIR_SUPPORT_INDENTEDOSTREAM_H_
103