Branch data     Line data    Source code

       1                 :            : // *****************************************************************************
       2                 :            : /*!
       3                 :            :   \file      src/Base/TeeBuf.hpp
       4                 :            :   \copyright 2012-2015 J. Bakosi,
       5                 :            :              2016-2018 Los Alamos National Security, LLC.,
       6                 :            :              2019-2021 Triad National Security, LLC.
       7                 :            :              All rights reserved. See the LICENSE file for details.
       8                 :            :   \brief     Tee stream buffer
       9                 :            :   \details   Tee stream buffer that can be used to duplicate a stream. Example:
      10                 :            :     \code{.cpp}
      11                 :            :       // createa a file
      12                 :            :       std::ofstream file( "data.txt" );
      13                 :            :       // tbuf will write to both file and cout
      14                 :            :       teebuf tbuf( file.rdbuf(), std::cout.rdbuf() );
      15                 :            :       // replace cout's streambuf with tbuf
      16                 :            :       scoped_streambuf_assignment ssa( std::cout, &tbuf );
      17                 :            :       // write to both std::cout and data.txt
      18                 :            :       std::cout << "Hello World" << std::endl;
      19                 :            :    \endcode
      20                 :            : */
      21                 :            : // *****************************************************************************
      22                 :            : #ifndef Tee_h
      23                 :            : #define Tee_h
      24                 :            : 
      25                 :            : #include <iostream>
      26                 :            : #include <streambuf>
      27                 :            : #include <fstream>
      28                 :            : #include <memory>
      29                 :            : 
      30                 :            : namespace tk {
      31                 :            : 
      32                 :            : template< class charT, class traits = std::char_traits<charT> >
      33                 :            : class basic_teebuf : public std::basic_streambuf< charT, traits > {
      34                 :            : 
      35                 :            :   public:
      36                 :            :     using char_type = charT;
      37                 :            :     using int_type = typename traits::int_type;
      38                 :            :     using pos_type = typename traits::pos_type;
      39                 :            :     using off_type = typename traits::off_type;
      40                 :            :     using traits_type = traits;
      41                 :            :     using streambuf_type = std::basic_streambuf< charT, traits >;
      42                 :            : 
      43                 :            :   private:
      44                 :            :     streambuf_type* m_sbuf1;
      45                 :            :     streambuf_type* m_sbuf2;
      46                 :            :     std::unique_ptr< char_type[] > m_buffer;
      47                 :            : 
      48                 :            :     enum { BUFFER_SIZE = 4096 / sizeof( char_type ) };
      49                 :            : 
      50                 :            :   public:
      51                 :      10636 :     basic_teebuf( streambuf_type *sbuf1, streambuf_type *sbuf2 )
      52                 :            :       : m_sbuf1( sbuf1 ), m_sbuf2( sbuf2 ),
      53         [ +  - ]:      10636 :         m_buffer( std::make_unique< char_type[] >( BUFFER_SIZE ) )
      54                 :      10636 :     { this->setp( m_buffer.get(), m_buffer.get() + BUFFER_SIZE ); }
      55                 :            : 
      56                 :       9593 :     ~basic_teebuf() override { this->pubsync(); }
      57                 :            : 
      58                 :            :   protected:
      59                 :      19471 :     virtual int_type overflow( int_type c = traits_type::eof() ) override {
      60                 :            :       // empty our buffer into m_sbuf1 and m_sbuf2
      61                 :      19471 :       std::streamsize n =
      62                 :      19471 :         static_cast< std::streamsize >( this->pptr() - this->pbase() );
      63                 :      19471 :       std::streamsize size1 = m_sbuf1->sputn( this->pbase(), n );
      64                 :      19471 :       std::streamsize size2 = m_sbuf2->sputn( this->pbase(), n );
      65         [ +  - ]:      19471 :       if ( size1 != n || size2 != n ) return traits_type::eof();
      66                 :            : 
      67                 :            :       // reset our buffer
      68         [ +  + ]:      19471 :       this->setp( m_buffer.get(), m_buffer.get() + BUFFER_SIZE );
      69                 :            : 
      70                 :            :       // write the passed character if necessary
      71         [ +  + ]:      19471 :       if ( !traits_type::eq_int_type(c, traits_type::eof()) ) {
      72                 :            :         traits_type::assign( *this->pptr(), traits_type::to_char_type(c) );
      73                 :            :         this->pbump(1);
      74                 :            :       }
      75                 :            : 
      76                 :            :       return traits_type::not_eof(c);
      77                 :            :     }
      78                 :            : 
      79                 :      19454 :     virtual int sync() override {
      80                 :            :       // flush our buffer into m_sbuf1 and m_sbuf2
      81                 :      19454 :       int_type c = this->overflow(traits_type::eof());
      82                 :            : 
      83                 :            :       // checking return for eof.
      84         [ +  - ]:      19454 :       if (traits_type::eq_int_type(c, traits_type::eof()))
      85                 :            :           return -1;
      86                 :            : 
      87                 :            :       // flush m_sbuf1 and m_sbuf2
      88 [ +  - ][ -  + ]:      19454 :       if (m_sbuf1->pubsync() == -1 || m_sbuf2->pubsync() == -1)
      89                 :          0 :           return -1;
      90                 :            : 
      91                 :            :       return 0;
      92                 :            :     }
      93                 :            : };
      94                 :            : 
      95                 :            : using teebuf = basic_teebuf< char >;
      96                 :            : using wteebuf = basic_teebuf< wchar_t >;
      97                 :            : 
      98                 :            : template< class charT, class traits = std::char_traits< charT > >
      99                 :            : struct scoped_basic_streambuf_assignment {
     100                 :            :   using stream_type = std::basic_ios< charT, traits >;
     101                 :            :   using streambuf_type = std::basic_streambuf< charT, traits >;
     102                 :            : 
     103                 :            :   stream_type& m_s;
     104                 :            :   streambuf_type* m_orig_sb;
     105                 :            : 
     106                 :      10636 :   scoped_basic_streambuf_assignment( stream_type &s, streambuf_type *new_sb )
     107         [ +  - ]:      10636 :     : m_s(s) { m_orig_sb = m_s.rdbuf( new_sb ); }
     108                 :            : 
     109                 :            :   ~scoped_basic_streambuf_assignment() { m_s.rdbuf( m_orig_sb ); }
     110                 :            : };
     111                 :            : 
     112                 :            : using scoped_streambuf_assignment = scoped_basic_streambuf_assignment<char>;
     113                 :            : using scoped_wstreambuf_assignment = scoped_basic_streambuf_assignment<wchar_t>;
     114                 :            : 
     115                 :            : } // tk::
     116                 :            : 
     117                 :            : #endif // TeeBuf_h