/root/doris/be/src/util/jsonb_writer.h

Source (jump to first uncovered line)
/*
 *  Copyright (c) 2014, Facebook, Inc.
 *  All rights reserved.
 *
 *  This source code is licensed under the BSD-style license found in the
 *  LICENSE file in the root directory of this source tree. An additional grant
 *  of patent rights can be found in the PATENTS file in the same directory.
 *
 */

/*
 * This file defines JsonbWriterT (template) and JsonbWriter.
 *
 * JsonbWriterT is a template class which implements an JSONB serializer.
 * Users call various write functions of JsonbWriterT object to write values
 * directly to JSONB packed bytes. All write functions of value or key return
 * the number of bytes written to JSONB, or 0 if there is an error. To write an
 * object, an array, or a string, you must call writeStart[..] before writing
 * values or key, and call writeEnd[..] after finishing at the end.
 *
 * By default, an JsonbWriterT object creates an output stream buffer.
 * Alternatively, you can also pass any output stream object to a writer, as
 * long as the stream object implements some basic functions of std::ostream
 * (such as JsonbOutStream, see JsonbStream.h).
 *
 * JsonbWriter specializes JsonbWriterT with JsonbOutStream type (see
 * JsonbStream.h). So unless you want to provide own a different output stream
 * type, use JsonbParser object.
 *
 * @author Tian Xia <tianx@fb.com>
 * this file is copied from 
 * https://github.com/facebook/mysql-5.6/blob/fb-mysql-5.6.35/fbson/FbsonWriter.h
 * and modified by Doris
 */

#ifndef JSONB_JSONBWRITER_H
#define JSONB_JSONBWRITER_H

#include <limits>
#include <stack>
#include <string>

#include "common/exception.h"
#include "jsonb_document.h"
#include "jsonb_stream.h"

namespace doris {

using int128_t = __int128;

template <class OS_TYPE>
class JsonbWriterT {
public:
    JsonbWriterT() : alloc_(true), hasHdr_(false), kvState_(WS_Value), str_pos_(0) {
        os_ = new OS_TYPE();
    }

    explicit JsonbWriterT(OS_TYPE& os)
            : os_(&os), alloc_(false), hasHdr_(false), kvState_(WS_Value), str_pos_(0) {}

    ~JsonbWriterT() {
        if (alloc_) {
            delete os_;
        }
    }

    void reset() {
        os_->clear();
        os_->seekp(0);
        hasHdr_ = false;
        kvState_ = WS_Value;
        first_ = true;
        for (; !stack_.empty(); stack_.pop())
            ;
    }

    uint32_t writeKey(const char* key, hDictInsert handler = nullptr) {
        return writeKey(key, strlen(key), handler);
    }

    // write a key string (or key id if an external dict is provided)
    uint32_t writeKey(const char* key, uint8_t len, hDictInsert handler = nullptr) {
        if (!stack_.empty() && verifyKeyState()) {
            int key_id = -1;
            if (handler) {
                key_id = handler(key, len);
            }

            uint32_t size = sizeof(uint8_t);
            if (key_id < 0) {
                os_->put(len);
                if (len == 0) {
                    // NOTE: we use sMaxKeyId to represent an empty key
                    JsonbKeyValue::keyid_type idx = JsonbKeyValue::sMaxKeyId;
                    os_->write((char*)&idx, sizeof(JsonbKeyValue::keyid_type));
                    size += sizeof(JsonbKeyValue::keyid_type);
                } else {
                    os_->write(key, len);
                    size += len;
                }
            } else if (key_id < JsonbKeyValue::sMaxKeyId) {
                JsonbKeyValue::keyid_type idx = key_id;
                os_->put(0);
                os_->write((char*)&idx, sizeof(JsonbKeyValue::keyid_type));
                size += sizeof(JsonbKeyValue::keyid_type);
            } else { // key id overflow
                assert(0);
                return 0;
            }

            kvState_ = WS_Key;
            return size;
        }

        return 0;
    }

    uint32_t writeValue(const JsonbValue* value) {
        if ((first_ && stack_.empty()) || (!stack_.empty() && verifyValueState())) {
            if (!writeFirstHeader()) return 0;
            os_->write((char*)value, value->numPackedBytes());
            kvState_ = WS_Value;
            return value->size();
        }
        return 0;
    }

    // write a key id
    uint32_t writeKey(JsonbKeyValue::keyid_type idx) {
        if (!stack_.empty() && verifyKeyState()) {
            os_->put(0);
            os_->write((char*)&idx, sizeof(JsonbKeyValue::keyid_type));
            kvState_ = WS_Key;
            return sizeof(uint8_t) + sizeof(JsonbKeyValue::keyid_type);
        }

        return 0;
    }

