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-exception
6//
7//===----------------------------------------------------------------------===//
8///
9/// This file implements a class that exposes a simple in-memory representation
10/// of a document of MsgPack objects, that can be read from MsgPack, written to
11/// MsgPack, and inspected and modified in memory. This is intended to be a
12/// lighter-weight (in terms of memory allocations) replacement for
13/// 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(Key: getDocument()->getNode(V: S));
32}
33
34/// Member access for MapDocNode. The string data must remain valid for the
35/// lifetime of the Document.
36DocNode &MapDocNode::operator[](StringRef S) {
37 return (*this)[getDocument()->getNode(V: 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(V: Key)];
54}
55DocNode &MapDocNode::operator[](unsigned Key) {
56 return (*this)[getDocument()->getNode(V: Key)];
57}
58DocNode &MapDocNode::operator[](int64_t Key) {
59 return (*this)[getDocument()->getNode(V: Key)];
60}
61DocNode &MapDocNode::operator[](uint64_t Key) {
62 return (*this)[getDocument()->getNode(V: 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(new_size: Index + 1, x: getDocument()->getEmptyNode());
70 }
71 return (*Array)[Index];
72}
73
74// Convenience assignment operators. This only works if the destination
75// DocNode has an associated Document, i.e. it was not constructed using the
76// default constructor. The string one does not copy, so the string must
77// remain valid for the lifetime of the Document. Use fromString to avoid
78// that restriction.
79DocNode &DocNode::operator=(StringRef Val) {
80 *this = getDocument()->getNode(V: Val);
81 return *this;
82}
83DocNode &DocNode::operator=(MemoryBufferRef Val) {
84 *this = getDocument()->getNode(V: Val);
85 return *this;
86}
87DocNode &DocNode::operator=(bool Val) {
88 *this = getDocument()->getNode(V: Val);
89 return *this;
90}
91DocNode &DocNode::operator=(int Val) {
92 *this = getDocument()->getNode(V: Val);
93 return *this;
94}
95DocNode &DocNode::operator=(unsigned Val) {
96 *this = getDocument()->getNode(V: Val);
97 return *this;
98}
99DocNode &DocNode::operator=(int64_t Val) {
100 *this = getDocument()->getNode(V: Val);
101 return *this;
102}
103DocNode &DocNode::operator=(uint64_t Val) {
104 *this = getDocument()->getNode(V: Val);
105 return *this;
106}
107DocNode &DocNode::operator=(double Val) {
108 *this = getDocument()->getNode(V: Val);
109 return *this;
110}
111
112// Equality operator. Compares recursively by value, supporting all node types
113// including Array and Map. Works correctly for nodes from different Documents.
114// This relies on operator< comparing scalar keys by value (not by document
115// identity), so that Map::find works across document boundaries.
116bool llvm::msgpack::operator==(const DocNode &Lhs, const DocNode &Rhs) {
117 if (Lhs.isEmpty() && Rhs.isEmpty())
118 return true;
119 if (Lhs.isEmpty() || Rhs.isEmpty())
120 return false;
121 if (Lhs.getKind() != Rhs.getKind())
122 return false;
123 switch (Lhs.getKind()) {
124 case Type::Nil:
125 return true;
126 case Type::Int:
127 return Lhs.Int == Rhs.Int;
128 case Type::UInt:
129 return Lhs.UInt == Rhs.UInt;
130 case Type::Boolean:
131 return Lhs.Bool == Rhs.Bool;
132 case Type::Float:
133 return Lhs.Float == Rhs.Float;
134 case Type::String:
135 case Type::Binary:
136 return Lhs.Raw == Rhs.Raw;
137 case Type::Array: {
138 if (Lhs.Array->size() != Rhs.Array->size())
139 return false;
140 for (size_t I = 0, E = Lhs.Array->size(); I != E; ++I)
141 if ((*Lhs.Array)[I] != (*Rhs.Array)[I])
142 return false;
143 return true;
144 }
145 case Type::Map: {
146 if (Lhs.Map->size() != Rhs.Map->size())
147 return false;
148 for (auto &Entry : *Lhs.Map) {
149 auto It = Rhs.Map->find(x: Entry.first);
150 if (It == Rhs.Map->end())
151 return false;
152 if (Entry.second != It->second)
153 return false;
154 }
155 return true;
156 }
157 default:
158 assert(false && "unhandled DocNode type in operator==");
159 return false;
160 }
161}
162
163/// Deep copy a DocNode from any Document into this Document.
164DocNode Document::copyNode(DocNode Src) {
165 if (Src.isEmpty())
166 return getEmptyNode();
167 switch (Src.getKind()) {
168 case Type::Nil:
169 return getNode();
170 case Type::Int:
171 return getNode(V: Src.getInt());
172 case Type::UInt:
173 return getNode(V: Src.getUInt());
174 case Type::Boolean:
175 return getNode(V: Src.getBool());
176 case Type::Float:
177 return getNode(V: Src.getFloat());
178 case Type::String:
179 // TODO: Restructure string interning so that no-copy strings from the
180 // source Document become no-copy strings in the destination Document,
181 // avoiding duplicate copies when the caller retains the source.
182 return getNode(V: Src.getString(), /*Copy=*/true);
183 case Type::Binary:
184 return getNode(V: Src.getBinary(), /*Copy=*/true);
185 case Type::Map: {
186 auto NewMap = getMapNode();
187 for (auto &Entry : Src.getMap())
188 NewMap[copyNode(Src: Entry.first)] = copyNode(Src: Entry.second);
189 return NewMap;
190 }
191 case Type::Array: {
192 auto NewArray = getArrayNode();
193 for (auto &Elem : Src.getArray())
194 NewArray.push_back(N: copyNode(Src: Elem));
195 return NewArray;
196 }
197 default:
198 assert(false && "unhandled DocNode type in copyNode");
199 return getEmptyNode();
200 }
201}
202
203// A level in the document reading stack.
204struct StackLevel {
205 StackLevel(DocNode Node, size_t StartIndex, size_t Length,
206 DocNode *MapEntry = nullptr)
207 : Node(Node), Index(StartIndex), End(StartIndex + Length),
208 MapEntry(MapEntry) {}
209 DocNode Node;
210 size_t Index;
211 size_t End;
212 // Points to map entry when we have just processed a map key.
213 DocNode *MapEntry;
214 DocNode MapKey;
215};
216
217// Read a document from a binary msgpack blob, merging into anything already in
218// the Document.
219// The blob data must remain valid for the lifetime of this Document (because a
220// string object in the document contains a StringRef into the original blob).
221// If Multi, then this sets root to an array and adds top-level objects to it.
222// If !Multi, then it only reads a single top-level object, even if there are
223// more, and sets root to that.
224// Returns false if failed due to illegal format or merge error.
225
226bool Document::readFromBlob(
227 StringRef Blob, bool Multi,
228 function_ref<int(DocNode *DestNode, DocNode SrcNode, DocNode MapKey)>
229 Merger) {
230 msgpack::Reader MPReader(Blob);
231 SmallVector<StackLevel, 4> Stack;
232 if (Multi) {
233 // Create the array for multiple top-level objects.
234 Root = getArrayNode();
235 Stack.push_back(Elt: StackLevel(Root, 0, (size_t)-1));
236 }
237 do {
238 // On to next element (or key if doing a map key next).
239 // Read the value.
240 Object Obj;
241 Expected<bool> ReadObj = MPReader.read(Obj);
242 if (!ReadObj) {
243 // FIXME: Propagate the Error to the caller.
244 consumeError(Err: ReadObj.takeError());
245 return false;
246 }
247 if (!ReadObj.get()) {
248 if (Multi && Stack.size() == 1) {
249 // OK to finish here as we've just done a top-level element with Multi
250 break;
251 }
252 return false; // Finished too early
253 }
254 // Convert it into a DocNode.
255 DocNode Node;
256 switch (Obj.Kind) {
257 case Type::Nil:
258 Node = getNode();
259 break;
260 case Type::Int:
261 Node = getNode(V: Obj.Int);
262 break;
263 case Type::UInt:
264 Node = getNode(V: Obj.UInt);
265 break;
266 case Type::Boolean:
267 Node = getNode(V: Obj.Bool);
268 break;
269 case Type::Float:
270 Node = getNode(V: Obj.Float);
271 break;
272 case Type::String:
273 Node = getNode(V: Obj.Raw);
274 break;
275 case Type::Binary:
276 Node = getNode(V: MemoryBufferRef(Obj.Raw, ""));
277 break;
278 case Type::Map:
279 Node = getMapNode();
280 break;
281 case Type::Array:
282 Node = getArrayNode();
283 break;
284 default:
285 return false; // Raw and Extension not supported
286 }
287
288 // Store it.
289 DocNode *DestNode = nullptr;
290 if (Stack.empty())
291 DestNode = &Root;
292 else if (Stack.back().Node.getKind() == Type::Array) {
293 // Reading an array entry.
294 auto &Array = Stack.back().Node.getArray();
295 DestNode = &Array[Stack.back().Index++];
296 } else {
297 auto &Map = Stack.back().Node.getMap();
298 if (!Stack.back().MapEntry) {
299 // Reading a map key.
300 Stack.back().MapKey = Node;
301 Stack.back().MapEntry = &Map[Node];
302 continue;
303 }
304 // Reading the value for the map key read in the last iteration.
305 DestNode = Stack.back().MapEntry;
306 Stack.back().MapEntry = nullptr;
307 ++Stack.back().Index;
308 }
309 int MergeResult = 0;
310 if (!DestNode->isEmpty()) {
311 // In a merge, there is already a value at this position. Call the
312 // callback to attempt to resolve the conflict. The resolution must result
313 // in an array or map if Node is an array or map respectively.
314 DocNode MapKey = !Stack.empty() && !Stack.back().MapKey.isEmpty()
315 ? Stack.back().MapKey
316 : getNode();
317 MergeResult = Merger(DestNode, Node, MapKey);
318 if (MergeResult < 0)
319 return false; // Merge conflict resolution failed
320 assert(!((Node.isMap() && !DestNode->isMap()) ||
321 (Node.isArray() && !DestNode->isArray())));
322 } else
323 *DestNode = Node;
324
325 // See if we're starting a new array or map.
326 switch (DestNode->getKind()) {
327 case msgpack::Type::Array:
328 case msgpack::Type::Map:
329 Stack.push_back(Elt: StackLevel(*DestNode, MergeResult, Obj.Length, nullptr));
330 break;
331 default:
332 break;
333 }
334
335 // Pop finished stack levels.
336 while (!Stack.empty()) {
337 if (Stack.back().MapEntry)
338 break;
339 if (Stack.back().Index != Stack.back().End)
340 break;
341 Stack.pop_back();
342 }
343 } while (!Stack.empty());
344 return true;
345}
346
347struct WriterStackLevel {
348 DocNode Node;
349 DocNode::MapTy::iterator MapIt;
350 DocNode::ArrayTy::iterator ArrayIt;
351 bool OnKey;
352};
353
354/// Write a MsgPack document to a binary MsgPack blob.
355void Document::writeToBlob(std::string &Blob) {
356 Blob.clear();
357 raw_string_ostream OS(Blob);
358 msgpack::Writer MPWriter(OS);
359 SmallVector<WriterStackLevel, 4> Stack;
360 DocNode Node = getRoot();
361 for (;;) {
362 switch (Node.getKind()) {
363 case Type::Array:
364 MPWriter.writeArraySize(Size: Node.getArray().size());
365 Stack.push_back(
366 Elt: {.Node: Node, .MapIt: DocNode::MapTy::iterator(), .ArrayIt: Node.getArray().begin(), .OnKey: false});
367 break;
368 case Type::Map:
369 MPWriter.writeMapSize(Size: Node.getMap().size());
370 Stack.push_back(
371 Elt: {.Node: Node, .MapIt: Node.getMap().begin(), .ArrayIt: DocNode::ArrayTy::iterator(), .OnKey: true});
372 break;
373 case Type::Nil:
374 MPWriter.writeNil();
375 break;
376 case Type::Boolean:
377 MPWriter.write(b: Node.getBool());
378 break;
379 case Type::Int:
380 MPWriter.write(i: Node.getInt());
381 break;
382 case Type::UInt:
383 MPWriter.write(u: Node.getUInt());
384 break;
385 case Type::String:
386 MPWriter.write(s: Node.getString());
387 break;
388 case Type::Binary:
389 MPWriter.write(Buffer: Node.getBinary());
390 break;
391 case Type::Float:
392 MPWriter.write(d: Node.getFloat());
393 break;
394 case Type::Empty:
395 llvm_unreachable("unhandled empty msgpack node");
396 default:
397 llvm_unreachable("unhandled msgpack object kind");
398 }
399 // Pop finished stack levels.
400 while (!Stack.empty()) {
401 if (Stack.back().Node.getKind() == Type::Map) {
402 if (Stack.back().MapIt != Stack.back().Node.getMap().end())
403 break;
404 } else {
405 if (Stack.back().ArrayIt != Stack.back().Node.getArray().end())
406 break;
407 }
408 Stack.pop_back();
409 }
410 if (Stack.empty())
411 break;
412 // Get the next value.
413 if (Stack.back().Node.getKind() == Type::Map) {
414 if (Stack.back().OnKey) {
415 // Do the key of a key,value pair in a map.
416 Node = Stack.back().MapIt->first;
417 Stack.back().OnKey = false;
418 } else {
419 Node = Stack.back().MapIt->second;
420 ++Stack.back().MapIt;
421 Stack.back().OnKey = true;
422 }
423 } else {
424 Node = *Stack.back().ArrayIt;
425 ++Stack.back().ArrayIt;
426 }
427 }
428}
429