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: |
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 = "" , bool indent = true) |
33 | : os(os), open(open), close(close), indent(indent) { |
34 | os << open; |
35 | if (indent) |
36 | os.indent(); |
37 | } |
38 | ~DelimitedScope() { |
39 | if (indent) |
40 | os.unindent(); |
41 | os << close; |
42 | } |
43 | |
44 | raw_indented_ostream &os; |
45 | |
46 | private: |
47 | StringRef open, close; |
48 | bool indent; |
49 | }; |
50 | |
51 | /// Returns the underlying (unindented) raw_ostream. |
52 | raw_ostream &getOStream() const { return os; } |
53 | |
54 | /// Returns DelimitedScope. |
55 | DelimitedScope scope(StringRef open = "" , StringRef close = "" , |
56 | bool indent = true) { |
57 | return DelimitedScope(*this, open, close, indent); |
58 | } |
59 | |
60 | /// Prints a string re-indented to the current indent. Re-indents by removing |
61 | /// the leading whitespace from the first non-empty line from every line of |
62 | /// the string, skipping over empty lines at the start. Prefixes each line |
63 | /// with extraPrefix after the indentation. |
64 | raw_indented_ostream &printReindented(StringRef str, |
65 | StringRef = "" ); |
66 | |
67 | /// Increases the indent and returning this raw_indented_ostream. |
68 | raw_indented_ostream &indent() { |
69 | currentIndent += indentSize; |
70 | return *this; |
71 | } |
72 | |
73 | /// Decreases the indent and returning this raw_indented_ostream. |
74 | raw_indented_ostream &unindent() { |
75 | currentIndent = std::max(a: 0, b: currentIndent - indentSize); |
76 | return *this; |
77 | } |
78 | |
79 | /// Emits whitespace and sets the indentation for the stream. |
80 | raw_indented_ostream &indent(int with) { |
81 | os.indent(NumSpaces: with); |
82 | atStartOfLine = false; |
83 | currentIndent = with; |
84 | return *this; |
85 | } |
86 | |
87 | private: |
88 | void write_impl(const char *ptr, size_t size) final; |
89 | |
90 | /// Return the current position within the stream, not counting the bytes |
91 | /// currently in the buffer. |
92 | uint64_t current_pos() const final { return os.tell(); } |
93 | |
94 | /// Constant indent added/removed. |
95 | static constexpr int indentSize = 2; |
96 | |
97 | /// Tracker for current indentation. |
98 | int currentIndent = 0; |
99 | |
100 | /// The leading whitespace of the string being printed, if reindent is used. |
101 | int leadingWs = 0; |
102 | |
103 | /// The extra prefix to be printed, if reindent is used. |
104 | StringRef ; |
105 | |
106 | /// Tracks whether at start of line and so indent is required or not. |
107 | bool atStartOfLine = true; |
108 | |
109 | /// The underlying raw_ostream. |
110 | raw_ostream &os; |
111 | }; |
112 | |
113 | inline raw_indented_ostream & |
114 | mlir::raw_indented_ostream::printReindented(StringRef str, |
115 | StringRef ) { |
116 | StringRef output = str; |
117 | // Skip empty lines. |
118 | while (!output.empty()) { |
119 | auto split = output.split(Separator: '\n'); |
120 | // Trim Windows \r characters from \r\n line endings. |
121 | auto firstTrimmed = split.first.rtrim(Char: '\r'); |
122 | size_t indent = firstTrimmed.find_first_not_of(Chars: " \t" ); |
123 | if (indent != StringRef::npos) { |
124 | // Set an initial value. |
125 | leadingWs = indent; |
126 | break; |
127 | } |
128 | output = split.second; |
129 | } |
130 | // Determine the maximum indent. |
131 | StringRef remaining = output; |
132 | while (!remaining.empty()) { |
133 | auto split = remaining.split(Separator: '\n'); |
134 | auto firstTrimmed = split.first.rtrim(Char: '\r'); |
135 | size_t indent = firstTrimmed.find_first_not_of(Chars: " \t" ); |
136 | if (indent != StringRef::npos) |
137 | leadingWs = std::min(a: leadingWs, b: static_cast<int>(indent)); |
138 | remaining = split.second; |
139 | } |
140 | // Print, skipping the empty lines. |
141 | std::swap(a&: currentExtraPrefix, b&: extraPrefix); |
142 | *this << output; |
143 | std::swap(a&: currentExtraPrefix, b&: extraPrefix); |
144 | leadingWs = 0; |
145 | return *this; |
146 | } |
147 | |
148 | inline void mlir::raw_indented_ostream::write_impl(const char *ptr, |
149 | size_t size) { |
150 | StringRef str(ptr, size); |
151 | // Print out indented. |
152 | auto print = [this](StringRef str) { |
153 | if (atStartOfLine) |
154 | os.indent(NumSpaces: currentIndent) << currentExtraPrefix << str.substr(Start: leadingWs); |
155 | else |
156 | os << str.substr(Start: leadingWs); |
157 | }; |
158 | |
159 | while (!str.empty()) { |
160 | size_t idx = str.find(C: '\n'); |
161 | if (idx == StringRef::npos) { |
162 | if (!str.substr(Start: leadingWs).empty()) { |
163 | print(str); |
164 | atStartOfLine = false; |
165 | } |
166 | break; |
167 | } |
168 | |
169 | auto split = |
170 | std::make_pair(x: str.slice(Start: 0, End: idx), y: str.slice(Start: idx + 1, End: StringRef::npos)); |
171 | // Print empty new line without spaces if line only has spaces and no extra |
172 | // prefix is requested. |
173 | if (!split.first.ltrim().empty() || !currentExtraPrefix.empty()) |
174 | print(split.first); |
175 | os << '\n'; |
176 | atStartOfLine = true; |
177 | str = split.second; |
178 | } |
179 | } |
180 | |
181 | } // namespace mlir |
182 | #endif // MLIR_SUPPORT_INDENTEDOSTREAM_H_ |
183 | |