line_wrapping_stream.h

Go to the documentation of this file.

#pragma once
 
#include <cassert>
#include <vector>
#include "vt_helper.h"
 
namespace ookii
{
 
    constexpr size_t use_console_width = std::numeric_limits<size_t>::max();
 
    template<typename CharType, typename Traits = std::char_traits<CharType>>
    class basic_line_wrapping_streambuf : public std::basic_streambuf<CharType, Traits>
    {
    public:
        using base_type = std::basic_streambuf<CharType, Traits>;
        using int_type = typename base_type::int_type;
        using char_type = typename base_type::char_type;
        using traits_type = typename base_type::traits_type;
 
        basic_line_wrapping_streambuf() noexcept = default;
 
        basic_line_wrapping_streambuf(base_type *streambuf, size_t max_line_length)
        {
            init(streambuf, max_line_length);
        }
 
        basic_line_wrapping_streambuf(basic_line_wrapping_streambuf &&other) noexcept
        {
            swap(other);
        }
 
        basic_line_wrapping_streambuf &operator=(basic_line_wrapping_streambuf &&other) noexcept
        {
            swap(other);
            return *this;
        }
 
        basic_line_wrapping_streambuf(basic_line_wrapping_streambuf &) = delete;
        basic_line_wrapping_streambuf &operator=(basic_line_wrapping_streambuf &) = delete;
 
        virtual ~basic_line_wrapping_streambuf()
        {
            sync(true);
        }
 
        void init(base_type *streambuf, size_t max_line_length, bool count_formatting = false) noexcept
        {
            if (_base_streambuf != nullptr)
            {
                flush_buffer(true);
            }
 
            _base_streambuf = streambuf;
            _count_formatting = count_formatting;
 
            // Check if the caller wants to use the console width.
            if (max_line_length == use_console_width)
            {
                _max_line_length = get_console_width();
                if (_max_line_length != 0)
                {
                    // Subtract one because it looks nicer.
                    --_max_line_length;
                }
            }
            else
            {
                _max_line_length = max_line_length;
            }
 
            // Don't allow super long lines, to avoid allocating large buffers.
            if (_max_line_length > c_max_allowed_line_length)
            {
                _max_line_length = c_max_allowed_line_length;
            }
 
            // Allocate a buffer only if line wrapping is used.
            if (_max_line_length > 0)
            {
                // Make the buffer twice the length of the line to simplify the line wrapping code.
                _buffer.resize(_max_line_length * 2);
                reset_put_area();
            }
            else
            {
                // Clear out the buffer if there was one.
                _buffer = {};
                this->setp(nullptr, nullptr);
            }
        }
 
        size_t indent() const
        {
            return _indent_count;
        }
 
        void indent(size_t indent) noexcept
        {
            if (_max_line_length > 0 && indent >= _max_line_length)
                _indent_count = 0;
            else
                _indent_count = indent;
        }
 
        bool reset_indent()
        {
            if (_base_streambuf == nullptr)
            {
                return false;
            }
 
            if (this->pubsync() != 0)
            {
                return false;
            }
 
            if (this->pptr() != this->pbase())
            {
                if (is_eof(this->sputc(_new_line)))
                {
                    return false;
                }
 
                if (this->pubsync() != 0)
                {
                    return false;
                }
 
                assert(this->pptr() == this->pbase());
            }
 
            _need_indent = false;
            return true;
        }
 
        void swap(basic_line_wrapping_streambuf &other) noexcept
        {
            if (this != std::addressof(other))
            {
                base_type::swap(other);
                std::swap(_base_streambuf, other._base_streambuf);
                std::swap(_max_line_length, other._max_line_length);
                std::swap(_buffer, other._buffer);
                std::swap(_new_line, other._new_line);
                std::swap(_space, other._space);
                std::swap(_indent_count, other._indent_count);
                std::swap(_need_indent, other._need_indent);
                std::swap(_count_formatting, other._count_formatting);
            }
        }
 
        virtual int sync(bool flush_last_line)
        {
            if (_base_streambuf == nullptr)
            {
                return -1;
            }
 
            // Nothing to flush if not using a buffer, but ask the base buffer to sync.
            if (_buffer.empty())
            {
                return _base_streambuf->pubsync();
            }
 
            // Attempt to flush the buffer if it's not empty.
            if (this->pptr() > this->pbase() && !flush_buffer(flush_last_line))
            {
                return -1;
            }
 
            return _base_streambuf->pubsync();
        }
 
