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 extraPrefix = ""); 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(0, 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(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 currentExtraPrefix; 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 extraPrefix) { 116 StringRef output = str; 117 // Skip empty lines. 118 while (!output.empty()) { 119 auto split = output.split('\n'); 120 // Trim Windows \r characters from \r\n line endings. 121 auto firstTrimmed = split.first.rtrim('\r'); 122 size_t indent = firstTrimmed.find_first_not_of(" \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('\n'); 134 auto firstTrimmed = split.first.rtrim('\r'); 135 size_t indent = firstTrimmed.find_first_not_of(" \t"); 136 if (indent != StringRef::npos) 137 leadingWs = std::min(leadingWs, static_cast<int>(indent)); 138 remaining = split.second; 139 } 140 // Print, skipping the empty lines. 141 std::swap(currentExtraPrefix, extraPrefix); 142 *this << output; 143 std::swap(currentExtraPrefix, 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(currentIndent) << currentExtraPrefix << str.substr(leadingWs); 155 else 156 os << str.substr(leadingWs); 157 }; 158 159 while (!str.empty()) { 160 size_t idx = str.find('\n'); 161 if (idx == StringRef::npos) { 162 if (!str.substr(leadingWs).empty()) { 163 print(str); 164 atStartOfLine = false; 165 } 166 break; 167 } 168 169 auto split = std::make_pair(str.substr(0, idx), str.substr(idx + 1)); 170 // Print empty new line without spaces if line only has spaces and no extra 171 // prefix is requested. 172 if (!split.first.ltrim().empty() || !currentExtraPrefix.empty()) 173 print(split.first); 174 os << '\n'; 175 atStartOfLine = true; 176 str = split.second; 177 } 178 } 179 180 } // namespace mlir 181 #endif // MLIR_SUPPORT_INDENTEDOSTREAM_H_ 182