rocksdb/third-party/fbson/FbsonWriter.h

436 lines
10 KiB
C
Raw Normal View History

/*
* Copyright (c) 2011-present, 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 FbsonWriterT (template) and FbsonWriter.
*
* FbsonWriterT is a template class which implements an FBSON serializer.
* Users call various write functions of FbsonWriterT object to write values
* directly to FBSON packed bytes. All write functions of value or key return
* the number of bytes written to FBSON, 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 FbsonWriterT 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 FbsonOutStream, see FbsonStream.h).
*
* FbsonWriter specializes FbsonWriterT with FbsonOutStream type (see
* FbsonStream.h). So unless you want to provide own a different output stream
* type, use FbsonParser object.
*
* @author Tian Xia <tianx@fb.com>
*/
#ifndef FBSON_FBSONWRITER_H
#define FBSON_FBSONWRITER_H
#include <stack>
#include "FbsonDocument.h"
#include "FbsonStream.h"
namespace fbson {
template <class OS_TYPE>
class FbsonWriterT {
public:
FbsonWriterT()
: alloc_(true), hasHdr_(false), kvState_(WS_Value), str_pos_(0) {
os_ = new OS_TYPE();
}
explicit FbsonWriterT(OS_TYPE& os)
: os_(&os),
alloc_(false),
hasHdr_(false),
kvState_(WS_Value),
str_pos_(0) {}
~FbsonWriterT() {
if (alloc_) {
delete os_;
}
}
void reset() {
os_->clear();
os_->seekp(0);
hasHdr_ = false;
kvState_ = WS_Value;
for (; !stack_.empty(); stack_.pop())
;
}
// 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 (len && !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);
os_->write(key, len);
size += len;
} else if (key_id <= FbsonKeyValue::sMaxKeyId) {
FbsonKeyValue::keyid_type idx = key_id;
os_->put(0);
os_->write((char*)&idx, sizeof(FbsonKeyValue::keyid_type));
size += sizeof(FbsonKeyValue::keyid_type);
} else { // key id overflow
assert(0);
return 0;
}
kvState_ = WS_Key;
return size;
}
return 0;
}
// write a key id
uint32_t writeKey(FbsonKeyValue::keyid_type idx) {
if (!stack_.empty() && verifyKeyState()) {
os_->put(0);
os_->write((char*)&idx, sizeof(FbsonKeyValue::keyid_type));
kvState_ = WS_Key;
return sizeof(uint8_t) + sizeof(FbsonKeyValue::keyid_type);
}
return 0;
}
uint32_t writeNull() {
if (!stack_.empty() && verifyValueState()) {
os_->put((FbsonTypeUnder)FbsonType::T_Null);
kvState_ = WS_Value;
return sizeof(FbsonValue);
}
return 0;
}
uint32_t writeBool(bool b) {
if (!stack_.empty() && verifyValueState()) {
if (b) {
os_->put((FbsonTypeUnder)FbsonType::T_True);
} else {
os_->put((FbsonTypeUnder)FbsonType::T_False);
}
kvState_ = WS_Value;
return sizeof(FbsonValue);
}
return 0;
}
uint32_t writeInt8(int8_t v) {
if (!stack_.empty() && verifyValueState()) {
os_->put((FbsonTypeUnder)FbsonType::T_Int8);
os_->put(v);
kvState_ = WS_Value;
return sizeof(Int8Val);
}
return 0;
}
uint32_t writeInt16(int16_t v) {
if (!stack_.empty() && verifyValueState()) {
os_->put((FbsonTypeUnder)FbsonType::T_Int16);
os_->write((char*)&v, sizeof(int16_t));
kvState_ = WS_Value;
return sizeof(Int16Val);
}
return 0;
}
uint32_t writeInt32(int32_t v) {
if (!stack_.empty() && verifyValueState()) {
os_->put((FbsonTypeUnder)FbsonType::T_Int32);
os_->write((char*)&v, sizeof(int32_t));
kvState_ = WS_Value;
return sizeof(Int32Val);
}
return 0;
}
uint32_t writeInt64(int64_t v) {
if (!stack_.empty() && verifyValueState()) {
os_->put((FbsonTypeUnder)FbsonType::T_Int64);
os_->write((char*)&v, sizeof(int64_t));
kvState_ = WS_Value;
return sizeof(Int64Val);
}
return 0;
}
uint32_t writeDouble(double v) {
if (!stack_.empty() && verifyValueState()) {
os_->put((FbsonTypeUnder)FbsonType::T_Double);
os_->write((char*)&v, sizeof(double));
kvState_ = WS_Value;
return sizeof(DoubleVal);
}
return 0;
}
// must call writeStartString before writing a string val
bool writeStartString() {
if (!stack_.empty() && verifyValueState()) {
os_->put((FbsonTypeUnder)FbsonType::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(char ch) {
if (kvState_ == WS_String) {
os_->put(ch);
return 1;
}
return 0;
}
// must call writeStartBinary before writing a binary val
bool writeStartBinary() {
if (!stack_.empty() && verifyValueState()) {
os_->put((FbsonTypeUnder)FbsonType::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 FBSON, write the header
if (!hasHdr_) {
writeHeader();
} else
return false;
}
os_->put((FbsonTypeUnder)FbsonType::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 FBSON, write the header
if (!hasHdr_) {
writeHeader();
} else
return false;
}
os_->put((FbsonTypeUnder)FbsonType::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_; }
private:
// verify we are in the right state before writing a value
bool verifyValueState() {
assert(!stack_.empty());
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(FBSON_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_;
};
typedef FbsonWriterT<FbsonOutStream> FbsonWriter;
} // namespace fbson
#endif // FBSON_FBSONWRITER_H