From 6bed79b3f00d3e2c273bc36ed350f802d76607b3 Mon Sep 17 00:00:00 2001
From: Rahul Joshi <rjoshi@nvidia.com>
Date: Mon, 23 Sep 2024 18:41:40 -0700
Subject: [PATCH] [Support] Add scaling support in `indent` (#109478)

Scaled indent is useful when indentation is always in steps of a fixed
number (the Scale) and still allow using the +/- operators to adjust
indentation.
---
 llvm/include/llvm/Support/raw_ostream.h     | 29 ++++++++++++++++-----
 llvm/unittests/Support/raw_ostream_test.cpp | 10 +++++++
 2 files changed, 32 insertions(+), 7 deletions(-)

diff --git a/llvm/include/llvm/Support/raw_ostream.h b/llvm/include/llvm/Support/raw_ostream.h
index 34f91cbe9551..c2f2299ed964 100644
--- a/llvm/include/llvm/Support/raw_ostream.h
+++ b/llvm/include/llvm/Support/raw_ostream.h
@@ -774,18 +774,33 @@ public:
 // you can use
 // OS << indent(6) << "more stuff";
 // which has better ergonomics (and clang-formats better as well).
+//
+// If indentation is always in increments of a fixed value, you can use Scale
+// to set that value once. So indent(1, 2) will add 2 spaces and
+// indent(1,2) + 1 will add 4 spaces.
 struct indent {
-  unsigned NumSpaces;
+  // Indentation is represented as `NumIndents` steps of size `Scale` each.
+  unsigned NumIndents;
+  unsigned Scale;
 
-  explicit indent(unsigned NumSpaces) : NumSpaces(NumSpaces) {}
-  void operator+=(unsigned N) { NumSpaces += N; }
-  void operator-=(unsigned N) { NumSpaces -= N; }
-  indent operator+(unsigned N) const { return indent(NumSpaces + N); }
-  indent operator-(unsigned N) const { return indent(NumSpaces - N); }
+  explicit indent(unsigned NumIndents, unsigned Scale = 1)
+      : NumIndents(NumIndents), Scale(Scale) {}
+
+  // These arithmeric operators preserve scale.
+  void operator+=(unsigned N) { NumIndents += N; }
+  void operator-=(unsigned N) {
+    assert(NumIndents >= N && "Indentation underflow");
+    NumIndents -= N;
+  }
+  indent operator+(unsigned N) const { return indent(NumIndents + N, Scale); }
+  indent operator-(unsigned N) const {
+    assert(NumIndents >= N && "Indentation undeflow");
+    return indent(NumIndents - N, Scale);
+  }
 };
 
 inline raw_ostream &operator<<(raw_ostream &OS, const indent &Indent) {
-  return OS.indent(Indent.NumSpaces);
+  return OS.indent(Indent.NumIndents * Indent.Scale);
 }
 
 class Error;
diff --git a/llvm/unittests/Support/raw_ostream_test.cpp b/llvm/unittests/Support/raw_ostream_test.cpp
index 99aa350adad7..a35edd616852 100644
--- a/llvm/unittests/Support/raw_ostream_test.cpp
+++ b/llvm/unittests/Support/raw_ostream_test.cpp
@@ -188,6 +188,16 @@ TEST(raw_ostreamTest, Indent) {
   EXPECT_EQ(Spaces(5), printToString(Indent));
   Indent -= 1;
   EXPECT_EQ(Spaces(4), printToString(Indent));
+
+  // Scaled indent.
+  indent Scaled(4, 2);
+  EXPECT_EQ(Spaces(8), printToString(Scaled));
+  EXPECT_EQ(Spaces(10), printToString(Scaled + 1));
+  EXPECT_EQ(Spaces(6), printToString(Scaled - 1));
+  Scaled += 1;
+  EXPECT_EQ(Spaces(10), printToString(Scaled));
+  Scaled -= 1;
+  EXPECT_EQ(Spaces(8), printToString(Scaled));
 }
 
 TEST(raw_ostreamTest, FormatHex) {