TeeStreambuf.hpp

/*------------------------------------------------------------------------*/
/*                 Copyright 2010 Sandia Corporation.                     */
/*  Under terms of Contract DE-AC04-94AL85000, there is a non-exclusive   */
/*  license for use of this work by or on behalf of the U.S. Government.  */
/*  Export of this program may require a license from the                 */
/*  United States Government.                                             */
/*------------------------------------------------------------------------*/

#ifndef STK_UTIL_UTIL_TEESTREAMBUF_HPP
#define STK_UTIL_UTIL_TEESTREAMBUF_HPP

#include <string>
#include <streambuf>
#include <ostream>
#include <set>
#include <map>

namespace stk {

template<class Ch, class Tr = std::char_traits<Ch> >
class basic_tee_streambuf : public std::basic_streambuf<Ch, Tr>
{
  typedef std::set<std::ostream *> StreamSet;
  typedef std::map<std::ostream *, int> StreamErrorMap;
  
public:
  basic_tee_streambuf()
  {}

  explicit basic_tee_streambuf(std::basic_ostream<Ch, Tr> *os) {
    add(os);
  }

  virtual ~basic_tee_streambuf()
  {}

  int eof() {
    return std::basic_streambuf<Ch, Tr>::traits_type::eof();
  }

  void add(std::ostream *os) {
    m_destinations.insert(os);
  }

  void remove(std::ostream *os) {
    m_destinations.erase(os);
  }

  void clear() {
    m_destinations.clear();
  }

private:
  virtual int sync() {
    if (m_destinations.empty())
      return 1;

    StreamErrorMap return_code;

    for (StreamSet::const_iterator it = m_destinations.begin(); it != m_destinations.end(); ++it) {
      int ret = (*it)->rdbuf()->pubsync();
      return_code[*it] = ret;
    }

    // Remove streambufs with errors
    for (StreamSet::iterator it = m_destinations.begin(); it != m_destinations.end(); ++it)
      if (return_code[*it] == eof())
        m_destinations.erase(it);

    if (m_destinations.empty())
      return 1;

    return 1;
  }

  virtual typename std::basic_streambuf<Ch, Tr>::int_type overflow(const int c) {
    if (m_destinations.empty())
      return 1;

    StreamErrorMap return_code;

    for (StreamSet::const_iterator it = m_destinations.begin(); it != m_destinations.end(); ++it) {
      int ret = (*it)->rdbuf()->sputc(c);
      return_code[*it] = ret;
    }

    // Remove streambufs with errors
    for (StreamSet::iterator it = m_destinations.begin(); it != m_destinations.end(); ++it) 
      if (return_code[*it] == eof())
        m_destinations.erase(it);

    if (m_destinations.empty())
      return 1;

    return 1;
  }

  virtual std::streamsize xsputn(char const *buffer, std::streamsize n) {
    if (m_destinations.empty())
      return n;

    StreamErrorMap return_code;
    
    for (StreamSet::const_iterator it = m_destinations.begin(); it != m_destinations.end(); ++it) {
      std::ostream *os = (*it);
      int ret = os->rdbuf()->sputn(buffer,n);
      return_code[*it] = ret;
    }

    // Remove ostreams with errors
    for (StreamSet::iterator it = m_destinations.begin(); it != m_destinations.end(); ++it) {
      if (return_code[*it] < 0) {
        m_destinations.erase(it);
      }
    }

    if (m_destinations.empty())
      return n;

    return n;
  }

private:
  StreamSet             m_destinations;    
};

typedef stk::basic_tee_streambuf<char> tee_streambuf;

} // namespace stk

#endif // STK_UTIL_UTIL_TEESTREAMBUF_HPP