From ac6415d465ab12c269629c1139f47c0cae7cc80a Mon Sep 17 00:00:00 2001
From: Snaipe <franklinmathieu@gmail.com>
Date: Sun, 13 Sep 2015 13:11:11 +0200
Subject: [PATCH] Added C++ stream interface for redirection functions

---
 include/criterion/redirect.h | 104 +++++++++++++++++++++++++++++++++++
 samples/redirect.cc          |  14 ++---
 2 files changed, 111 insertions(+), 7 deletions(-)

diff --git a/include/criterion/redirect.h b/include/criterion/redirect.h
index cf1efcf..eacd0bb 100644
--- a/include/criterion/redirect.h
+++ b/include/criterion/redirect.h
@@ -29,6 +29,12 @@
 
 # ifdef __cplusplus
 #  include <cstdio>
+#  include <memory>
+#  include <fstream>
+
+#  ifdef __GNUC__
+#   include <ext/stdio_filebuf.h>
+#  endif
 # else
 #  include <stdio.h>
 # endif
@@ -73,4 +79,102 @@ CR_END_C_API
 # define cr_assert_file_contents_not_match_str(...) CR_EXPAND(cr_assert_redir_op_va_(CR_FAIL_ABORT_,     cr_file_match_str, !=, __VA_ARGS__))
 # define cr_expect_file_contents_not_match_str(...) CR_EXPAND(cr_assert_redir_op_va_(CR_FAIL_CONTINUES_, cr_file_match_str, !=, __VA_ARGS__))
 
+# ifdef __cplusplus
+namespace criterion {
+
+    template <typename CharT>
+    class basic_ofstream : public std::basic_ofstream<CharT> {
+    public:
+        basic_ofstream(FILE* f)
+#  ifdef __GNUC__
+            : std::ofstream()
+            , fbuf(new ::__gnu_cxx::stdio_filebuf<CharT>(f, std::ios::out))
+#  else
+            : std::ofstream(f)
+#  endif
+            , file(f)
+        {
+            std::ios::rdbuf(&*fbuf);
+        }
+
+        void close(void) {
+            std::basic_ofstream<CharT>::flush();
+            std::basic_ofstream<CharT>::close();
+            std::fclose(file);
+        }
+
+    private:
+#  ifdef __GNUC__
+        std::unique_ptr<::__gnu_cxx::stdio_filebuf<CharT>> fbuf;
+#  endif
+        std::FILE* file;
+    };
+
+    template <typename CharT>
+    class basic_ifstream : public std::basic_ifstream<CharT> {
+    public:
+        basic_ifstream(FILE* f)
+#  ifdef __GNUC__
+            : std::ifstream()
+            , fbuf(new ::__gnu_cxx::stdio_filebuf<CharT>(f, std::ios::in))
+#  else
+            : std::ifstream(f)
+#  endif
+            , file(f)
+        {
+            std::ios::rdbuf(&*fbuf);
+        }
+
+        void close(void) {
+            std::basic_ifstream<CharT>::flush();
+            std::basic_ifstream<CharT>::close();
+            std::fclose(file);
+        }
+
+    private:
+#  ifdef __GNUC__
+        std::unique_ptr<::__gnu_cxx::stdio_filebuf<CharT>> fbuf;
+#  endif
+        std::FILE* file;
+    };
+
+    template <typename CharT>
+    struct get_redirected_out_stream_ {
+        static inline basic_ofstream<CharT>& call(std::FILE* f) {
+            static std::unique_ptr<basic_ofstream<CharT>> stream;
+
+            if (!stream)
+                stream.reset(new basic_ofstream<CharT>(f));
+            return *stream;
+        }
+
+    };
+
+    template <typename CharT>
+    struct get_redirected_in_stream_ {
+        static inline basic_ifstream<CharT>& call(std::FILE* f) {
+            static std::unique_ptr<basic_ifstream<CharT>> stream;
+            if (!stream)
+                stream.reset(new basic_ifstream<CharT>(f));
+            return *stream;
+        }
+    };
+
+    using ofstream = basic_ofstream<char>;
+    using ifstream = basic_ifstream<char>;
+
+    static inline ofstream& get_redirected_cin(void) {
+        return get_redirected_out_stream_<char>::call(cr_get_redirected_stdin());
+    }
+
+    static inline ifstream& get_redirected_cout(void) {
+        return get_redirected_in_stream_<char>::call(cr_get_redirected_stdout());
+    }
+
+    static inline ifstream& get_redirected_cerr(void) {
+        return get_redirected_in_stream_<char>::call(cr_get_redirected_stderr());
+    }
+}
+# endif
+
 #endif /* !CRITERION_REDIRECT_H_ */
diff --git a/samples/redirect.cc b/samples/redirect.cc
index f4eae63..29ee036 100644
--- a/samples/redirect.cc
+++ b/samples/redirect.cc
@@ -6,6 +6,10 @@
 #include <fstream>
 #include <cctype>
 
+#ifdef __GNUC__
+# include <ext/stdio_filebuf.h>
+#endif
+
 // Testing stdout/stderr
 
 void redirect_all_std(void) {
@@ -14,16 +18,13 @@ void redirect_all_std(void) {
 }
 
 Test(redirect, test_outputs, .init = redirect_all_std) {
-    std::FILE* f_stdout = cr_get_redirected_stdout();
-
     std::cout << "foo" << std::flush;
+    std::cerr << "bar" << std::flush;
 
+    std::FILE* f_stdout = cr_get_redirected_stdout();
     cr_assert_file_contents_match_str(f_stdout, "foo");
 
     std::FILE* f_stderr = cr_get_redirected_stderr();
-
-    std::cerr << "bar" << std::flush;
-
     cr_assert_file_contents_match_str(f_stderr, "bar");
 }
 
@@ -43,8 +44,7 @@ void rot13_io(void) {
 }
 
 Test(redirect, rot13, .init = cr_redirect_stdout) {
-    std::FILE* f_stdin = cr_get_redirected_stdin();
-    std::ofstream f_cin(f_stdin);
+    auto& f_cin = criterion::get_redirected_cin();
 
     f_cin << "the quick brown fox jumps over the lazy dog";
     f_cin.close();