More cleanup in CodeWriter

- regular file, stdout and string are all written using ostream.
- factory functions GetFileWriter and GetStringWriter are refactored
to as static member functions ForFile and ForString

Bug: 110967839
Test: m -j, device boots to the UI and basic functionality works
Test: runtest.sh is successful

Change-Id: Id2ce1b7d5c2ff740b3241313c4bc5f6bd80f07bf
diff --git a/ast_cpp_unittest.cpp b/ast_cpp_unittest.cpp
index ce11c1c..188665c 100644
--- a/ast_cpp_unittest.cpp
+++ b/ast_cpp_unittest.cpp
@@ -95,8 +95,7 @@
   void CompareGeneratedCode(const AstNode& node,
                             const string& expected_output) {
     string actual_output;
-    CodeWriterPtr writer = GetStringWriter(&actual_output);
-    node.Write(writer.get());
+    node.Write(CodeWriter::ForString(&actual_output).get());
     EXPECT_EQ(expected_output, actual_output);
   }
 };  // class AstCppTests
diff --git a/ast_java_unittest.cpp b/ast_java_unittest.cpp
index b2aad23..a9baf37 100644
--- a/ast_java_unittest.cpp
+++ b/ast_java_unittest.cpp
@@ -53,8 +53,7 @@
   a_class.extends = &extend_type;
 
   string actual_output;
-  CodeWriterPtr writer = GetStringWriter(&actual_output);
-  a_class.Write(writer.get());
+  a_class.Write(CodeWriter::ForString(&actual_output).get());
   EXPECT_EQ(string(kExpectedClassOutput), actual_output);
 }
 
diff --git a/code_writer.cpp b/code_writer.cpp
index 048dade..e75d7a5 100644
--- a/code_writer.cpp
+++ b/code_writer.cpp
@@ -13,12 +13,12 @@
  * See the License for the specific language governing permissions and
  * limitations under the License.
  */
-
 #include "code_writer.h"
 
+#include <stdarg.h>
 #include <fstream>
 #include <iostream>
-#include <stdarg.h>
+#include <sstream>
 #include <vector>
 
 #include <android-base/stringprintf.h>