    bool writeFirstHeader() {
        if (first_ && stack_.empty()) {
            first_ = false;
            // if this is a new JSONB, write the header
            if (!hasHdr_) {
                writeHeader();
                return true;
            } else {
                return false;
            }
        } else {
            return true;
        }
    }

    uint32_t writeNull() {
        if ((first_ && stack_.empty()) || (!stack_.empty() && verifyValueState())) {
            if (!writeFirstHeader()) return 0;
            os_->put((JsonbTypeUnder)JsonbType::T_Null);
            kvState_ = WS_Value;
            return sizeof(JsonbValue);
        }

        return 0;
    }

    uint32_t writeBool(bool b) {
        if ((first_ && stack_.empty()) || (!stack_.empty() && verifyValueState())) {
            if (!writeFirstHeader()) return 0;
            if (b) {
                os_->put((JsonbTypeUnder)JsonbType::T_True);
            } else {
                os_->put((JsonbTypeUnder)JsonbType::T_False);
            }

            kvState_ = WS_Value;
            return sizeof(JsonbValue);
        }

        return 0;
    }

    // This function is a helper. It will make use of smallest space to
    // write an int
    uint32_t writeInt(int64_t val) {
        if (val >= std::numeric_limits<int8_t>::min() &&
            val <= std::numeric_limits<int8_t>::max()) {
            return writeInt8((int8_t)val);
        } else if (val >= std::numeric_limits<int16_t>::min() &&
                   val <= std::numeric_limits<int16_t>::max()) {
            return writeInt16((int16_t)val);
        } else if (val >= std::numeric_limits<int32_t>::min() &&
                   val <= std::numeric_limits<int32_t>::max()) {
            return writeInt32((int32_t)val);
        } else {
            return writeInt64(val);
        }
    }

    uint32_t writeInt8(int8_t v) {
        if ((first_ && stack_.empty()) || (!stack_.empty() && verifyValueState())) {
            if (!writeFirstHeader()) return 0;
            os_->put((JsonbTypeUnder)JsonbType::T_Int8);
            os_->put(v);
            kvState_ = WS_Value;
            return sizeof(JsonbInt8Val);
        }

        return 0;
    }

    uint32_t writeInt16(int16_t v) {
        if ((first_ && stack_.empty()) || (!stack_.empty() && verifyValueState())) {
            if (!writeFirstHeader()) return 0;
            os_->put((JsonbTypeUnder)JsonbType::T_Int16);
            os_->write((char*)&v, sizeof(int16_t));
            kvState_ = WS_Value;
            return sizeof(JsonbInt16Val);
        }

        return 0;
    }

    uint32_t writeInt32(int32_t v) {
        if ((first_ && stack_.empty()) || (!stack_.empty() && verifyValueState())) {
            if (!writeFirstHeader()) return 0;
            os_->put((JsonbTypeUnder)JsonbType::T_Int32);
            os_->write((char*)&v, sizeof(int32_t));
            kvState_ = WS_Value;
            return sizeof(JsonbInt32Val);
        }

        return 0;
    }

    uint32_t writeInt64(int64_t v) {
        if ((first_ && stack_.empty()) || (!stack_.empty() && verifyValueState())) {
            if (!writeFirstHeader()) return 0;
            os_->put((JsonbTypeUnder)JsonbType::T_Int64);
            os_->write((char*)&v, sizeof(int64_t));
            kvState_ = WS_Value;
            return sizeof(JsonbInt64Val);
        }

        return 0;
    }

    uint32_t writeInt128(int128_t v) {
        if ((first_ && stack_.empty()) || (!stack_.empty() && verifyValueState())) {
            if (!writeFirstHeader()) return 0;
            os_->put((JsonbTypeUnder)JsonbType::T_Int128);
            os_->write((char*)&v, sizeof(int128_t));
            kvState_ = WS_Value;
            return sizeof(JsonbInt128Val);
        }

        return 0;
    }

    uint32_t writeDouble(double v) {
        if ((first_ && stack_.empty()) || (!stack_.empty() && verifyValueState())) {
            if (!writeFirstHeader()) return 0;
            os_->put((JsonbTypeUnder)JsonbType::T_Double);
            os_->write((char*)&v, sizeof(double));
            kvState_ = WS_Value;
            return sizeof(JsonbDoubleVal);
        }

        return 0;
    }

    uint32_t writeFloat(float v) {
        if ((first_ && stack_.empty()) || (!stack_.empty() && verifyValueState())) {
            if (!writeFirstHeader()) return 0;
            os_->put((JsonbTypeUnder)JsonbType::T_Float);
            os_->write((char*)&v, sizeof(float));
            kvState_ = WS_Value;
            return sizeof(JsonbFloatVal);
        }

        return 0;
    }