    protected:
        virtual int_type overflow(int_type ch = traits_type::eof()) override
        {
            if (_base_streambuf == nullptr)
            {
                return traits_type::eof();
            }
 
            // If there is no maximum line length set, just forward the characters to the base stream.
            if (_buffer.size() == 0)
            {
                // Nothing to do if this is eof.
                if (is_eof(ch))
                {
                    return traits_type::not_eof(ch);
                }
 
                // If a line break, indicate the next line needs indentation, unless the previous
                // line was blank.
                if (is_new_line(ch))
                {
                    _need_indent = !_blank_line;
                    _blank_line = true;
                }
                else
                {
                    // Write indentation if needed.
                    // N.B. Don't indent blank lines.
                    if (_need_indent && !write_indent())
                    {
                       return traits_type::eof();
                    }
 
                    _blank_line = false;
                }
 
                // Write the character to the base stream.
                return _base_streambuf->sputc(traits_type::to_char_type(ch));
            }
 
            // If the buffer is not empty full, try to flush it.
            if (this->pptr() > this->pbase() && !flush_buffer())
            {
                return traits_type::eof();
            }
 
            if (!is_eof(ch))
            {
                // If the buffer is still full, try to grow it.
                if (this->pptr() == this->epptr() && !grow_buffer())
                {
                    return traits_type::eof();
                }
 
                // The buffer must now have space for at least one character.
                assert(this->pptr() != this->epptr());
                this->sputc(traits_type::to_char_type(ch));
            }
 
            // Return a character other than eof to indicate success.
            return traits_type::not_eof(ch);
        }
 
        virtual int sync() override
        {
            return sync(false);
        }
 
        virtual void imbue(const std::locale &loc) override
        {
            if (_base_streambuf != nullptr)
            {
                // Change the locale of the underlying buffer.
                _base_streambuf->pubimbue(loc);
                update_locale_characters(loc);
            }
        }
 
    private:
        using vt_helper_type = vt::details::vt_helper<CharType, Traits>;
 
        // Write the contents of the buffer to the underlying stream buffer, wrapping lines and
        // adding indentation as necessary.
        bool flush_buffer(bool flush_last_line = false)
        {
            if (_base_streambuf == nullptr)
            {
                return false;
            }
 
            auto start = this->pbase();
            auto end = this->pptr();
 
            // Only called when there is a buffer.
            assert(start != nullptr && end != nullptr);
 
            std::streamsize count;
            char_type *potential_line_break{};
            char_type *new_start{};
            auto locale = this->getloc();
            size_t line_length{};
            if (_need_indent)
            {
                line_length = _indent_count;
            }
 
            // Loop over the contents of the buffer.
            for (auto current = start; current < end; ++current)
            {
                // Check if this character can be used as a line break.
                if (std::isspace(*current, locale))
                {
                    potential_line_break = current;
                }
 
                if (!_count_formatting && traits_type::eq(*current, vt_helper_type::c_escape))
                {
                    auto vt_end = vt_helper_type::find_sequence_end(current + 1, end, locale);
                    if (vt_end == nullptr)
                    {
                        // Incomplete VT sequence, so we can't flush until the end is found.
                        break;
                    }
 
                    // Continue with the next character after the sequence.
                    current = vt_end;
                    continue;
                }
 
                // Check if we need to output a line.
                if (line_length >= _max_line_length || traits_type::eq(*current, _new_line))
                {
                    // Before writing the line, add indentation if necessary.
                    // N.B. Don't indent empty lines.
                    if (_need_indent)
                    {
                        line_length -= _indent_count;
                        if (line_length > 0 && !write_indent())
                        {
                            return false;
                        }
                    }
 
                    // If no place was found to break the line, break it here.
                    if (potential_line_break == nullptr)
                    {
                        potential_line_break = current;
                        new_start = current;
                    }
                    else
                    {
                        new_start = potential_line_break + 1;
                    }
 
                    // Write the line up to the break spot.
                    count = potential_line_break - start;
                    if (_base_streambuf->sputn(start, count) < count)
                    {
                        return false;
                    }
 
                    // Write the line break.
                    if (is_eof(_base_streambuf->sputc(_new_line)))
                    {
                        return false;
                    }
 
                    // Update the state for the new line.
                    start = new_start;
                    potential_line_break = nullptr;
                    _need_indent = line_length > 0;
                    line_length = ((current + 1) - start);
                    if (_need_indent)
                    {
                        line_length += _indent_count;
                    }
                }
                else
                {
                    ++line_length;
                }
            }
 
            if (flush_last_line)
            {
                if (end > start)
                {
                    if (_need_indent && !write_indent())
                    {
                        return false;
                    }
 
                    count = end - start;
                    // Write the remainder of the buffer plus a new line.
                    if (_base_streambuf->sputn(start, count) < count ||
                        is_eof(_base_streambuf->sputc(_new_line)))
                    {
                        return false;
                    }
                }
 
                reset_put_area(0);
            }
            else if (start > this->pbase())
            {
                // If we flushed any characters, move the remaining characters to the front for the
                // next flush.
                // N.B. line_length can't be used for this because it includes indent character
                // count.
                count = 0;
                if (end > start)
                {
                    count = end - start;
                    traits_type::move(this->pbase(), start, static_cast<size_t>(count));
                }
 
                // Reset the put area, and indicate the number of characters moved to the front.
                reset_put_area(static_cast<int>(count));
            }
 
            return true;
        }
 