@@ -26,6 +26,8 @@
 namespace android {
 namespace aidl {
 
+CodeWriter::CodeWriter(std::unique_ptr<std::ostream> ostream) : ostream_(std::move(ostream)) {}
+
 std::string CodeWriter::ApplyIndent(const std::string& str) {
   std::string output;
   if (!start_of_line_ || str == "\n") {
@@ -62,72 +64,51 @@
   for (auto line : lines) {
     indented.append(ApplyIndent(line));
   }
-  return Output(indented);
+
+  (*ostream_) << indented;
+  return !ostream_->fail();
 }
 
-class StringCodeWriter : public CodeWriter {
- public:
-  explicit StringCodeWriter(std::string* output_buffer) : output_(output_buffer) {}
-  virtual ~StringCodeWriter() = default;
-
-  bool Output(const std::string& str) override {
-    output_->append(str);
-    return true;
+bool CodeWriter::Close() {
+  if (ostream_.get()->rdbuf() != std::cout.rdbuf()) {
+    // if the steam is for file (not stdout), do the close.
+    static_cast<std::fstream*>(ostream_.get())->close();
+    return !ostream_->fail();
   }
-
-  bool Close() override { return true; }
-
- private:
-  std::string* output_;
-};  // class StringCodeWriter
-
-class FileCodeWriter : public CodeWriter {
- public:
-  explicit FileCodeWriter(const std::string& filename) : cout_(std::cout) {
-    if (filename == "-") {
-      to_stdout_ = true;
-    } else {
-      to_stdout_ = false;
-      fileout_.open(filename, std::ofstream::out |
-                    std::ofstream::binary);
-      if (fileout_.fail()) {
-        std::cerr << "unable to open " << filename << " for write" << std::endl;
-      }
-    }
-  }
-
-  bool Output(const std::string& str) override {
-    if (to_stdout_) {
-      cout_ << str;
-    } else {
-      fileout_ << str;
-    }
-    return TestSuccess();
-  }
-
-  bool Close() override {
-    if (!to_stdout_) {
-      fileout_.close();
-    }
-    return TestSuccess();
-  }
-
-  bool TestSuccess() const {
-    return to_stdout_ ? true : !fileout_.fail();
-  }
-
- private:
-  std::ostream& cout_;
-  std::ofstream fileout_;
-  bool to_stdout_;
-};  // class StringCodeWriter
-
-CodeWriterPtr GetFileWriter(const std::string& output_file) {
-  return CodeWriterPtr(new FileCodeWriter(output_file));
+  return true;
 }
 
-CodeWriterPtr GetStringWriter(std::string* output_buffer) {
-  return CodeWriterPtr(new StringCodeWriter(output_buffer));
+CodeWriterPtr CodeWriter::ForFile(const std::string& filename) {
+  std::unique_ptr<std::ostream> stream;
+  if (filename == "-") {
+    stream = std::unique_ptr<std::ostream>(new std::ostream(std::cout.rdbuf()));
+  } else {
+    stream = std::unique_ptr<std::ostream>(
+        new std::fstream(filename, std::fstream::out | std::fstream::binary));
+  }
+  return CodeWriterPtr(new CodeWriter(std::move(stream)));
+}
+
+CodeWriterPtr CodeWriter::ForString(std::string* buf) {
+  // This class is defined inside this static function of CodeWriter
+  // in order to have access to private constructor and private member
+  // ostream_.
+  class StringCodeWriter : public CodeWriter {
+   public:
+    StringCodeWriter(std::string* buf)
+        : CodeWriter(std::unique_ptr<std::ostream>(new std::stringstream())), buf_(buf) {}
+    ~StringCodeWriter() { Close(); }
+    bool Close() override {
+      // extract whats written to the stringstream to the external buffer.
+      // we are sure that ostream_ is indeed stringstream.
+      *buf_ = static_cast<std::stringstream*>(ostream_.get())->str();
+      return true;
+    }
+
+   private:
+    std::string* buf_;
+  };
+  return CodeWriterPtr(new StringCodeWriter(buf));
 }
 
 }  // namespace aidl
diff --git a/code_writer.h b/code_writer.h
index 73e84e3..2a81dac 100644
--- a/code_writer.h
+++ b/code_writer.h
@@ -27,34 +27,34 @@
 namespace android {
 namespace aidl {
 
+class CodeWriter;
+using CodeWriterPtr = std::unique_ptr<CodeWriter>;
+
 class CodeWriter {
  public:
+  // Get a CodeWriter that writes to a file. When filename is "-",
+  // it is written to stdout.
+  static CodeWriterPtr ForFile(const std::string& filename);
+  // Get a CodeWriter that writes to a string buffer.
+  // The buffer gets updated only after Close() is called or the CodeWriter
+  // is deleted -- much like a real file.
+  static CodeWriterPtr ForString(std::string* buf);
   // Write a formatted string to this writer in the usual printf sense.
   // Returns false on error.
   virtual bool Write(const char* format, ...);
-  virtual bool Close() = 0;
-  virtual ~CodeWriter() = default;
   inline void Indent() { indent_level_++; }
   inline void Dedent() { indent_level_--; }
- protected:
-  // Actuall writes str which is formatted and properly indented
-  // to the actual medium (file, string, etc.)
-  virtual bool Output(const std::string& str) = 0;
+  virtual bool Close();
+  virtual ~CodeWriter() = default;
+  CodeWriter() = default;
+
  private:
+  CodeWriter(std::unique_ptr<std::ostream> ostream);
   std::string ApplyIndent(const std::string& str);
+  const std::unique_ptr<std::ostream> ostream_;
   int indent_level_ {0};
   bool start_of_line_ {true};
-};  // class CodeWriter
-
-using CodeWriterPtr = std::unique_ptr<CodeWriter>;
-
-// Get a CodeWriter that writes to |output_file|.
-CodeWriterPtr GetFileWriter(const std::string& output_file);
-
-// Get a CodeWriter that writes to a string buffer.
-// Caller retains ownership of the buffer.
-// The buffer must outlive the CodeWriter.
-CodeWriterPtr GetStringWriter(std::string* output_buffer);
+};
 
 }  // namespace aidl
 }  // namespace android
diff --git a/generate_cpp_unittest.cpp b/generate_cpp_unittest.cpp
index 1e15f6b..aa3bd25 100644
--- a/generate_cpp_unittest.cpp
+++ b/generate_cpp_unittest.cpp
@@ -1314,9 +1314,7 @@
 
   void Compare(Document* doc, const char* expected) {
     string output;
-    unique_ptr<CodeWriter> cw = GetStringWriter(&output);
-
-    doc->Write(cw.get());
+    doc->Write(CodeWriter::ForString(&output).get());
 
     if (expected == output) {
       return; // Success
diff --git a/io_delegate.cpp b/io_delegate.cpp
index 62be3f5..bd6eb09 100644
--- a/io_delegate.cpp
+++ b/io_delegate.cpp
@@ -175,7 +175,7 @@
 
 unique_ptr<CodeWriter> IoDelegate::GetCodeWriter(
     const string& file_path) const {
-  return GetFileWriter(file_path);
+  return CodeWriter::ForFile(file_path);
 }
 
 void IoDelegate::RemovePath(const std::string& file_path) const {
diff --git a/tests/fake_io_delegate.cpp b/tests/fake_io_delegate.cpp
index c4b3584..68f0ff2 100644
--- a/tests/fake_io_delegate.cpp
+++ b/tests/fake_io_delegate.cpp
@@ -36,7 +36,6 @@
 class BrokenCodeWriter : public CodeWriter {
   bool Write(const char* /* format */, ...) override {  return true; }
   bool Close() override { return false; }
-  bool Output(const std::string& /*str*/) override { return true; }
   virtual ~BrokenCodeWriter() = default;
 };  // class BrokenCodeWriter
 
@@ -84,7 +83,7 @@
   }
   removed_files_.erase(file_path);
   written_file_contents_[file_path] = "";
-  return GetStringWriter(&written_file_contents_[file_path]);
+  return CodeWriter::ForString(&written_file_contents_[file_path]);
 }
 
 void FakeIoDelegate::RemovePath(const std::string& file_path) const {