    // must call writeStartString before writing a string val
    bool writeStartString() {
        if ((first_ && stack_.empty()) || (!stack_.empty() && verifyValueState())) {
            if (!writeFirstHeader()) return 0;
            os_->put((JsonbTypeUnder)JsonbType::T_String);
            str_pos_ = os_->tellp();

            // fill the size bytes with 0 for now
            uint32_t size = 0;
            os_->write((char*)&size, sizeof(uint32_t));

            kvState_ = WS_String;
            return true;
        }

        return false;
    }

    // finish writing a string val
    bool writeEndString() {
        if (kvState_ == WS_String) {
            std::streampos cur_pos = os_->tellp();
            int32_t size = (int32_t)(cur_pos - str_pos_ - sizeof(uint32_t));
            assert(size >= 0);

            os_->seekp(str_pos_);
            os_->write((char*)&size, sizeof(uint32_t));
            os_->seekp(cur_pos);

            kvState_ = WS_Value;
            return true;
        }

        return false;
    }

    uint32_t writeString(const char* str, uint32_t len) {
        if (kvState_ == WS_String) {
            os_->write(str, len);
            return len;
        }

        return 0;
    }

    uint32_t writeString(const std::string& str) {
        return writeString(str.c_str(), (uint32_t)str.size());
    }
    uint32_t writeString(char ch) {
        if (kvState_ == WS_String) {
            os_->put(ch);
            return 1;
        }

        return 0;
    }

    // must call writeStartBinary before writing a binary val
    bool writeStartBinary() {
        if ((first_ && stack_.empty()) || (!stack_.empty() && verifyValueState())) {
            if (!writeFirstHeader()) return 0;
            os_->put((JsonbTypeUnder)JsonbType::T_Binary);
            str_pos_ = os_->tellp();

            // fill the size bytes with 0 for now
            uint32_t size = 0;
            os_->write((char*)&size, sizeof(uint32_t));

            kvState_ = WS_Binary;
            return true;
        }

        return false;
    }

    // finish writing a binary val
    bool writeEndBinary() {
        if (kvState_ == WS_Binary) {
            std::streampos cur_pos = os_->tellp();
            int32_t size = (int32_t)(cur_pos - str_pos_ - sizeof(uint32_t));
            assert(size >= 0);

            os_->seekp(str_pos_);
            os_->write((char*)&size, sizeof(uint32_t));
            os_->seekp(cur_pos);

            kvState_ = WS_Value;
            return true;
        }

        return false;
    }

    uint32_t writeBinary(const char* bin, uint32_t len) {
        if (kvState_ == WS_Binary) {
            os_->write(bin, len);
            return len;
        }

        return 0;
    }

    // must call writeStartObject before writing an object val
    bool writeStartObject() {
        if (stack_.empty() || verifyValueState()) {
            if (stack_.empty()) {
                // if this is a new JSONB, write the header
                if (!hasHdr_) {
                    writeHeader();
                } else
                    return false;
            }

            // check if the object exceeds the maximum nesting level
            if (stack_.size() >= MaxNestingLevel) return false;

            os_->put((JsonbTypeUnder)JsonbType::T_Object);
            // save the size position
            stack_.push(WriteInfo({WS_Object, os_->tellp()}));

            // fill the size bytes with 0 for now
            uint32_t size = 0;
            os_->write((char*)&size, sizeof(uint32_t));

            kvState_ = WS_Value;
            return true;
        }

        return false;
    }

    // finish writing an object val
    bool writeEndObject() {
        if (!stack_.empty() && stack_.top().state == WS_Object && kvState_ == WS_Value) {
            WriteInfo& ci = stack_.top();
            std::streampos cur_pos = os_->tellp();
            int32_t size = (int32_t)(cur_pos - ci.sz_pos - sizeof(uint32_t));
            assert(size >= 0);

            os_->seekp(ci.sz_pos);
            os_->write((char*)&size, sizeof(uint32_t));
            os_->seekp(cur_pos);
            stack_.pop();

            return true;
        }

        return false;
    }

    // must call writeStartArray before writing an array val
    bool writeStartArray() {
        if (stack_.empty() || verifyValueState()) {
            if (stack_.empty()) {
                // if this is a new JSONB, write the header
                if (!hasHdr_) {
                    writeHeader();
                } else
                    return false;
            }

            // check if the array exceeds the maximum nesting level
            if (stack_.size() >= MaxNestingLevel) return false;

            os_->put((JsonbTypeUnder)JsonbType::T_Array);
            // save the size position
            stack_.push(WriteInfo({WS_Array, os_->tellp()}));

            // fill the size bytes with 0 for now
            uint32_t size = 0;
            os_->write((char*)&size, sizeof(uint32_t));

            kvState_ = WS_Value;
            return true;
        }

        return false;
    }