        bool grow_buffer()
        {
            assert(this->epptr() == this->pptr());
            auto old_size = static_cast<int>(_buffer.size());
            auto new_size = old_size * 2;
            // Cap consumed memory at 2GB if overflow happened.
            if (new_size < old_size)
            {
                new_size = std::numeric_limits<std::int32_t>::max();
            }
 
            if (new_size == old_size)
            {
                return false;
            }
 
            _buffer.resize(new_size);
            reset_put_area(old_size);
            return true;
        }
 
        // Reset the put area pointers.
        void reset_put_area(int valid_data = 0)
        {
            this->setp(_buffer.data(), _buffer.data() + _buffer.size());
            if (valid_data > 0)
            {
                this->pbump(valid_data);
            }
        }
 
        // Update the new line and space characters based on the specified locale.
        void update_locale_characters(const std::locale &loc)
        {
            auto &ctype = std::use_facet<std::ctype<char_type>>(loc);
            _new_line = ctype.widen('\n');
            _space = ctype.widen(' ');
        }
 
        // Write indentation to the underlying stream buffer.
        bool write_indent()
        {
            if (_base_streambuf == nullptr)
            {
                return false;
            }
 
            assert(_need_indent);
            for (size_t i = 0; i < _indent_count; ++i)
            {
                if (is_eof(_base_streambuf->sputc(_space)))
                {
                    return false;
                }
            }
 
            _need_indent = false;
            return true;
        }
 
        // Checks if the character (in integer representation) equals a new line.
        bool is_new_line(int_type ch) const
        {
            return traits_type::eq_int_type(ch, traits_type::to_int_type(_new_line));
        }
 
        // Checks if the character (in integer representation) is an end-of-file character.
        static bool is_eof(int_type ch)
        {
            return traits_type::eq_int_type(ch, traits_type::eof());
        }
 
        static constexpr size_t c_max_allowed_line_length = 65536;
 
        base_type *_base_streambuf{};
        size_t _max_line_length{};
        std::vector<char_type> _buffer;
        char_type _new_line{};
        char_type _space{};
        size_t _indent_count{};
        bool _need_indent{};
        bool _blank_line{true};
        bool _count_formatting{};
    };
 
    using line_wrapping_streambuf = basic_line_wrapping_streambuf<char>;
    using wline_wrapping_streambuf = basic_line_wrapping_streambuf<wchar_t>;
 
    namespace details
    {
 
        // Returns a stream's buffer, or NULL if the stream doesn't use a basic_line_wrapping_streambuf.
        //
        // N.B. This may require RTTI.
        template<typename CharType, typename Traits>
        auto get_line_wrapping_streambuf(std::basic_ostream<CharType, Traits> &stream)
        {
            return dynamic_cast<basic_line_wrapping_streambuf<CharType, Traits>*>(stream.rdbuf());
        }
 
        // Stream manipulator to change the indent.
        struct set_indent_helper
        {
            size_t indent;
        };
 
        // Changes the indent size of the stream.
        //
        // N.B. Has no effect if the stream does not use a basic_line_wrapping_streambuf.
        template<typename CharType, typename Traits>
        std::basic_ostream<CharType, Traits> &operator<<(std::basic_ostream<CharType, Traits> &stream, const set_indent_helper &helper)
        {
            auto buffer = get_line_wrapping_streambuf(stream);
            if (buffer != nullptr)
            {
                buffer->indent(helper.indent);
            }
 
            return stream;
        }
 
        struct flush_helper
        {
            bool flush_last_line;
        };
 
        template<typename CharType, typename Traits>
        std::basic_ostream<CharType, Traits> &operator<<(std::basic_ostream<CharType, Traits> &stream, const flush_helper &helper)
        {
            auto buffer = get_line_wrapping_streambuf(stream);
            if (buffer != nullptr)
            {
                buffer->sync(helper.flush_last_line);
            }
            else
            {
                stream.flush();
            }
 
            return stream;
        }
    }
 
    inline details::set_indent_helper set_indent(size_t indent)
    {
        return details::set_indent_helper{indent};
    }
 
    template<typename CharType, typename Traits>
    std::basic_ostream<CharType, Traits> &reset_indent(std::basic_ostream<CharType, Traits> &stream)
    {
        auto buffer = details::get_line_wrapping_streambuf(stream);
        if (buffer != nullptr)
        {
            if (!buffer->reset_indent())
            {
                stream.setstate(std::ios::failbit);
            }
        }
 
        return stream;
    }
 
