ydb/library/cpp/yson_pull/detail/writer.h

#pragma once

#include "byte_writer.h"
#include "cescape.h"
#include "percent_scalar.h"
#include "stream_counter.h"
#include "symbols.h"
#include "varint.h"

#include <library/cpp/yson_pull/consumer.h>
#include <library/cpp/yson_pull/event.h>
#include <library/cpp/yson_pull/output.h>
#include <library/cpp/yson_pull/stream_type.h>
#include <library/cpp/yson_pull/writer.h>

#include <util/generic/vector.h>
#include <util/system/yassert.h>

#include <cmath>

namespace NYsonPull {
    namespace NDetail {
        class writer: public IConsumer {
            enum class state {
                maybe_key,
                maybe_value,
                value,
                value_noattr,
                before_begin,
                before_end,
                after_end,
            };

            byte_writer<stream_counter<false>> stream_;
            TVector<EEventType> stack_;
            bool need_item_separator_ = false;
            EStreamType mode_ = EStreamType::ListFragment;
            state state_ = state::before_begin;

        public:
            void OnBeginStream() override {
                update_state(EEventType::BeginStream);
            }

            void OnEndStream() override {
                update_state(EEventType::EndStream);
                stream_.flush_buffer();
            }

            void OnBeginList() override {
                begin_node();
                write(NSymbol::begin_list);
                update_state(EEventType::BeginList);
                begin_collection(collection_type::list);
            }

            void OnEndList() override {
                update_state(EEventType::EndList);
                end_collection(collection_type::list);
                write(NSymbol::end_list);
                end_node();
            }

            void OnBeginMap() override {
                begin_node();
                write(NSymbol::begin_map);
                update_state(EEventType::BeginMap);
                begin_collection(collection_type::map);
            }

            void OnEndMap() override {
                update_state(EEventType::EndMap);
                end_collection(collection_type::map);
                write(NSymbol::end_map);
                end_node();
            }

            void OnBeginAttributes() override {
                begin_node();
                write(NSymbol::begin_attributes);
                update_state(EEventType::BeginAttributes);
                begin_collection(collection_type::attributes);
            }

            void OnEndAttributes() override {
                update_state(EEventType::EndAttributes);
                end_collection(collection_type::attributes);
                write(NSymbol::end_attributes);
                // no end_node
            }

            void OnEntity() override {
                begin_node();
                update_state(EEventType::Scalar);
                write(NSymbol::entity);
                end_node();
            }

        protected:
            enum class collection_type {
                list,
                map,
                attributes,
            };

            writer(NYsonPull::NOutput::IStream& stream, EStreamType mode)
                : stream_(stream)
                , mode_{mode} {
            }

            bool need_item_separator() const {
                return need_item_separator_;
            }
            void need_item_separator(bool value) {
                need_item_separator_ = value;
            }

            size_t depth() const {
                Y_ASSERT(!stack_.empty());
                if (mode_ == EStreamType::Node) {
                    return stack_.size() - 1;
                } else {
                    return stack_.size() - 2;
                }
            }
            EStreamType mode() const {
                return mode_;
            }

            void write(ui8 c) {
                stream_.write(c);
            }

            void write(TStringBuf value) {
                write_raw(value.data(), value.size());
            }

            void write_raw(const void* ptr, size_t len) {
                stream_.write(static_cast<const ui8*>(ptr), len);
            }

            template <typename T>
            void write_varint(T value) {
                NVarInt::write(stream_, value);
            }

            void write_escaped_string(TStringBuf value) {
                write(NSymbol::quote);
                NCEscape::encode(stream_, value);
                write(NSymbol::quote);
            }

            void push(EEventType type) {
                stack_.push_back(type);
            }

            void pop(EEventType type) {
                if (stack_.empty()) {
                    fail("Unpaired events: empty event stack");
                }
                if (stack_.back() != type) {
                    fail("Unpaired events: expected ", type, ", got ", stack_.back());
                }
                stack_.pop_back();
            }