    // finish writing an array val
    bool writeEndArray() {
        if (!stack_.empty() && stack_.top().state == WS_Array && kvState_ == WS_Value) {
            WriteInfo& ci = stack_.top();
            std::streampos cur_pos = os_->tellp();
            int32_t size = (int32_t)(cur_pos - ci.sz_pos - sizeof(uint32_t));
            assert(size >= 0);

            os_->seekp(ci.sz_pos);
            os_->write((char*)&size, sizeof(uint32_t));
            os_->seekp(cur_pos);
            stack_.pop();

            return true;
        }

        return false;
    }

    OS_TYPE* getOutput() { return os_; }
    JsonbDocument* getDocument() {
        JsonbDocument* doc = nullptr;
        auto st = JsonbDocument::checkAndCreateDocument(getOutput()->getBuffer(),
                                                        getOutput()->getSize(), &doc);
        if (st.ok()) {
            return doc;
        } else {
            throw doris::Exception(st);
        }
    }

    JsonbValue* getValue() {
        return JsonbDocument::createValue(getOutput()->getBuffer(),
                                          (uint32_t)getOutput()->getSize());
    }

    bool writeEnd() {
        while (!stack_.empty()) {
            bool ok = false;
            switch (stack_.top().state) {
            case WS_Array:
                ok = writeEndArray();
                break;
            case WS_Object:
                ok = writeEndObject();
                break;
            case WS_String:
                ok = writeEndString();
                break;
            case WS_Binary:
                ok = writeEndBinary();
                break;
            default:
                ok = false;
                break;
            }
            if (ok == false) return false;
        }
        return true;
    }

private:
    // verify we are in the right state before writing a value
    bool verifyValueState() {
        assert(!stack_.empty());
        // The document can only be an Object or an Array which follows
        // the standard.
        return (stack_.top().state == WS_Object && kvState_ == WS_Key) ||
               (stack_.top().state == WS_Array && kvState_ == WS_Value);
    }

    // verify we are in the right state before writing a key
    bool verifyKeyState() {
        assert(!stack_.empty());
        return stack_.top().state == WS_Object && kvState_ == WS_Value;
    }

    void writeHeader() {
        os_->put(JSONB_VER);
        hasHdr_ = true;
    }

private:
    enum WriteState {
        WS_NONE,
        WS_Array,
        WS_Object,
        WS_Key,
        WS_Value,
        WS_String,
        WS_Binary,
    };

    struct WriteInfo {
        WriteState state;
        std::streampos sz_pos;
    };

private:
    OS_TYPE* os_ = nullptr;
    bool alloc_;
    bool hasHdr_;
    WriteState kvState_; // key or value state
    std::streampos str_pos_;
    std::stack<WriteInfo> stack_;
    bool first_ = true;
};

typedef JsonbWriterT<JsonbOutStream> JsonbWriter;

} // namespace doris

#endif // JSONB_JSONBWRITER_H

Coverage Report

Created: 2025-06-18 18:52