    inline details::flush_helper flush(bool flush_last_line)
    {
        return {flush_last_line};
    }
 
    template<typename CharType, typename Traits = std::char_traits<CharType>>
    class basic_line_wrapping_ostream : public std::basic_ostream<CharType, Traits>
    {
    public:
        using base_type = std::basic_ostream<CharType, Traits>;
        using streambuf_type = std::basic_streambuf<CharType, Traits>;
 
        basic_line_wrapping_ostream()
            : base_type{std::addressof(_buffer)}
        {
        }
 
        basic_line_wrapping_ostream(base_type &base_stream, size_t max_line_length, bool count_formatting = false)
            : base_type{std::addressof(_buffer)}
        {
            init(base_stream, max_line_length, count_formatting);
        }
 
        void init(base_type &base_stream, size_t max_line_length, bool count_formatting = false)
        {
            _buffer.init(base_stream.rdbuf(), max_line_length, count_formatting);
            this->imbue(base_stream.rdbuf()->getloc());
        }
 
        basic_line_wrapping_ostream(basic_line_wrapping_ostream &&other) noexcept
            : base_type{std::addressof(_buffer)}
        {
            swap(other);
        }
 
        basic_line_wrapping_ostream &operator=(basic_line_wrapping_ostream &&other) noexcept
        {
            swap(other);
            return *this;
        }
 
        static basic_line_wrapping_ostream for_cout(short default_width = 80)
        {
            return {console_stream<CharType>::cout(), static_cast<size_t>(get_console_width(default_width))};
        }
 
        static basic_line_wrapping_ostream for_cerr(short default_width = 80)
        {
            return {console_stream<CharType>::cerr(), static_cast<size_t>(get_console_width(default_width))};
        }
 
        void swap(basic_line_wrapping_ostream &other) noexcept
        {
            if (this != std::addressof(other))
            {
                base_type::swap(other);
                _buffer.swap(other._buffer);
            }
        }
 
        void flush(bool flush_last_line)
        {
            if (_buffer.sync(flush_last_line) < 0)
            {
                this->setstate(basic_line_wrapping_ostream::badbit);
            }
        }
 
    private:
        basic_line_wrapping_streambuf<CharType, Traits> _buffer;
    };
 
    using line_wrapping_ostream = basic_line_wrapping_ostream<char>;
    using wline_wrapping_ostream = basic_line_wrapping_ostream<wchar_t>;
 
    namespace details
    {
        template <class T, class = void>
        struct is_allocator : std::false_type
        {
        };
 
        template <class T>
        struct is_allocator<T, std::void_t<typename T::value_type, decltype(std::declval<T&>().deallocate(
            std::declval<T&>().allocate(size_t{ 1 }), size_t{ 1 })) >>
            : std::true_type
        {
        };
    }
 
    template<typename CharType, typename Traits = std::char_traits<CharType>, typename Allocator = std::allocator<CharType>>
    class basic_line_wrapping_ostringstream : public basic_line_wrapping_ostream<CharType, Traits>
    {
    public:
        using base_type = basic_line_wrapping_ostream<CharType, Traits>;
        using ostringstream_type = std::basic_ostringstream<CharType, Traits, Allocator>;
        using string_type = std::basic_string<CharType, Traits, Allocator>;
 
        basic_line_wrapping_ostringstream(size_t max_line_length, bool count_formatting = false)
        {
            this->init(_stringstream, max_line_length, count_formatting);
        }
 
        string_type str() const&
        {
            return _stringstream.str();
        }
 
        template<typename SAlloc, std::enable_if_t<details::is_allocator<SAlloc>::value, int> = 0>
        std::basic_string<CharType, Traits, SAlloc> str(const SAlloc &alloc) const
        {
            return _stringstream.str(alloc);
        }
 
        string_type str() &&
        {
            return std::move(_stringstream).str();
        }
 
        void str(const string_type &s)
        {
            _stringstream.str(s);
        }
 
        template<typename SAlloc>
        void str(const std::basic_string<CharType, Traits, SAlloc> &s)
        {
            _stringstream.str(s);
        }
 
        void str(string_type &&s)
        {
            _stringstream.str(std::move(s));
        }
 
        std::basic_string_view<CharType, Traits> view() const noexcept
        {
            return _stringstream.view();
        }
 
    private:
        ostringstream_type _stringstream;
    };
 
    using line_wrapping_ostringstream = basic_line_wrapping_ostringstream<char>;
    using wline_wrapping_ostringstream = basic_line_wrapping_ostringstream<wchar_t>;
}