blob: eed1ada89406c857e430b07701ee1475a9b16c43 [file] [log] [blame]
/*
* 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 "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() {
return JsonbDocument::createDocument(getOutput()->getBuffer(), getOutput()->getSize());
}
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_;
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