Line	Count	Source (jump to first uncovered line)
1		/*
2		* Copyright (c) 2014, Facebook, Inc.
3		* All rights reserved.
4		*
5		* This source code is licensed under the BSD-style license found in the
6		* LICENSE file in the root directory of this source tree. An additional grant
7		* of patent rights can be found in the PATENTS file in the same directory.
8		*
9		*/
10
11		/*
12		* This file defines JsonbWriterT (template) and JsonbWriter.
13		*
14		* JsonbWriterT is a template class which implements an JSONB serializer.
15		* Users call various write functions of JsonbWriterT object to write values
16		* directly to JSONB packed bytes. All write functions of value or key return
17		* the number of bytes written to JSONB, or 0 if there is an error. To write an
18		* object, an array, or a string, you must call writeStart[..] before writing
19		* values or key, and call writeEnd[..] after finishing at the end.
20		*
21		* By default, an JsonbWriterT object creates an output stream buffer.
22		* Alternatively, you can also pass any output stream object to a writer, as
23		* long as the stream object implements some basic functions of std::ostream
24		* (such as JsonbOutStream, see JsonbStream.h).
25		*
26		* JsonbWriter specializes JsonbWriterT with JsonbOutStream type (see
27		* JsonbStream.h). So unless you want to provide own a different output stream
28		* type, use JsonbParser object.
29		*
30		* @author Tian Xia <tianx@fb.com>
31		* this file is copied from
32		* https://github.com/facebook/mysql-5.6/blob/fb-mysql-5.6.35/fbson/FbsonWriter.h
33		* and modified by Doris
34		*/
35
36		#ifndef JSONB_JSONBWRITER_H
37		#define JSONB_JSONBWRITER_H
38
39		#include <limits>
40		#include <stack>
41		#include <string>
42
43		#include "common/exception.h"
44		#include "jsonb_document.h"
45		#include "jsonb_stream.h"
46
47		namespace doris {
48
49		using int128_t = __int128;
50
51		template <class OS_TYPE>
52		class JsonbWriterT {
53		public:
54	2.15k	JsonbWriterT() : alloc_(true), hasHdr_(false), kvState_(WS_Value), str_pos_(0) {
55	2.15k	os_ = new OS_TYPE();
56	2.15k	}
57
58		explicit JsonbWriterT(OS_TYPE& os)
59		: os_(&os), alloc_(false), hasHdr_(false), kvState_(WS_Value), str_pos_(0) {}
60
61	2.15k	~JsonbWriterT() {
62	2.15k	if (alloc_) {
63	2.15k	delete os_;
64	2.15k	}
65	2.15k	}
66
67	1.30k	void reset() {
68	1.30k	os_->clear();
69	1.30k	os_->seekp(0);
70	1.30k	hasHdr_ = false;
71	1.30k	kvState_ = WS_Value;
72	1.30k	first_ = true;
73	1.32k	for (; !stack_.empty(); stack_.pop())
74	11	;
75	1.30k	}
76
77		uint32_t writeKey(const char* key, hDictInsert handler = nullptr) {
78		return writeKey(key, strlen(key), handler);
79		}
80
81		// write a key string (or key id if an external dict is provided)
82	424	uint32_t writeKey(const char* key, uint8_t len, hDictInsert handler = nullptr) {
83	424	if (!stack_.empty() && verifyKeyState()) {
84	424	int key_id = -1;
85	424	if (handler) {
86	0	key_id = handler(key, len);
87	0	}
88
89	424	uint32_t size = sizeof(uint8_t);
90	424	if (key_id < 0) {
91	424	os_->put(len);
92	424	if (len == 0) {
93		// NOTE: we use sMaxKeyId to represent an empty key
94	1	JsonbKeyValue::keyid_type idx = JsonbKeyValue::sMaxKeyId;
95	1	os_->write((char*)&idx, sizeof(JsonbKeyValue::keyid_type));
96	1	size += sizeof(JsonbKeyValue::keyid_type);
97	423	} else {
98	423	os_->write(key, len);
99	423	size += len;
100	423	}
101	424	} else if (key_id < JsonbKeyValue::sMaxKeyId) {
102	0	JsonbKeyValue::keyid_type idx = key_id;
103	0	os_->put(0);
104	0	os_->write((char*)&idx, sizeof(JsonbKeyValue::keyid_type));
105	0	size += sizeof(JsonbKeyValue::keyid_type);
106	0	} else { // key id overflow
107	0	assert(0);
108	0	return 0;
109	0	}
110
111	424	kvState_ = WS_Key;
112	424	return size;
113	424	}
114
115	0	return 0;
116	424	}
117
118	35	uint32_t writeValue(const JsonbValue* value) {
119	35	if ((first_ && stack_.empty()) \|\| (!stack_.empty() && verifyValueState())) {
120	35	if (!writeFirstHeader()) return 0;
121	35	os_->write((char*)value, value->numPackedBytes());
122	35	kvState_ = WS_Value;
123	35	return value->size();
124	35	}
125	0	return 0;
126	35	}
127
128		// write a key id
129	10.3k	uint32_t writeKey(JsonbKeyValue::keyid_type idx) {
130	10.3k	if (!stack_.empty() && verifyKeyState()) {
131	9.29k	os_->put(0);
132	9.29k	os_->write((char*)&idx, sizeof(JsonbKeyValue::keyid_type));
133	9.29k	kvState_ = WS_Key;
134	9.29k	return sizeof(uint8_t) + sizeof(JsonbKeyValue::keyid_type);
135	9.29k	}
136
137	1.02k	return 0;
138	10.3k	}
139
140	11.6k	bool writeFirstHeader() {
141	11.6k	if (first_ && stack_.empty()) {
142	632	first_ = false;
143		// if this is a new JSONB, write the header
144	632	if (!hasHdr_) {
145	632	writeHeader();
146	632	return true;
147	632	} else {
148	0	return false;
149	0	}
150	10.9k	} else {
151	10.9k	return true;
152	10.9k	}
153	11.6k	}
154
155	2.22k	uint32_t writeNull() {
156	2.22k	if ((first_ && stack_.empty()) \|\| (!stack_.empty() && verifyValueState())) {
157	2.22k	if (!writeFirstHeader()) return 0;
158	2.22k	os_->put((JsonbTypeUnder)JsonbType::T_Null);
159	2.22k	kvState_ = WS_Value;
160	2.22k	return sizeof(JsonbValue);
161	2.22k	}
162
163	0	return 0;
164	2.22k	}
165
166	355	uint32_t writeBool(bool b) {
167	355	if ((first_ && stack_.empty()) \|\| (!stack_.empty() && verifyValueState())) {
168	355	if (!writeFirstHeader()) return 0;
169	355	if (b) {
170	181	os_->put((JsonbTypeUnder)JsonbType::T_True);
171	181	} else {
172	174	os_->put((JsonbTypeUnder)JsonbType::T_False);
173	174	}
174
175	355	kvState_ = WS_Value;
176	355	return sizeof(JsonbValue);
177	355	}
178
179	0	return 0;
180	355	}
181
182		// This function is a helper. It will make use of smallest space to
183		// write an int
184		uint32_t writeInt(int64_t val) {
185		if (val >= std::numeric_limits<int8_t>::min() &&
186		val <= std::numeric_limits<int8_t>::max()) {
187		return writeInt8((int8_t)val);
188		} else if (val >= std::numeric_limits<int16_t>::min() &&
189		val <= std::numeric_limits<int16_t>::max()) {
190		return writeInt16((int16_t)val);
191		} else if (val >= std::numeric_limits<int32_t>::min() &&
192		val <= std::numeric_limits<int32_t>::max()) {
193		return writeInt32((int32_t)val);
194		} else {
195		return writeInt64(val);
196		}
197		}
198
199	395	uint32_t writeInt8(int8_t v) {
200	395	if ((first_ && stack_.empty()) \|\| (!stack_.empty() && verifyValueState())) {
201	395	if (!writeFirstHeader()) return 0;
202	395	os_->put((JsonbTypeUnder)JsonbType::T_Int8);
203	395	os_->put(v);
204	395	kvState_ = WS_Value;
205	395	return sizeof(JsonbInt8Val);
206	395	}
207
208	0	return 0;
209	395	}
210
211	303	uint32_t writeInt16(int16_t v) {
212	303	if ((first_ && stack_.empty()) \|\| (!stack_.empty() && verifyValueState())) {
213	303	if (!writeFirstHeader()) return 0;
214	303	os_->put((JsonbTypeUnder)JsonbType::T_Int16);
215	303	os_->write((char*)&v, sizeof(int16_t));
216	303	kvState_ = WS_Value;
217	303	return sizeof(JsonbInt16Val);
218	303	}
219
220	0	return 0;
221	303	}
222
223	3.17k	uint32_t writeInt32(int32_t v) {
224	3.17k	if ((first_ && stack_.empty()) \|\| (!stack_.empty() && verifyValueState())) {
225	3.17k	if (!writeFirstHeader()) return 0;
226	3.17k	os_->put((JsonbTypeUnder)JsonbType::T_Int32);
227	3.17k	os_->write((char*)&v, sizeof(int32_t));
228	3.17k	kvState_ = WS_Value;
229	3.17k	return sizeof(JsonbInt32Val);
230	3.17k	}
231
232	0	return 0;
233	3.17k	}
234
235	68	uint32_t writeInt64(int64_t v) {
236	68	if ((first_ && stack_.empty()) \|\| (!stack_.empty() && verifyValueState())) {
237	68	if (!writeFirstHeader()) return 0;
238	68	os_->put((JsonbTypeUnder)JsonbType::T_Int64);
239	68	os_->write((char*)&v, sizeof(int64_t));
240	68	kvState_ = WS_Value;
241	68	return sizeof(JsonbInt64Val);
242	68	}
243
244	0	return 0;
245	68	}
246
247	1.02k	uint32_t writeInt128(int128_t v) {
248	1.02k	if ((first_ && stack_.empty()) \|\| (!stack_.empty() && verifyValueState())) {
249	1.02k	if (!writeFirstHeader()) return 0;
250	1.02k	os_->put((JsonbTypeUnder)JsonbType::T_Int128);
251	1.02k	os_->write((char*)&v, sizeof(int128_t));
252	1.02k	kvState_ = WS_Value;
253	1.02k	return sizeof(JsonbInt128Val);
254	1.02k	}
255
256	0	return 0;
257	1.02k	}
258
259	297	uint32_t writeDouble(double v) {
260	297	if ((first_ && stack_.empty()) \|\| (!stack_.empty() && verifyValueState())) {
261	297	if (!writeFirstHeader()) return 0;
262	297	os_->put((JsonbTypeUnder)JsonbType::T_Double);
263	297	os_->write((char*)&v, sizeof(double));
264	297	kvState_ = WS_Value;
265	297	return sizeof(JsonbDoubleVal);
266	297	}
267
268	0	return 0;
269	297	}
270
271	0	uint32_t writeFloat(float v) {
272	0	if ((first_ && stack_.empty()) \|\| (!stack_.empty() && verifyValueState())) {
273	0	if (!writeFirstHeader()) return 0;
274	0	os_->put((JsonbTypeUnder)JsonbType::T_Float);
275	0	os_->write((char*)&v, sizeof(float));
276	0	kvState_ = WS_Value;
277	0	return sizeof(JsonbFloatVal);
278	0	}
279
280	0	return 0;
281	0	}
282
283		// must call writeStartString before writing a string val
284	612	bool writeStartString() {
285	612	if ((first_ && stack_.empty()) \|\| (!stack_.empty() && verifyValueState())) {
286	612	if (!writeFirstHeader()) return 0;
287	612	os_->put((JsonbTypeUnder)JsonbType::T_String);
288	612	str_pos_ = os_->tellp();
289
290		// fill the size bytes with 0 for now
291	612	uint32_t size = 0;
292	612	os_->write((char*)&size, sizeof(uint32_t));
293
294	612	kvState_ = WS_String;
295	612	return true;
296	612	}
297
298	0	return false;
299	612	}
300
301		// finish writing a string val
302	612	bool writeEndString() {
303	612	if (kvState_ == WS_String) {
304	612	std::streampos cur_pos = os_->tellp();
305	612	int32_t size = (int32_t)(cur_pos - str_pos_ - sizeof(uint32_t));
306	612	assert(size >= 0);
307
308	0	os_->seekp(str_pos_);
309	612	os_->write((char*)&size, sizeof(uint32_t));
310	612	os_->seekp(cur_pos);
311
312	612	kvState_ = WS_Value;
313	612	return true;
314	612	}
315
316	0	return false;
317	612	}
318
319	612	uint32_t writeString(const char* str, uint32_t len) {
320	612	if (kvState_ == WS_String) {
321	612	os_->write(str, len);
322	612	return len;
323	612	}
324
325	0	return 0;
326	612	}
327
328	2	uint32_t writeString(const std::string& str) {
329	2	return writeString(str.c_str(), (uint32_t)str.size());
330	2	}
331		uint32_t writeString(char ch) {
332		if (kvState_ == WS_String) {
333		os_->put(ch);
334		return 1;
335		}
336
337		return 0;
338		}
339
340		// must call writeStartBinary before writing a binary val
341	3.11k	bool writeStartBinary() {
342	3.11k	if ((first_ && stack_.empty()) \|\| (!stack_.empty() && verifyValueState())) {
343	3.11k	if (!writeFirstHeader()) return 0;
344	3.11k	os_->put((JsonbTypeUnder)JsonbType::T_Binary);
345	3.11k	str_pos_ = os_->tellp();
346
347		// fill the size bytes with 0 for now
348	3.11k	uint32_t size = 0;
349	3.11k	os_->write((char*)&size, sizeof(uint32_t));
350
351	3.11k	kvState_ = WS_Binary;
352	3.11k	return true;
353	3.11k	}
354
355	0	return false;
356	3.11k	}
357
358		// finish writing a binary val
359	3.11k	bool writeEndBinary() {
360	3.11k	if (kvState_ == WS_Binary) {
361	3.11k	std::streampos cur_pos = os_->tellp();
362	3.11k	int32_t size = (int32_t)(cur_pos - str_pos_ - sizeof(uint32_t));
363	3.11k	assert(size >= 0);
364
365	0	os_->seekp(str_pos_);
366	3.11k	os_->write((char*)&size, sizeof(uint32_t));
367	3.11k	os_->seekp(cur_pos);
368
369	3.11k	kvState_ = WS_Value;
370	3.11k	return true;
371	3.11k	}
372
373	0	return false;
374	3.11k	}
375
376	3.11k	uint32_t writeBinary(const char* bin, uint32_t len) {
377	3.11k	if (kvState_ == WS_Binary) {
378	3.11k	os_->write(bin, len);
379	3.11k	return len;
380	3.11k	}
381
382	0	return 0;
383	3.11k	}
384
385		// must call writeStartObject before writing an object val
386	1.38k	bool writeStartObject() {
387	1.38k	if (stack_.empty() \|\| verifyValueState()) {
388	1.38k	if (stack_.empty()) {
389		// if this is a new JSONB, write the header
390	1.26k	if (!hasHdr_) {
391	1.26k	writeHeader();
392	1.26k	} else
393	0	return false;
394	1.26k	}
395
396		// check if the object exceeds the maximum nesting level
397	1.38k	if (stack_.size() >= MaxNestingLevel) return false;
398
399	1.38k	os_->put((JsonbTypeUnder)JsonbType::T_Object);
400		// save the size position
401	1.38k	stack_.push(WriteInfo({WS_Object, os_->tellp()}));
402
403		// fill the size bytes with 0 for now
404	1.38k	uint32_t size = 0;
405	1.38k	os_->write((char*)&size, sizeof(uint32_t));
406
407	1.38k	kvState_ = WS_Value;
408	1.38k	return true;
409	1.38k	}
410
411	0	return false;
412	1.38k	}
413
414		// finish writing an object val
415	1.36k	bool writeEndObject() {
416	1.36k	if (!stack_.empty() && stack_.top().state == WS_Object && kvState_ == WS_Value) {
417	1.36k	WriteInfo& ci = stack_.top();
418	1.36k	std::streampos cur_pos = os_->tellp();
419	1.36k	int32_t size = (int32_t)(cur_pos - ci.sz_pos - sizeof(uint32_t));
420	1.36k	assert(size >= 0);
421
422	0	os_->seekp(ci.sz_pos);
423	1.36k	os_->write((char*)&size, sizeof(uint32_t));
424	1.36k	os_->seekp(cur_pos);
425	1.36k	stack_.pop();
426
427	1.36k	return true;
428	1.36k	}
429
430	0	return false;
431	1.36k	}
432
433		// must call writeStartArray before writing an array val
434	491	bool writeStartArray() {
435	491	if (stack_.empty() \|\| verifyValueState()) {
436	491	if (stack_.empty()) {
437		// if this is a new JSONB, write the header
438	474	if (!hasHdr_) {
439	474	writeHeader();
440	474	} else
441	0	return false;
442	474	}
443
444		// check if the array exceeds the maximum nesting level
445	491	if (stack_.size() >= MaxNestingLevel) return false;
446
447	491	os_->put((JsonbTypeUnder)JsonbType::T_Array);
448		// save the size position
449	491	stack_.push(WriteInfo({WS_Array, os_->tellp()}));
450
451		// fill the size bytes with 0 for now
452	491	uint32_t size = 0;
453	491	os_->write((char*)&size, sizeof(uint32_t));
454
455	491	kvState_ = WS_Value;
456	491	return true;
457	491	}
458
459	0	return false;
460	491	}
461
462		// finish writing an array val
463	479	bool writeEndArray() {
464	479	if (!stack_.empty() && stack_.top().state == WS_Array && kvState_ == WS_Value) {
465	479	WriteInfo& ci = stack_.top();
466	479	std::streampos cur_pos = os_->tellp();
467	479	int32_t size = (int32_t)(cur_pos - ci.sz_pos - sizeof(uint32_t));
468	479	assert(size >= 0);
469
470	0	os_->seekp(ci.sz_pos);
471	479	os_->write((char*)&size, sizeof(uint32_t));
472	479	os_->seekp(cur_pos);
473	479	stack_.pop();
474
475	479	return true;
476	479	}
477
478	0	return false;
479	479	}
480
481	4.59k	OS_TYPE* getOutput() { return os_; }
482		JsonbDocument* getDocument() {
483		JsonbDocument* doc = nullptr;
484		auto st = JsonbDocument::checkAndCreateDocument(getOutput()->getBuffer(),
485		getOutput()->getSize(), &doc);
486		if (st.ok()) {
487		return doc;
488		} else {
489		throw doris::Exception(st);
490		}
491		}
492
493		JsonbValue* getValue() {
494		return JsonbDocument::createValue(getOutput()->getBuffer(),
495		(uint32_t)getOutput()->getSize());
496		}
497
498		bool writeEnd() {
499		while (!stack_.empty()) {
500		bool ok = false;
501		switch (stack_.top().state) {
502		case WS_Array:
503		ok = writeEndArray();
504		break;
505		case WS_Object:
506		ok = writeEndObject();
507		break;
508		case WS_String:
509		ok = writeEndString();
510		break;
511		case WS_Binary:
512		ok = writeEndBinary();
513		break;
514		default:
515		ok = false;
516		break;
517		}
518		if (ok == false) return false;
519		}
520		return true;
521		}
522
523		private:
524		// verify we are in the right state before writing a value
525	11.1k	bool verifyValueState() {
526	11.1k	assert(!stack_.empty());
527		// The document can only be an Object or an Array which follows
528		// the standard.
529	11.1k	return (stack_.top().state == WS_Object && kvState_ == WS_Key) \|\|
530	11.1k	(stack_.top().state == WS_Array && kvState_ == WS_Value);
531	11.1k	}
532
533		// verify we are in the right state before writing a key
534	10.7k	bool verifyKeyState() {
535	10.7k	assert(!stack_.empty());
536	10.7k	return stack_.top().state == WS_Object && kvState_ == WS_Value;
537	10.7k	}
538
539	2.37k	void writeHeader() {
540	2.37k	os_->put(JSONB_VER);
541	2.37k	hasHdr_ = true;
542	2.37k	}
543
544		private:
545		enum WriteState {
546		WS_NONE,
547		WS_Array,
548		WS_Object,
549		WS_Key,
550		WS_Value,
551		WS_String,
552		WS_Binary,
553		};
554
555		struct WriteInfo {
556		WriteState state;
557		std::streampos sz_pos;
558		};
559
560		private:
561		OS_TYPE* os_ = nullptr;
562		bool alloc_;
563		bool hasHdr_;
564		WriteState kvState_; // key or value state
565		std::streampos str_pos_;
566		std::stack<WriteInfo> stack_;
567		bool first_ = true;
568		};
569
570		typedef JsonbWriterT<JsonbOutStream> JsonbWriter;
571
572		} // namespace doris
573
574		#endif // JSONB_JSONBWRITER_H