            void update_state(EEventType event) {
                switch (state_) {
                    case state::before_begin:
                        if (event != EEventType::BeginStream) {
                            fail("Expected begin_stream, got ", event);
                        }
                        begin_stream();
                        return;

                    case state::before_end:
                        if (event != EEventType::EndStream) {
                            fail("Expected end_stream, got ", event);
                        }
                        end_stream();
                        return;

                    case state::after_end:
                        fail("Attempted write past stream end");

                    case state::maybe_key:
                        if (event == EEventType::Key) {
                            state_ = state::value;
                            return;
                        }

                        switch (event) {
                            case EEventType::EndStream:
                                end_stream();
                                return;

                            case EEventType::EndMap:
                                pop(EEventType::BeginMap);
                                next_state();
                                return;

                            case EEventType::EndAttributes:
                                pop(EEventType::BeginAttributes);
                                state_ = state::value_noattr;
                                return;

                            default:
                                fail("Unexpected event ", event, " in maybe_key");
                        }
                        break;

                    case state::maybe_value:
                        switch (event) {
                            case EEventType::EndList:
                                pop(EEventType::BeginList);
                                next_state();
                                return;

                            case EEventType::EndStream:
                                end_stream();
                                return;

                            default:
                                break;
                        }
                        [[fallthrough]];
                    case state::value:
                        if (event == EEventType::BeginAttributes) {
                            push(EEventType::BeginAttributes);
                            next_state();
                            return;
                        }
                        [[fallthrough]];
                    case state::value_noattr:
                        switch (event) {
                            case EEventType::Scalar:
                                next_state();
                                return;

                            case EEventType::BeginList:
                                push(EEventType::BeginList);
                                next_state();
                                return;

                            case EEventType::BeginMap:
                                push(EEventType::BeginMap);
                                next_state();
                                return;

                            default:
                                fail("Unexpected event ", event, " (in value_*)");
                        }
                        break;
                }
            }

            void next_state() {
                Y_ASSERT(!stack_.empty());
                switch (stack_.back()) {
                    case EEventType::BeginMap:
                    case EEventType::BeginAttributes:
                        state_ = state::maybe_key;
                        break;

                    case EEventType::BeginList:
                        state_ = state::maybe_value;
                        break;

                    case EEventType::BeginStream:
                        state_ = state::before_end;
                        break;

                    default:
                        Y_UNREACHABLE();
                }
            }

            void begin_stream() {
                push(EEventType::BeginStream);
                switch (mode_) {
                    case EStreamType::ListFragment:
                        push(EEventType::BeginList);
                        state_ = state::maybe_value;
                        break;

                    case EStreamType::MapFragment:
                        push(EEventType::BeginMap);
                        state_ = state::maybe_key;
                        break;

                    case EStreamType::Node:
                        state_ = state::value;
                        break;
                }
            }

            void end_stream() {
                switch (mode_) {
                    case EStreamType::ListFragment:
                        pop(EEventType::BeginList);
                        break;

                    case EStreamType::MapFragment:
                        pop(EEventType::BeginMap);
                        break;

                    case EStreamType::Node:
                        break;
                }
                pop(EEventType::BeginStream);
                state_ = state::after_end;
            }

            virtual void begin_node() {
                if (need_item_separator_) {
                    write(NSymbol::item_separator);
                }
            }

            virtual void end_node() {
                need_item_separator_ = true;
            }

            virtual void begin_key() {
                begin_node();
            }

            virtual void end_key() {
                need_item_separator_ = false;
                write(NSymbol::key_value_separator);
            }

            virtual void begin_collection(collection_type type) {
                Y_UNUSED(type);
                need_item_separator_ = false;
            }

            virtual void end_collection(collection_type type) {
                need_item_separator_ = (type != collection_type::attributes);
            }

            template <typename... Args>
            ATTRIBUTE(noinline, cold)
            void fail[[noreturn]](const char* msg, Args&&... args) {
                auto formatted_message = format_string(
                    msg,
                    std::forward<Args>(args)...);
                throw NException::TBadOutput(
                    formatted_message,
                    stream_.counter().info());
            }
        };

        class TBinaryWriterImpl final: public writer {
        public:
            TBinaryWriterImpl(NYsonPull::NOutput::IStream& stream, EStreamType mode)
                : writer(stream, mode)
            {
            }

            void OnScalarBoolean(bool value) override {
                update_state(EEventType::Scalar);

                begin_node();
                write(value ? NSymbol::true_marker : NSymbol::false_marker);
                end_node();
            }

            void OnScalarInt64(i64 value) override {
                update_state(EEventType::Scalar);

                begin_node();
                write(NSymbol::int64_marker);
                write_varint(value);
                end_node();
            }

