2 files changed, 19 insertions, 15 deletions

diff --git a/include/llvm/Support/FormattedStream.h b/include/llvm/Support/FormattedStream.h
index 11a0e18db4..a74ee6f7d9 100644
--- a/include/llvm/Support/FormattedStream.h
+++ b/include/llvm/Support/FormattedStream.h
@@ -101,13 +101,11 @@ namespace llvm
 
     ~formatted_raw_ostream() {
       flush();
-      if (DeleteStream)
-        delete TheStream;
+      releaseStream();
     }
-    
+
     void setStream(raw_ostream &Stream, bool Delete = false) {
-      if (DeleteStream)
-        delete TheStream;
+      releaseStream();
 
       TheStream = &Stream;
       DeleteStream = Delete;
@@ -118,6 +116,8 @@ namespace llvm
       // had been using, and tell TheStream not to do its own buffering.
       if (size_t BufferSize = TheStream->GetBufferSize())
         SetBufferSize(BufferSize);
+      else
+        SetUnbuffered();
       TheStream->SetUnbuffered();
 
       Scanned = begin();
@@ -130,6 +130,20 @@ namespace llvm
     /// recent I/O, even if the current column + minpad > newcol.
     ///
     void PadToColumn(unsigned NewCol, unsigned MinPad = 0);
+
+  private:
+    void releaseStream() {
+      // Delete the stream if needed. Otherwise, transfer the buffer
+      // settings from this raw_ostream back to the underlying stream.
+      if (!TheStream)
+        return;
+      if (DeleteStream)
+        delete TheStream;
+      else if (size_t BufferSize = GetBufferSize())
+        TheStream->SetBufferSize(BufferSize);
+      else
+        TheStream->SetUnbuffered();
+    }
   };
 
 /// fouts() - This returns a reference to a formatted_raw_ostream for
diff --git a/lib/VMCore/AsmWriter.cpp b/lib/VMCore/AsmWriter.cpp
index 35aae41e34..68a45221ac 100644
--- a/lib/VMCore/AsmWriter.cpp
+++ b/lib/VMCore/AsmWriter.cpp
@@ -2005,14 +2005,9 @@ void Module::print(std::ostream &o, AssemblyAnnotationWriter *AAW) const {
 }
 void Module::print(raw_ostream &ROS, AssemblyAnnotationWriter *AAW) const {
   SlotTracker SlotTable(this);
-  size_t OldBufferSize = ROS.GetBufferSize();
   formatted_raw_ostream OS(ROS);
   AssemblyWriter W(OS, SlotTable, this, AAW);
   W.write(this);
-  // formatted_raw_ostream forces the underlying raw_ostream to be
-  // unbuffered. Reset it to its original buffer size.
-  if (OldBufferSize != 0)
-    ROS.SetBufferSize(OldBufferSize);
 }
 
 void Type::print(std::ostream &o) const {
@@ -2033,7 +2028,6 @@ void Value::print(raw_ostream &ROS, AssemblyAnnotationWriter *AAW) const {
     ROS << "printing a <null> value\n";
     return;
   }
-  size_t OldBufferSize = ROS.GetBufferSize();
   formatted_raw_ostream OS(ROS);
   if (const Instruction *I = dyn_cast<Instruction>(this)) {
     const Function *F = I->getParent() ? I->getParent()->getParent() : 0;
@@ -2089,10 +2083,6 @@ void Value::print(raw_ostream &ROS, AssemblyAnnotationWriter *AAW) const {
   } else {
     llvm_unreachable("Unknown value to print out!");
   }
-  // formatted_raw_ostream forces the underlying raw_ostream to be
-  // unbuffered. Reset it to its original buffer size.
-  if (OldBufferSize != 0)
-    ROS.SetBufferSize(OldBufferSize);
 }
 
 void Value::print(std::ostream &O, AssemblyAnnotationWriter *AAW) const {