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
20 using namespace llvm;
21 using namespace msgpack;
22
23 // Convert this DocNode into an empty array.
convertToArray()24 void DocNode::convertToArray() { *this = getDocument()->getArrayNode(); }
25
26 // Convert this DocNode into an empty map.
convertToMap()27 void DocNode::convertToMap() { *this = getDocument()->getMapNode(); }
28
29 /// Find the key in the MapDocNode.
find(StringRef S)30 DocNode::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 the
35 /// lifetime of the Document.
operator [](StringRef S)36 DocNode &MapDocNode::operator[](StringRef S) {
37 return (*this)[getDocument()->getNode(S)];
38 }
39
40 /// Member access for MapDocNode.
operator [](DocNode Key)41 DocNode &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.
operator [](int Key)52 DocNode &MapDocNode::operator[](int Key) {
53 return (*this)[getDocument()->getNode(Key)];
54 }
operator [](unsigned Key)55 DocNode &MapDocNode::operator[](unsigned Key) {
56 return (*this)[getDocument()->getNode(Key)];
57 }
operator [](int64_t Key)58 DocNode &MapDocNode::operator[](int64_t Key) {
59 return (*this)[getDocument()->getNode(Key)];
60 }
operator [](uint64_t Key)61 DocNode &MapDocNode::operator[](uint64_t Key) {
62 return (*this)[getDocument()->getNode(Key)];
63 }
64
65 /// Array element access. This extends the array if necessary.
operator [](size_t Index)66 DocNode &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 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.
operator =(StringRef Val)79 DocNode &DocNode::operator=(StringRef Val) {
80 *this = getDocument()->getNode(Val);
81 return *this;
82 }
operator =(MemoryBufferRef Val)83 DocNode &DocNode::operator=(MemoryBufferRef Val) {
84 *this = getDocument()->getNode(Val);
85 return *this;
86 }
operator =(bool Val)87 DocNode &DocNode::operator=(bool Val) {
88 *this = getDocument()->getNode(Val);
89 return *this;
90 }
operator =(int Val)91 DocNode &DocNode::operator=(int Val) {
92 *this = getDocument()->getNode(Val);
93 return *this;
94 }
operator =(unsigned Val)95 DocNode &DocNode::operator=(unsigned Val) {
96 *this = getDocument()->getNode(Val);
97 return *this;
98 }
operator =(int64_t Val)99 DocNode &DocNode::operator=(int64_t Val) {
100 *this = getDocument()->getNode(Val);
101 return *this;
102 }
operator =(uint64_t Val)103 DocNode &DocNode::operator=(uint64_t Val) {
104 *this = getDocument()->getNode(Val);
105 return *this;
106 }
107
108 // A level in the document reading stack.
109 struct StackLevel {
StackLevelStackLevel110 StackLevel(DocNode Node, size_t StartIndex, size_t Length,
111 DocNode *MapEntry = nullptr)
112 : Node(Node), Index(StartIndex), End(StartIndex + Length),
113 MapEntry(MapEntry) {}
114 DocNode Node;
115 size_t Index;
116 size_t End;
117 // Points to map entry when we have just processed a map key.
118 DocNode *MapEntry;
119 DocNode MapKey;
120 };
121
122 // Read a document from a binary msgpack blob, merging into anything already in
123 // the Document.
124 // The blob data must remain valid for the lifetime of this Document (because a
125 // string object in the document contains a StringRef into the original blob).
126 // If Multi, then this sets root to an array and adds top-level objects to it.
127 // If !Multi, then it only reads a single top-level object, even if there are
128 // more, and sets root to that.
129 // Returns false if failed due to illegal format or merge error.
130
readFromBlob(StringRef Blob,bool Multi,function_ref<int (DocNode * DestNode,DocNode SrcNode,DocNode MapKey)> Merger)131 bool Document::readFromBlob(
132 StringRef Blob, bool Multi,
133 function_ref<int(DocNode *DestNode, DocNode SrcNode, DocNode MapKey)>
134 Merger) {
135 msgpack::Reader MPReader(Blob);
136 SmallVector<StackLevel, 4> Stack;
137 if (Multi) {
138 // Create the array for multiple top-level objects.
139 Root = getArrayNode();
140 Stack.push_back(StackLevel(Root, 0, (size_t)-1));
141 }
142 do {
143 // On to next element (or key if doing a map key next).
144 // Read the value.
145 Object Obj;
146 Expected<bool> ReadObj = MPReader.read(Obj);
147 if (!ReadObj) {
148 // FIXME: Propagate the Error to the caller.
149 consumeError(ReadObj.takeError());
150 return false;
151 }
152 if (!ReadObj.get()) {
153 if (Multi && Stack.size() == 1) {
154 // OK to finish here as we've just done a top-level element with Multi
155 break;
156 }
157 return false; // Finished too early
158 }
159 // Convert it into a DocNode.
160 DocNode Node;
161 switch (Obj.Kind) {
162 case Type::Nil:
163 Node = getNode();
164 break;
165 case Type::Int:
166 Node = getNode(Obj.Int);
167 break;
168 case Type::UInt:
169 Node = getNode(Obj.UInt);
170 break;
171 case Type::Boolean:
172 Node = getNode(Obj.Bool);
173 break;
174 case Type::Float:
175 Node = getNode(Obj.Float);
176 break;
177 case Type::String:
178 Node = getNode(Obj.Raw);
179 break;
180 case Type::Binary:
181 Node = getNode(MemoryBufferRef(Obj.Raw, ""));
182 break;
183 case Type::Map:
184 Node = getMapNode();
185 break;
186 case Type::Array:
187 Node = getArrayNode();
188 break;
189 default:
190 return false; // Raw and Extension not supported
191 }
192
193 // Store it.
194 DocNode *DestNode = nullptr;
195 if (Stack.empty())
196 DestNode = &Root;
197 else if (Stack.back().Node.getKind() == Type::Array) {
198 // Reading an array entry.
199 auto &Array = Stack.back().Node.getArray();
200 DestNode = &Array[Stack.back().Index++];
201 } else {
202 auto &Map = Stack.back().Node.getMap();
203 if (!Stack.back().MapEntry) {
204 // Reading a map key.
205 Stack.back().MapKey = Node;
206 Stack.back().MapEntry = &Map[Node];
207 continue;
208 }
209 // Reading the value for the map key read in the last iteration.
210 DestNode = Stack.back().MapEntry;
211 Stack.back().MapEntry = nullptr;
212 ++Stack.back().Index;
213 }
214 int MergeResult = 0;
215 if (!DestNode->isEmpty()) {
216 // In a merge, there is already a value at this position. Call the
217 // callback to attempt to resolve the conflict. The resolution must result
218 // in an array or map if Node is an array or map respectively.
219 DocNode MapKey = !Stack.empty() && !Stack.back().MapKey.isEmpty()
220 ? Stack.back().MapKey
221 : getNode();
222 MergeResult = Merger(DestNode, Node, MapKey);
223 if (MergeResult < 0)
224 return false; // Merge conflict resolution failed
225 assert(!((Node.isMap() && !DestNode->isMap()) ||
226 (Node.isArray() && !DestNode->isArray())));
227 } else
228 *DestNode = Node;
229
230 // See if we're starting a new array or map.
231 switch (DestNode->getKind()) {
232 case msgpack::Type::Array:
233 case msgpack::Type::Map:
234 Stack.push_back(StackLevel(*DestNode, MergeResult, Obj.Length, nullptr));
235 break;
236 default:
237 break;
238 }
239
240 // Pop finished stack levels.
241 while (!Stack.empty()) {
242 if (Stack.back().MapEntry)
243 break;
244 if (Stack.back().Index != Stack.back().End)
245 break;
246 Stack.pop_back();
247 }
248 } while (!Stack.empty());
249 return true;
250 }
251
252 struct WriterStackLevel {
253 DocNode Node;
254 DocNode::MapTy::iterator MapIt;
255 DocNode::ArrayTy::iterator ArrayIt;
256 bool OnKey;
257 };
258
259 /// Write a MsgPack document to a binary MsgPack blob.
writeToBlob(std::string & Blob)260 void Document::writeToBlob(std::string &Blob) {
261 Blob.clear();
262 raw_string_ostream OS(Blob);
263 msgpack::Writer MPWriter(OS);
264 SmallVector<WriterStackLevel, 4> Stack;
265 DocNode Node = getRoot();
266 for (;;) {
267 switch (Node.getKind()) {
268 case Type::Array:
269 MPWriter.writeArraySize(Node.getArray().size());
270 Stack.push_back(
271 {Node, DocNode::MapTy::iterator(), Node.getArray().begin(), false});
272 break;
273 case Type::Map:
274 MPWriter.writeMapSize(Node.getMap().size());
275 Stack.push_back(
276 {Node, Node.getMap().begin(), DocNode::ArrayTy::iterator(), true});
277 break;
278 case Type::Nil:
279 MPWriter.writeNil();
280 break;
281 case Type::Boolean:
282 MPWriter.write(Node.getBool());
283 break;
284 case Type::Int:
285 MPWriter.write(Node.getInt());
286 break;
287 case Type::UInt:
288 MPWriter.write(Node.getUInt());
289 break;
290 case Type::String:
291 MPWriter.write(Node.getString());
292 break;
293 case Type::Binary:
294 MPWriter.write(Node.getBinary());
295 break;
296 case Type::Empty:
297 llvm_unreachable("unhandled empty msgpack node");
298 default:
299 llvm_unreachable("unhandled msgpack object kind");
300 }
301 // Pop finished stack levels.
302 while (!Stack.empty()) {
303 if (Stack.back().Node.getKind() == Type::Map) {
304 if (Stack.back().MapIt != Stack.back().Node.getMap().end())
305 break;
306 } else {
307 if (Stack.back().ArrayIt != Stack.back().Node.getArray().end())
308 break;
309 }
310 Stack.pop_back();
311 }
312 if (Stack.empty())
313 break;
314 // Get the next value.
315 if (Stack.back().Node.getKind() == Type::Map) {
316 if (Stack.back().OnKey) {
317 // Do the key of a key,value pair in a map.
318 Node = Stack.back().MapIt->first;
319 Stack.back().OnKey = false;
320 } else {
321 Node = Stack.back().MapIt->second;
322 ++Stack.back().MapIt;
323 Stack.back().OnKey = true;
324 }
325 } else {
326 Node = *Stack.back().ArrayIt;
327 ++Stack.back().ArrayIt;
328 }
329 }
330 }
331