            void OnScalarUInt64(ui64 value) override {
                update_state(EEventType::Scalar);

                begin_node();
                write(NSymbol::uint64_marker);
                write_varint(value);
                end_node();
            }

            void OnScalarFloat64(double value) override {
                update_state(EEventType::Scalar);

                begin_node();
                write(NSymbol::double_marker);
                write_raw(&value, sizeof value);
                end_node();
            }

            void OnScalarString(TStringBuf value) override {
                update_state(EEventType::Scalar);

                begin_node();
                write(NSymbol::string_marker);
                write_varint(static_cast<i32>(value.size()));
                write_raw(value.data(), value.size());
                end_node();
            }

            void OnKey(TStringBuf name) override {
                update_state(EEventType::Key);

                begin_key();
                write(NSymbol::string_marker);
                write_varint(static_cast<i32>(name.size()));
                write_raw(name.data(), name.size());
                end_key();
            }
        };

        class TTextWriterImpl: public writer {
        public:
            TTextWriterImpl(NYsonPull::NOutput::IStream& stream, EStreamType mode)
                : writer(stream, mode)
            {
            }

            void OnScalarBoolean(bool value) override {
                update_state(EEventType::Scalar);

                begin_node();
                write(value ? percent_scalar::true_literal : percent_scalar::false_literal);
                end_node();
            }

            void OnScalarInt64(i64 value) override {
                update_state(EEventType::Scalar);

                char buf[32];
                auto len = ::snprintf(buf, sizeof(buf), "%" PRIi64, value);

                begin_node();
                write_raw(buf, len);
                end_node();
            }

            void OnScalarUInt64(ui64 value) override {
                update_state(EEventType::Scalar);

                char buf[32];
                auto len = ::snprintf(buf, sizeof(buf), "%" PRIu64, value);

                begin_node();
                write_raw(buf, len);
                write('u');
                end_node();
            }

            void OnScalarFloat64(double value) override {
                update_state(EEventType::Scalar);

                begin_node();

                if (std::isfinite(value)) {
                    char buf[32];
                    auto len = ::snprintf(buf, sizeof(buf), "%#.17lg", value);
                    write_raw(buf, len);
                } else if (std::isnan(value)) {
                    write(percent_scalar::nan_literal);
                } else if (value > 0) {
                    write(percent_scalar::positive_inf_literal);
                } else {
                    write(percent_scalar::negative_inf_literal);
                }

                end_node();
            }

            void OnScalarString(TStringBuf value) override {
                update_state(EEventType::Scalar);

                begin_node();
                write_escaped_string(value);
                end_node();
            }

            void OnKey(TStringBuf name) override {
                update_state(EEventType::Key);

                begin_key();
                write_escaped_string(name);
                end_key();
            }

        protected:
            void begin_node() override {
                if (need_item_separator()) {
                    write(NSymbol::item_separator);
                    write(' ');
                }
            }

            void end_node() override {
                if (mode() != EStreamType::Node && depth() == 0) {
                    write(NSymbol::item_separator);
                    write('\n');
                    need_item_separator(false);
                } else {
                    writer::end_node();
                }
            }

            void end_key() override {
                write(' ');
                writer::end_key();
                write(' ');
            }
        };

        class TPrettyWriterImpl final: public TTextWriterImpl {
            size_t indent_size_;

        public:
            TPrettyWriterImpl(
                NYsonPull::NOutput::IStream& stream,
                EStreamType mode,
                size_t indent_size)
                : TTextWriterImpl(stream, mode)
                , indent_size_{indent_size} {
            }

        protected:
            void begin_node() override {
                if (need_item_separator()) {
                    write(NSymbol::item_separator);
                    newline();
                }
            }

            void begin_collection(collection_type type) override {
                TTextWriterImpl::begin_collection(type);
                newline();
            }

            void end_collection(collection_type type) override {
                TTextWriterImpl::end_collection(type);
                newline();
            }

            void newline() {
                write('\n');
                indent(depth());
            }

            void indent(size_t count) {
                for (size_t i = 0; i < count * indent_size_; ++i) {
                    write(' ');
                }
            }
        };

        template <typename T, typename... Args>
        NYsonPull::TWriter make_writer(
            THolder<NYsonPull::NOutput::IStream> stream,
            Args&&... args) {
            auto impl = MakeHolder<T>(*stream, std::forward<Args>(args)...);
            return NYsonPull::TWriter(std::move(stream), std::move(impl));
        }
    }
}