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// 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 in
127// the Document.
128// The blob data must remain valid for the lifetime of this Document (because a
129// 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 are
132// 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(Elt: 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(Err: 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 Multi
159 break;
160 }
161 return false; // Finished too early
162 }
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(V: Obj.Int);
171 break;
172 case Type::UInt:
173 Node = getNode(V: Obj.UInt);
174 break;
175 case Type::Boolean:
176 Node = getNode(V: Obj.Bool);
177 break;
178 case Type::Float:
179 Node = getNode(V: Obj.Float);
180 break;
181 case Type::String:
182 Node = getNode(V: Obj.Raw);
183 break;
184 case Type::Binary:
185 Node = getNode(V: 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 supported
195 }
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 the
221 // callback to attempt to resolve the conflict. The resolution must result
222 // 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().MapKey
225 : getNode();
226 MergeResult = Merger(DestNode, Node, MapKey);
227 if (MergeResult < 0)
228 return false; // Merge conflict resolution failed
229 assert(!((Node.isMap() && !DestNode->isMap()) ||
230 (Node.isArray() && !DestNode->isArray())));
231 } else
232 *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(Elt: 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(Size: Node.getArray().size());
274 Stack.push_back(
275 Elt: {.Node: Node, .MapIt: DocNode::MapTy::iterator(), .ArrayIt: Node.getArray().begin(), .OnKey: false});
276 break;
277 case Type::Map:
278 MPWriter.writeMapSize(Size: Node.getMap().size());
279 Stack.push_back(
280 Elt: {.Node: Node, .MapIt: Node.getMap().begin(), .ArrayIt: DocNode::ArrayTy::iterator(), .OnKey: true});
281 break;
282 case Type::Nil:
283 MPWriter.writeNil();
284 break;
285 case Type::Boolean:
286 MPWriter.write(b: Node.getBool());
287 break;
288 case Type::Int:
289 MPWriter.write(i: Node.getInt());
290 break;
291 case Type::UInt:
292 MPWriter.write(u: Node.getUInt());
293 break;
294 case Type::String:
295 MPWriter.write(s: Node.getString());
296 break;
297 case Type::Binary:
298 MPWriter.write(Buffer: Node.getBinary());
299 break;
300 case Type::Float:
301 MPWriter.write(d: 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