338 lines · cpp
1//===-- MsgPackDocument.cpp - MsgPack Document --------------------------*-===//2//3// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.4// See https://llvm.org/LICENSE.txt for license information.5// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception6//7//===----------------------------------------------------------------------===//8///9/// This file implements a class that exposes a simple in-memory representation10/// of a document of MsgPack objects, that can be read from MsgPack, written to11/// MsgPack, and inspected and modified in memory. This is intended to be a12/// lighter-weight (in terms of memory allocations) replacement for13/// MsgPackTypes.14///15//===----------------------------------------------------------------------===//16 17#include "llvm/BinaryFormat/MsgPackDocument.h"18#include "llvm/BinaryFormat/MsgPackWriter.h"19 20using namespace llvm;21using namespace msgpack;22 23// Convert this DocNode into an empty array.24void DocNode::convertToArray() { *this = getDocument()->getArrayNode(); }25 26// Convert this DocNode into an empty map.27void DocNode::convertToMap() { *this = getDocument()->getMapNode(); }28 29/// Find the key in the MapDocNode.30DocNode::MapTy::iterator MapDocNode::find(StringRef S) {31 return find(getDocument()->getNode(S));32}33 34/// Member access for MapDocNode. The string data must remain valid for the35/// lifetime of the Document.36DocNode &MapDocNode::operator[](StringRef S) {37 return (*this)[getDocument()->getNode(S)];38}39 40/// Member access for MapDocNode.41DocNode &MapDocNode::operator[](DocNode Key) {42 assert(!Key.isEmpty());43 DocNode &N = (*Map)[Key];44 if (N.isEmpty()) {45 // Ensure a new element has its KindAndDoc initialized.46 N = getDocument()->getEmptyNode();47 }48 return N;49}50 51/// Member access for MapDocNode for integer key.52DocNode &MapDocNode::operator[](int Key) {53 return (*this)[getDocument()->getNode(Key)];54}55DocNode &MapDocNode::operator[](unsigned Key) {56 return (*this)[getDocument()->getNode(Key)];57}58DocNode &MapDocNode::operator[](int64_t Key) {59 return (*this)[getDocument()->getNode(Key)];60}61DocNode &MapDocNode::operator[](uint64_t Key) {62 return (*this)[getDocument()->getNode(Key)];63}64 65/// Array element access. This extends the array if necessary.66DocNode &ArrayDocNode::operator[](size_t Index) {67 if (size() <= Index) {68 // Ensure new elements have their KindAndDoc initialized.69 Array->resize(Index + 1, getDocument()->getEmptyNode());70 }71 return (*Array)[Index];72}73 74// Convenience assignment operators. This only works if the destination75// DocNode has an associated Document, i.e. it was not constructed using the76// default constructor. The string one does not copy, so the string must77// remain valid for the lifetime of the Document. Use fromString to avoid78// that restriction.79DocNode &DocNode::operator=(StringRef Val) {80 *this = getDocument()->getNode(Val);81 return *this;82}83DocNode &DocNode::operator=(MemoryBufferRef Val) {84 *this = getDocument()->getNode(Val);85 return *this;86}87DocNode &DocNode::operator=(bool Val) {88 *this = getDocument()->getNode(Val);89 return *this;90}91DocNode &DocNode::operator=(int Val) {92 *this = getDocument()->getNode(Val);93 return *this;94}95DocNode &DocNode::operator=(unsigned Val) {96 *this = getDocument()->getNode(Val);97 return *this;98}99DocNode &DocNode::operator=(int64_t Val) {100 *this = getDocument()->getNode(Val);101 return *this;102}103DocNode &DocNode::operator=(uint64_t Val) {104 *this = getDocument()->getNode(Val);105 return *this;106}107DocNode &DocNode::operator=(double Val) {108 *this = getDocument()->getNode(Val);109 return *this;110}111 112// A level in the document reading stack.113struct StackLevel {114 StackLevel(DocNode Node, size_t StartIndex, size_t Length,115 DocNode *MapEntry = nullptr)116 : Node(Node), Index(StartIndex), End(StartIndex + Length),117 MapEntry(MapEntry) {}118 DocNode Node;119 size_t Index;120 size_t End;121 // Points to map entry when we have just processed a map key.122 DocNode *MapEntry;123 DocNode MapKey;124};125 126// Read a document from a binary msgpack blob, merging into anything already in127// the Document.128// The blob data must remain valid for the lifetime of this Document (because a129// string object in the document contains a StringRef into the original blob).130// If Multi, then this sets root to an array and adds top-level objects to it.131// If !Multi, then it only reads a single top-level object, even if there are132// more, and sets root to that.133// Returns false if failed due to illegal format or merge error.134 135bool Document::readFromBlob(136 StringRef Blob, bool Multi,137 function_ref<int(DocNode *DestNode, DocNode SrcNode, DocNode MapKey)>138 Merger) {139 msgpack::Reader MPReader(Blob);140 SmallVector<StackLevel, 4> Stack;141 if (Multi) {142 // Create the array for multiple top-level objects.143 Root = getArrayNode();144 Stack.push_back(StackLevel(Root, 0, (size_t)-1));145 }146 do {147 // On to next element (or key if doing a map key next).148 // Read the value.149 Object Obj;150 Expected<bool> ReadObj = MPReader.read(Obj);151 if (!ReadObj) {152 // FIXME: Propagate the Error to the caller.153 consumeError(ReadObj.takeError());154 return false;155 }156 if (!ReadObj.get()) {157 if (Multi && Stack.size() == 1) {158 // OK to finish here as we've just done a top-level element with Multi159 break;160 }161 return false; // Finished too early162 }163 // Convert it into a DocNode.164 DocNode Node;165 switch (Obj.Kind) {166 case Type::Nil:167 Node = getNode();168 break;169 case Type::Int:170 Node = getNode(Obj.Int);171 break;172 case Type::UInt:173 Node = getNode(Obj.UInt);174 break;175 case Type::Boolean:176 Node = getNode(Obj.Bool);177 break;178 case Type::Float:179 Node = getNode(Obj.Float);180 break;181 case Type::String:182 Node = getNode(Obj.Raw);183 break;184 case Type::Binary:185 Node = getNode(MemoryBufferRef(Obj.Raw, ""));186 break;187 case Type::Map:188 Node = getMapNode();189 break;190 case Type::Array:191 Node = getArrayNode();192 break;193 default:194 return false; // Raw and Extension not supported195 }196 197 // Store it.198 DocNode *DestNode = nullptr;199 if (Stack.empty())200 DestNode = &Root;201 else if (Stack.back().Node.getKind() == Type::Array) {202 // Reading an array entry.203 auto &Array = Stack.back().Node.getArray();204 DestNode = &Array[Stack.back().Index++];205 } else {206 auto &Map = Stack.back().Node.getMap();207 if (!Stack.back().MapEntry) {208 // Reading a map key.209 Stack.back().MapKey = Node;210 Stack.back().MapEntry = &Map[Node];211 continue;212 }213 // Reading the value for the map key read in the last iteration.214 DestNode = Stack.back().MapEntry;215 Stack.back().MapEntry = nullptr;216 ++Stack.back().Index;217 }218 int MergeResult = 0;219 if (!DestNode->isEmpty()) {220 // In a merge, there is already a value at this position. Call the221 // callback to attempt to resolve the conflict. The resolution must result222 // in an array or map if Node is an array or map respectively.223 DocNode MapKey = !Stack.empty() && !Stack.back().MapKey.isEmpty()224 ? Stack.back().MapKey225 : getNode();226 MergeResult = Merger(DestNode, Node, MapKey);227 if (MergeResult < 0)228 return false; // Merge conflict resolution failed229 assert(!((Node.isMap() && !DestNode->isMap()) ||230 (Node.isArray() && !DestNode->isArray())));231 } else232 *DestNode = Node;233 234 // See if we're starting a new array or map.235 switch (DestNode->getKind()) {236 case msgpack::Type::Array:237 case msgpack::Type::Map:238 Stack.push_back(StackLevel(*DestNode, MergeResult, Obj.Length, nullptr));239 break;240 default:241 break;242 }243 244 // Pop finished stack levels.245 while (!Stack.empty()) {246 if (Stack.back().MapEntry)247 break;248 if (Stack.back().Index != Stack.back().End)249 break;250 Stack.pop_back();251 }252 } while (!Stack.empty());253 return true;254}255 256struct WriterStackLevel {257 DocNode Node;258 DocNode::MapTy::iterator MapIt;259 DocNode::ArrayTy::iterator ArrayIt;260 bool OnKey;261};262 263/// Write a MsgPack document to a binary MsgPack blob.264void Document::writeToBlob(std::string &Blob) {265 Blob.clear();266 raw_string_ostream OS(Blob);267 msgpack::Writer MPWriter(OS);268 SmallVector<WriterStackLevel, 4> Stack;269 DocNode Node = getRoot();270 for (;;) {271 switch (Node.getKind()) {272 case Type::Array:273 MPWriter.writeArraySize(Node.getArray().size());274 Stack.push_back(275 {Node, DocNode::MapTy::iterator(), Node.getArray().begin(), false});276 break;277 case Type::Map:278 MPWriter.writeMapSize(Node.getMap().size());279 Stack.push_back(280 {Node, Node.getMap().begin(), DocNode::ArrayTy::iterator(), true});281 break;282 case Type::Nil:283 MPWriter.writeNil();284 break;285 case Type::Boolean:286 MPWriter.write(Node.getBool());287 break;288 case Type::Int:289 MPWriter.write(Node.getInt());290 break;291 case Type::UInt:292 MPWriter.write(Node.getUInt());293 break;294 case Type::String:295 MPWriter.write(Node.getString());296 break;297 case Type::Binary:298 MPWriter.write(Node.getBinary());299 break;300 case Type::Float:301 MPWriter.write(Node.getFloat());302 break;303 case Type::Empty:304 llvm_unreachable("unhandled empty msgpack node");305 default:306 llvm_unreachable("unhandled msgpack object kind");307 }308 // Pop finished stack levels.309 while (!Stack.empty()) {310 if (Stack.back().Node.getKind() == Type::Map) {311 if (Stack.back().MapIt != Stack.back().Node.getMap().end())312 break;313 } else {314 if (Stack.back().ArrayIt != Stack.back().Node.getArray().end())315 break;316 }317 Stack.pop_back();318 }319 if (Stack.empty())320 break;321 // Get the next value.322 if (Stack.back().Node.getKind() == Type::Map) {323 if (Stack.back().OnKey) {324 // Do the key of a key,value pair in a map.325 Node = Stack.back().MapIt->first;326 Stack.back().OnKey = false;327 } else {328 Node = Stack.back().MapIt->second;329 ++Stack.back().MapIt;330 Stack.back().OnKey = true;331 }332 } else {333 Node = *Stack.back().ArrayIt;334 ++Stack.back().ArrayIt;335 }336 }337}338