1 files changed, 87 insertions, 1 deletions

diff --git a/llvm/lib/Demangle/Utility.h b/llvm/lib/Demangle/Utility.h
index 62c0e7430a5..f764663436d 100644
--- a/llvm/lib/Demangle/Utility.h
+++ b/llvm/lib/Demangle/Utility.h
@@ -12,6 +12,8 @@
 #ifndef LLVM_DEMANGLE_UTILITY_H
 #define LLVM_DEMANGLE_UTILITY_H
 
+#include "StringView.h"
+
 #include <cstdlib>
 #include <cstring>
 #include <limits>
@@ -33,6 +35,27 @@ class OutputStream {
     }
   }
 
+  void writeUnsigned(uint64_t N, bool isNeg = false) {
+    // Handle special case...
+    if (N == 0) {
+      *this << '0';
+      return;
+    }
+
+    char Temp[21];
+    char *TempPtr = std::end(Temp);
+
+    while (N) {
+      *--TempPtr = '0' + char(N % 10);
+      N /= 10;
+    }
+
+    // Add negative sign...
+    if (isNeg)
+      *--TempPtr = '-';
+    this->operator<<(StringView(TempPtr, std::end(Temp)));
+  }
+
 public:
   OutputStream(char *StartBuf, size_t Size)
       : Buffer(StartBuf), CurrentPosition(0), BufferCapacity(Size) {}
@@ -43,6 +66,20 @@ public:
     BufferCapacity = BufferCapacity_;
   }
 
+  /// Create an OutputStream from a buffer and a size.  If either of these are
+  /// null a buffer is allocated.
+  static OutputStream create(char *StartBuf, size_t *Size, size_t AllocSize) {
+    OutputStream Result;
+
+    if (!StartBuf || !Size) {
+      StartBuf = static_cast<char *>(std::malloc(AllocSize));
+      Size = &AllocSize;
+    }
+
+    Result.reset(StartBuf, *Size);
+    return Result;
+  }
+
   /// If a ParameterPackExpansion (or similar type) is encountered, the offset
   /// into the pack that we're currently printing.
   unsigned CurrentPackIndex = std::numeric_limits<unsigned>::max();
@@ -64,6 +101,39 @@ public:
     return *this;
   }
 
+  OutputStream &operator<<(StringView R) { return (*this += R); }
+
+  OutputStream &operator<<(char C) { return (*this += C); }
+
+  OutputStream &operator<<(long long N) {
+    if (N < 0)
+      writeUnsigned(static_cast<unsigned long long>(-N), true);
+    else
+      writeUnsigned(static_cast<unsigned long long>(N));
+    return *this;
+  }
+
+  OutputStream &operator<<(unsigned long long N) {
+    writeUnsigned(N, false);
+    return *this;
+  }
+
+  OutputStream &operator<<(long N) {
+    return this->operator<<(static_cast<long long>(N));
+  }
+
+  OutputStream &operator<<(unsigned long N) {
+    return this->operator<<(static_cast<unsigned long long>(N));
+  }
+
+  OutputStream &operator<<(int N) {
+    return this->operator<<(static_cast<long long>(N));
+  }
+
+  OutputStream &operator<<(unsigned int N) {
+    return this->operator<<(static_cast<unsigned long long>(N));
+  }
+
   size_t getCurrentPosition() const { return CurrentPosition; }
   void setCurrentPosition(size_t NewPos) { CurrentPosition = NewPos; }
 
@@ -81,13 +151,29 @@ public:
 template <class T> class SwapAndRestore {
   T &Restore;
   T OriginalValue;
+  bool ShouldRestore = true;
 
 public:
+  SwapAndRestore(T &Restore_) : SwapAndRestore(Restore_, Restore_) {}
+
   SwapAndRestore(T &Restore_, T NewVal)
       : Restore(Restore_), OriginalValue(Restore) {
     Restore = std::move(NewVal);
   }
-  ~SwapAndRestore() { Restore = std::move(OriginalValue); }
+  ~SwapAndRestore() {
+    if (ShouldRestore)
+      Restore = std::move(OriginalValue);
+  }
+
+  void shouldRestore(bool ShouldRestore_) { ShouldRestore = ShouldRestore_; }
+
+  void restoreNow(bool Force) {
+    if (!Force && !ShouldRestore)
+      return;
+
+    Restore = std::move(OriginalValue);
+    ShouldRestore = false;
+  }
 
   SwapAndRestore(const SwapAndRestore &) = delete;
   SwapAndRestore &operator=(const SwapAndRestore &) = delete;